import os
from dotenv import load_dotenv

from vectorstore import create_vectorstore, load_documents
from agent import create_agent

def main():
    load_dotenv()  # загружает TAVILY_API_KEY из .env

    persist_dir = "./chroma_db"
    vectorstore = create_vectorstore(persist_directory=persist_dir)

    docs_dir = "./documents"
    load_documents(docs_dir, vectorstore)

    agent_executor = create_agent(vectorstore)

    print("RAG-агент готов. Введите 'exit' для выхода.")
    while True:
        query = input("\nЗапрос: ").strip()
        if query.lower() in ("exit", "quit"):
            break
        if not query:
            continue

        result = agent_executor.invoke({"input": query})
        # Ожидаем, что агент вернёт dict с полями 'output' и, опционально, 'source'
        if isinstance(result, dict):
            output = result.get("output", "")
            source = result.get("source", "unknown")
        else:
            output = str(result)
            source = "unknown"

        print(f"\n{output}")
        print(f"Источник: {source}")

if __name__ == "__main__":
    main()