39 lines
1.2 KiB
Python
39 lines
1.2 KiB
Python
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() |