add main.py

main.py

@@ -1,64 +1,61 @@
 import os
+from pathlib import Path
 from dotenv import load_dotenv
-from agent import create_agent
+
+from langchain_ollama import ChatOllama
+from langchain.agents import Tool, AgentExecutor, create_openai_tools_agent
+from langchain.tools import tool
+from langchain.schema import HumanMessage
+
 from vectorstore import create_vectorstore, load_documents
 
-# Load environment variables
+# Load env
 load_dotenv()
+TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
+if not TAVILY_API_KEY:
+    raise RuntimeError("TAVILY_API_KEY not set in .env")
 
+# Setup vector store
+vectorstore = create_vectorstore()
+# Load documents if not already loaded
+if not Path("./chroma_db/chroma-collections.jsonl").exists():
+    load_documents("documents", vectorstore)
 
-def initialize_knowledge_base(documents_dir: str = "./documents"):
-    """Initialize the vectorstore and load documents if not already loaded."""
-    vectorstore = create_vectorstore()
-    
-    # Check if vectorstore is empty
-    collection_count = vectorstore._collection.count()
-    if collection_count == 0 and os.path.exists(documents_dir):
-        print(f"Loading documents from {documents_dir}...")
-        load_documents(documents_dir, vectorstore)
-    else:
-        print(f"Vectorstore already contains {collection_count} documents")
-    
-    return vectorstore
+retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
 
+@tool(name="search_local_kb", description="Search local knowledge base in ChromaDB")
+def search_local_kb(query: str, top_k: int = 3) -> str:
+    docs = retriever.invoke({"query": query, "k": top_k})
+    return "\n---\n".join([d.page_content for d in docs])
 
-def main():
-    """Main CLI chat loop."""
-    print("=" * 50)
-    print("RAG Agent with ChromaDB and Web Search")
-    print("=" * 50)
-    
-    # Initialize knowledge base
-    initialize_knowledge_base()
-    
-    # Create agent
-    agent = create_agent()
-    
-    print("\nAgent ready. Type 'exit' to quit.\n")
-    
-    while True:
-        try:
-            query = input("Запрос: ").strip()
-            
-            if query.lower() == "exit":
-                print("Goodbye!")
-                break
-            
-            if not query:
-                continue
-            
-            # Run agent
-            result = agent.invoke({"input": query})
-            answer = result.get("output", "No answer generated")
-            
-            print(f"\n{answer}\n")
-            
-        except KeyboardInterrupt:
-            print("\nGoodbye!")
-            break
-        except Exception as e:
-            print(f"Error: {e}")
+@tool(name="web_search", description="Search the web via Tavily")
+def web_search(query: str) -> str:
+    from tavily import TavilyClient
+    client = TavilyClient(api_key=TAVILY_API_KEY)
+    results = client.search(query, max_results=3)
+    return "\n---\n".join([f"{r.title}\n{r.url}" for r in results])
 
+tools = [search_local_kb, web_search]
 
-if __name__ == "__main__":
-    main()
+system_prompt = (
+    "You are an AI assistant that answers user questions.
+    If the answer can be found in local documents, use search_local_kb.\n"
+    "If the question is about recent events or requires up-to-date info, use web_search.\n"
+    "Always indicate the source of your answer: chromadb or tavily."
+)
+
+agent = create_openai_tools_agent(
+    llm=ChatOllama(model="llama3", temperature=0),
+    tools=tools,
+    system_message=system_prompt,
+)
+executor = AgentExecutor(agent=agent, tools=tools, verbose=True)
+
+print("RAG agent ready. Type 'exit' to quit.")
+while True:
+    user_input = input("Query: ")
+    if user_input.lower() in {"exit", "quit"}:
+        break
+    response = executor.invoke({"input": user_input})
+    print(response["output"])
+print("Goodbye!")