add main.py

main.py

@@ -1,61 +1,72 @@
+"""
+Simple CLI for the RAG‑agent.
+
+The agent automatically chooses between a local ChromaDB search and a web search via Tavily.
+It prints the answer together with the source label.
+"""
 import os
 from pathlib import Path
 from dotenv import load_dotenv
 
-from langchain_ollama import ChatOllama
-from langchain.agents import Tool, AgentExecutor, create_openai_tools_agent
+from langchain_ollama import OllamaLLM
+from langchain_tavily import TavilySearchResults
 from langchain.tools import tool
-from langchain.schema import HumanMessage
+from langchain.agents import initialize_agent, AgentType
+from langchain.chains import RetrievalQA
 
 from vectorstore import create_vectorstore, load_documents
 
-# Load env
+# Load env for Tavily API key
 load_dotenv()
 TAVILY_API_KEY = os.getenv("TAVILY_API_KEY")
 if not TAVILY_API_KEY:
-    raise RuntimeError("TAVILY_API_KEY not set in .env")
+    raise RuntimeError("TAVILY_API_KEY is missing in .env")
 
-# Setup vector store
+# 1. Vector store and retriever
 vectorstore = create_vectorstore()
-# Load documents if not already loaded
-if not Path("./chroma_db/chroma-collections.jsonl").exists():
-    load_documents("documents", vectorstore)
-
+load_documents("documents", vectorstore)
 retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
 
-@tool(name="search_local_kb", description="Search local knowledge base in ChromaDB")
+# 2. Tools
+@tool(name="search_local_kb", description="Search the local knowledge base using 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])
+    docs = retriever.invoke({"query": query}) if hasattr(retriever, "invoke") else retriever.get_relevant_documents(query)
+    return "\n---\n".join([f"{d.metadata.get('source')}:\n{d.page_content[:200]}…" for d in docs])
 
-@tool(name="web_search", description="Search the web via Tavily")
+@tool(name="web_search", description="Search the web using 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])
+    tavily = TavilySearchResults(api_key=TAVILY_API_KEY, max_results=3)
+    results = tavily.run(query)
+    return "\n---\n".join([f"{r['title']} ({r['url']}):\n{r.get('content', '')[:200]}…" for r in results])
 
-tools = [search_local_kb, web_search]
+# 3. Agent with simple routing logic
+from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
+from langchain.chat_models import ChatOllama
 
-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."
+chat = ChatOllama(model="llama3")
+
+system_prompt = "You are an assistant that can answer questions using either a local knowledge base or the web. If the question is about recent events or news, use web_search; otherwise use search_local_kb. Return the answer and specify the source as either chromadb or tavily."
+prompt = ChatPromptTemplate.from_messages([
+    SystemMessagePromptTemplate.from_template(system_prompt),
+    HumanMessagePromptTemplate.from_template("{input}")
+])
+
+agent_chain = initialize_agent(
+    tools=[search_local_kb, web_search],
+    llm=chat,
+    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
+    verbose=True,
 )
 
-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"}:
+# 4. CLI loop
+if __name__ == "__main__":
+    print("RAG Agent ready. Type 'exit' to quit.")
+    while True:
+        try:
+            q = input("Query: ")
+        except EOFError:
             break
-    response = executor.invoke({"input": user_input})
-    print(response["output"])
-print("Goodbye!")
+        if q.strip().lower() in {"exit", "quit"}:
+            break
+        response = agent_chain.run(q)
+        print(f"\nAnswer:\n{response}\n")