82 lines
2.9 KiB
Python
82 lines
2.9 KiB
Python
"""RAG agent CLI.
|
||
|
||
This script demonstrates a simple chat loop with a LangChain agent that
|
||
searches either a local Chroma vector store or the web via Tavily. The
|
||
agent automatically decides which tool to use based on the user query.
|
||
|
||
Prerequisites:
|
||
* Ollama must be running locally with the ``llama3`` model and the
|
||
``nomic-embed-text`` embedding model.
|
||
* A valid Tavily API key must be set in the environment variable
|
||
``TAVILY_API_KEY``.
|
||
* The ``documents`` directory should contain the source text files.
|
||
"""
|
||
|
||
import os
|
||
from pathlib import Path
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from langchain.agents import create_agent
|
||
|
||
from vectorstore import create_vectorstore, load_documents
|
||
from tools import search_local_kb, web_search
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Configuration
|
||
# ---------------------------------------------------------------------------
|
||
CHROMA_DIR = "./chroma_db"
|
||
DOCS_DIR = "./documents"
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Initialise vector store and load documents
|
||
# ---------------------------------------------------------------------------
|
||
print("Initializing Chroma vector store…")
|
||
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
|
||
print("Loading documents…")
|
||
load_documents(DOCS_DIR, vectorstore)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent setup
|
||
# ---------------------------------------------------------------------------
|
||
# System prompt that tells the model how to choose a tool.
|
||
SYSTEM_PROMPT = (
|
||
"You are an AI assistant that can answer questions using two tools. "
|
||
"If the answer requires up‑to‑date information, use the web_search tool. "
|
||
"Otherwise, use the search_local_kb tool. "
|
||
"When you call a tool, the tool will return the answer. "
|
||
"Respond with the final answer and include the source tag (chromadb or tavily)."
|
||
)
|
||
|
||
llm = ChatOllama(model="llama3", temperature=0)
|
||
|
||
# Tools list
|
||
TOOLS = [search_local_kb, web_search]
|
||
|
||
agent = create_agent(
|
||
model=llm,
|
||
tools=TOOLS,
|
||
system_prompt=SYSTEM_PROMPT,
|
||
)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Chat loop
|
||
# ---------------------------------------------------------------------------
|
||
print("\n--- RAG Agent CLI ---")
|
||
print("Type 'exit' or 'quit' to end.")
|
||
while True:
|
||
user_input = input("\nUser: ")
|
||
if user_input.lower() in {"exit", "quit", "q"}:
|
||
print("Goodbye!")
|
||
break
|
||
# Invoke the agent
|
||
try:
|
||
response = agent.invoke({"messages": [{"role": "user", "content": user_input}]})
|
||
# The response is a dict with a "messages" key
|
||
assistant_msg = next(
|
||
m for m in response["messages"] if m["role"] == "assistant"
|
||
)
|
||
print("\nAssistant:", assistant_msg["content"].strip())
|
||
except Exception as e:
|
||
print("Error:", e)
|
||
|
||
"""End of main.py""" |