diff --git a/main.py b/main.py new file mode 100644 index 0000000..7e59d3a --- /dev/null +++ b/main.py @@ -0,0 +1,39 @@ +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() \ No newline at end of file