62 lines
2.1 KiB
Python
62 lines
2.1 KiB
Python
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.tools import tool
|
|
from langchain.schema import HumanMessage
|
|
|
|
from vectorstore import create_vectorstore, load_documents
|
|
|
|
# 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)
|
|
|
|
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])
|
|
|
|
@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]
|
|
|
|
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!")
|