100 lines
3.4 KiB
Python
100 lines
3.4 KiB
Python
"""Main script for the RAG agent with ChromaDB and Tavily.
|
||
|
||
The script:
|
||
1. Loads or creates the Chroma vector store.
|
||
2. Loads documents from the `documents/` folder.
|
||
3. Sets up the LangChain agent with two tools: `search_local_kb` and `web_search`.
|
||
4. Runs a simple CLI loop.
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from langchain.agents import initialize_agent, AgentType
|
||
from langchain.tools import Tool
|
||
|
||
from vectorstore import create_vectorstore, load_documents
|
||
from rag_tools import search_local_kb, web_search
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
CHROMA_DIR = "./chroma_db"
|
||
DOCS_DIR = "./documents"
|
||
MODEL = "llama3"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Helper: load or create vector store
|
||
# ---------------------------------------------------------------------------
|
||
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
|
||
|
||
# Load documents – we always load; Chroma will deduplicate by ID if same content
|
||
print("Loading documents into ChromaDB (if not already present)...")
|
||
load_documents(DOCS_DIR, vectorstore)
|
||
print("Documents loaded.")
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Define tools – pass the vectorstore to the local search tool
|
||
# ---------------------------------------------------------------------------
|
||
# We wrap the tool functions to include the vectorstore argument
|
||
|
||
def local_kb_tool(query: str, top_k: int = 3):
|
||
return search_local_kb(query=query, top_k=top_k, vectorstore=vectorstore)
|
||
|
||
# Create LangChain Tool objects
|
||
local_tool = Tool(
|
||
name="search_local_kb",
|
||
func=local_kb_tool,
|
||
description="Semantic search in the local knowledge base. Use when the answer is in the local documents.",
|
||
)
|
||
web_tool = Tool(
|
||
name="web_search",
|
||
func=web_search,
|
||
description="Search the web using Tavily. Use for up‑to‑date facts or news.",
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent setup
|
||
# ---------------------------------------------------------------------------
|
||
llm = ChatOllama(model=MODEL, temperature=0.0)
|
||
|
||
system_prompt = (
|
||
"You are an assistant that answers user questions. "
|
||
"If the answer is likely to be in the local knowledge base, use the tool "
|
||
"search_local_kb. If the answer requires up‑to‑date information, use the "
|
||
"web_search tool. After retrieving information, provide the answer and "
|
||
"state the source: either 'chromadb' or 'tavily'."
|
||
)
|
||
|
||
agent = initialize_agent(
|
||
tools=[local_tool, web_tool],
|
||
llm=llm,
|
||
agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
|
||
verbose=True,
|
||
prefix=system_prompt,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# CLI loop
|
||
# ---------------------------------------------------------------------------
|
||
print("\nRAG Agent ready. Type your question (or 'exit' to quit).\n")
|
||
while True:
|
||
try:
|
||
query = input("Query: ")
|
||
except (KeyboardInterrupt, EOFError):
|
||
print("\nExiting.")
|
||
break
|
||
if query.strip().lower() in {"exit", "quit", "q"}:
|
||
print("Exiting.")
|
||
break
|
||
if not query.strip():
|
||
continue
|
||
# Run the agent
|
||
try:
|
||
result = agent.run(query)
|
||
print(f"\nAnswer:\n{result}\n")
|
||
except Exception as e:
|
||
print(f"Error: {e}")
|
||
continue
|