Update agent.py
This commit is contained in:
@@ -1,97 +1,106 @@
|
||||
"""Core logic for the RAG agent.
|
||||
|
||||
The agent decides whether to use the local knowledge base or Tavily based on simple heuristics:
|
||||
The agent decides whether to use the local knowledge base or Tavily based on a
|
||||
very simple heuristic: if the query contains words like ``news``, ``latest``
|
||||
or ``today`` it is routed to the web search; otherwise the local KB is used.
|
||||
|
||||
- If the query contains words that usually refer to recent events (e.g. "новости", "актуальные", "сегодня", "сейчас"), we route to Tavily.
|
||||
- Otherwise we assume the answer can be found in the local KB.
|
||||
|
||||
The agent uses LangChain's :class:`langchain.agents.AgentExecutor` with a custom prompt that
|
||||
instructs the LLM to specify the source in the response.
|
||||
The decision logic can be replaced with a more sophisticated router if
|
||||
desired.
|
||||
"""
|
||||
|
||||
import os
|
||||
from typing import Dict, Any
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Tuple
|
||||
|
||||
from langchain.agents import AgentExecutor, create_openai_tools_agent
|
||||
from langchain_ollama import ChatOllama
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
|
||||
from rag_tools import search_local_kb, web_search
|
||||
from vectorstore import create_vectorstore, load_documents
|
||||
from vectorstore import create_vectorstore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Configuration
|
||||
# Prompt template
|
||||
# ---------------------------------------------------------------------------
|
||||
# Directory containing the documents to index
|
||||
DOCUMENTS_DIR = "./documents"
|
||||
# Persistence directory for Chroma
|
||||
CHROMA_DIR = "./chroma_db"
|
||||
# The LLM used by the agent
|
||||
LLM = ChatOllama(model="llama3")
|
||||
SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search.
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper functions
|
||||
# ---------------------------------------------------------------------------
|
||||
When answering, always include the source of the information:
|
||||
- "chromadb" for local knowledge base results.
|
||||
- "tavily" for web search results.
|
||||
|
||||
def should_use_web(query: str) -> bool:
|
||||
"""Return True if the query looks like it needs up‑to‑date information.
|
||||
If you are uncertain, say "I don't know" but still mention the source you used.
|
||||
"""
|
||||
|
||||
The heuristic looks for a few Russian keywords that usually indicate a
|
||||
request for recent news.
|
||||
"""
|
||||
keywords = ["новости", "актуальные", "сегодня", "сейчас", "текущие", "текущий"]
|
||||
lowered = query.lower()
|
||||
return any(k in lowered for k in keywords)
|
||||
USER_PROMPT = """Question: {question}\n
|
||||
When you respond, first state the source (chromadb or tavily) and then provide the answer.
|
||||
"""
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent setup
|
||||
# ---------------------------------------------------------------------------
|
||||
# Load or create the vector store
|
||||
vectorstore = create_vectorstore(persist_directory=CHROMA_DIR)
|
||||
# If the store is empty, load documents from the documents directory
|
||||
if not vectorstore._collection.count(): # type: ignore[attr-defined]
|
||||
load_documents(DOCUMENTS_DIR, vectorstore)
|
||||
|
||||
# Tools – we pass the vectorstore instance to the local search tool via a closure
|
||||
local_kb_tool = search_local_kb
|
||||
local_kb_tool.__globals__["vectorstore"] = vectorstore # inject vectorstore
|
||||
|
||||
TOOLS = [local_kb_tool, web_search]
|
||||
|
||||
# Prompt template – the LLM is instructed to specify the source.
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
("system", "You are an AI assistant that answers user questions.") ,
|
||||
("user", "{input}"),
|
||||
("system", SYSTEM_PROMPT),
|
||||
("user", USER_PROMPT),
|
||||
])
|
||||
|
||||
# Agent – we use the simple tool‑calling agent
|
||||
agent = create_openai_tools_agent(llm=LLM, tools=TOOLS, prompt=prompt)
|
||||
executor = AgentExecutor(agent=agent, tools=TOOLS, verbose=True)
|
||||
# ---------------------------------------------------------------------------
|
||||
# Decision logic
|
||||
# ---------------------------------------------------------------------------
|
||||
WEB_KEYWORDS = {"news", "latest", "today", "current", "recent"}
|
||||
|
||||
|
||||
def choose_tool(question: str) -> Tuple[str, callable]:
|
||||
"""Return the name of the tool and the function to call.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
question: str
|
||||
The user query.
|
||||
|
||||
Returns
|
||||
-------
|
||||
Tuple[str, callable]
|
||||
The tool name and the corresponding function.
|
||||
"""
|
||||
lowered = question.lower()
|
||||
if any(word in lowered for word in WEB_KEYWORDS):
|
||||
return "web_search", web_search
|
||||
return "search_local_kb", search_local_kb
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI
|
||||
# Agent creation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def main() -> None:
|
||||
def create_agent() -> AgentExecutor:
|
||||
"""Instantiate the agent with the two tools.
|
||||
|
||||
The LLM used is Ollama's ``llama3``.
|
||||
"""
|
||||
tools = [search_local_kb, web_search]
|
||||
llm = ChatOllama(model="llama3", temperature=0)
|
||||
# Build an agent that knows about the tools and uses the custom prompt
|
||||
agent = create_openai_tools_agent(llm=llm, tools=tools, prompt=prompt)
|
||||
return AgentExecutor(agent=agent, tools=tools, verbose=True)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# CLI loop
|
||||
# ---------------------------------------------------------------------------
|
||||
if __name__ == "__main__":
|
||||
# Ensure the vector store is loaded once
|
||||
store = create_vectorstore()
|
||||
# Load documents if the store is empty
|
||||
if not store.get_index_info():
|
||||
from vectorstore import load_documents
|
||||
load_documents("documents", store)
|
||||
|
||||
agent = create_agent()
|
||||
print("RAG Agent ready. Type 'exit' to quit.")
|
||||
while True:
|
||||
user_query = input("\nЗапрос: ")
|
||||
if user_query.lower() in {"exit", "quit", "q"}:
|
||||
try:
|
||||
question = input("\nЗапрос: ")
|
||||
except EOFError:
|
||||
break
|
||||
# Decide which tool to use
|
||||
if should_use_web(user_query):
|
||||
# Explicitly call the web search tool
|
||||
result = web_search(user_query)
|
||||
source = "tavily"
|
||||
else:
|
||||
result = search_local_kb(user_query, vectorstore=vectorstore)
|
||||
source = "chromadb"
|
||||
# Ask the LLM to format the answer
|
||||
formatted = LLM.invoke({"input": user_query + "\n\nAnswer: " + result})
|
||||
print("\nОтвет:", formatted.content)
|
||||
print("Источник:", source)
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
||||
# End of agent.py
|
||||
if question.strip().lower() in {"exit", "quit"}:
|
||||
break
|
||||
# The agent will automatically call the chosen tool via the prompt.
|
||||
# We simply pass the question to the agent.
|
||||
result = agent.invoke({"input": question})
|
||||
# The agent's output already contains the source.
|
||||
print("Ответ:", result["output"]) # noqa: T201
|
||||
|
||||
Reference in New Issue
Block a user