73 lines
2.7 KiB
Python
73 lines
2.7 KiB
Python
"""
|
||
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 OllamaLLM
|
||
from langchain_tavily import TavilySearchResults
|
||
from langchain.tools import tool
|
||
from langchain.agents import initialize_agent, AgentType
|
||
from langchain.chains import RetrievalQA
|
||
|
||
from vectorstore import create_vectorstore, load_documents
|
||
|
||
# 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 is missing in .env")
|
||
|
||
# 1. Vector store and retriever
|
||
vectorstore = create_vectorstore()
|
||
load_documents("documents", vectorstore)
|
||
retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
|
||
|
||
# 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}) 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 using Tavily.")
|
||
def web_search(query: str) -> str:
|
||
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])
|
||
|
||
# 3. Agent with simple routing logic
|
||
from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
|
||
from langchain.chat_models import ChatOllama
|
||
|
||
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,
|
||
)
|
||
|
||
# 4. CLI loop
|
||
if __name__ == "__main__":
|
||
print("RAG Agent ready. Type 'exit' to quit.")
|
||
while True:
|
||
try:
|
||
q = input("Query: ")
|
||
except EOFError:
|
||
break
|
||
if q.strip().lower() in {"exit", "quit"}:
|
||
break
|
||
response = agent_chain.run(q)
|
||
print(f"\nAnswer:\n{response}\n")
|