diff --git a/main.py b/main.py index 092db74..ab7388b 100644 --- a/main.py +++ b/main.py @@ -1,25 +1,63 @@ -"""Entry point for the RAG agent. +""" +CLI entry point for the RAG agent. -Running ``python main.py`` starts a simple CLI that accepts user queries and -prints the agent's answer along with the source. +The script loads / creates the vector store, populates it from the ``documents`` +folder and starts an interactive chat loop. """ -from agent import create_agent, create_vectorstore -from vectorstore import load_documents +import os +import sys +from pathlib import Path -if __name__ == "__main__": - # Load or create the vector store - store = create_vectorstore() - if not store.get_index_info(): - load_documents("documents", store) - agent = create_agent() - print("RAG Agent ready. Type 'exit' to quit.") +from dotenv import load_dotenv + +# Load environment variables (TAVILY_API_KEY, etc.) +load_dotenv() + +# Import our modules +from vectorstore import create_vectorstore, load_documents +from agent import create_agent + +# --------------------------------------------------------------------------- +# Helper: populate vector store +# --------------------------------------------------------------------------- + +def init_vectorstore(persist_dir: str = "./chroma_db", docs_dir: str = "./documents"): + """Create or load the vector store and load documents if needed.""" + vectorstore = create_vectorstore(persist_directory=persist_dir) + # Always load documents – Chroma will deduplicate if already present. + print("Loading documents into ChromaDB…") + load_documents(docs_dir, vectorstore) + return vectorstore + +# --------------------------------------------------------------------------- +# Main chat loop +# --------------------------------------------------------------------------- + +def main(): + print("Initializing RAG agent…") + vectorstore = init_vectorstore() + agent = create_agent(vectorstore) + + print("RAG agent ready. Type your question (or 'exit' to quit).") while True: try: - query = input("\nЗапрос: ") - except EOFError: + user_input = input("\n> ") + except (EOFError, KeyboardInterrupt): + print("\nGoodbye!") break - if query.strip().lower() in {"exit", "quit"}: + if user_input.strip().lower() in {"exit", "quit", "q"}: + print("Goodbye!") break - result = agent.invoke({"input": query}) - print("Ответ:", result["output"]) # noqa: T201 + if not user_input.strip(): + continue + # Run the agent and capture the output + result = agent.run(user_input) + print("\nAnswer:\n", result) + +if __name__ == "__main__": + main() + +# --------------------------------------------------------------------------- +# End of script +# --------------------------------------------------------------------------- \ No newline at end of file