Update agent.py
This commit is contained in:
@@ -1,106 +1,72 @@
|
||||
"""Core logic for the RAG agent.
|
||||
"""
|
||||
Agent creation for the RAG system.
|
||||
|
||||
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.
|
||||
|
||||
The decision logic can be replaced with a more sophisticated router if
|
||||
desired.
|
||||
Provides a function ``create_agent`` that returns an ``AgentExecutor`` capable of
|
||||
choosing between the local KB search and the Tavily web search.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
from typing import List
|
||||
|
||||
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 langchain.agents import AgentExecutor, create_openai_functions_agent
|
||||
from langchain.tools import Tool
|
||||
|
||||
from rag_tools import search_local_kb, web_search
|
||||
from vectorstore import create_vectorstore
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Prompt template
|
||||
# ---------------------------------------------------------------------------
|
||||
SYSTEM_PROMPT = """You are an AI assistant that can answer questions using either a local knowledge base or real‑time web search.
|
||||
|
||||
When answering, always include the source of the information:
|
||||
- "chromadb" for local knowledge base results.
|
||||
- "tavily" for web search results.
|
||||
|
||||
If you are uncertain, say "I don't know" but still mention the source you used.
|
||||
"""
|
||||
|
||||
USER_PROMPT = """Question: {question}\n
|
||||
When you respond, first state the source (chromadb or tavily) and then provide the answer.
|
||||
"""
|
||||
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
("system", SYSTEM_PROMPT),
|
||||
("user", USER_PROMPT),
|
||||
])
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 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
|
||||
# Import the tools defined in tools.py
|
||||
from tools import search_local_kb, web_search
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent creation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def create_agent() -> AgentExecutor:
|
||||
"""Instantiate the agent with the two tools.
|
||||
def create_agent(vectorstore_instance) -> AgentExecutor:
|
||||
"""Create an agent that can decide between local KB and web search.
|
||||
|
||||
The LLM used is Ollama's ``llama3``.
|
||||
Parameters
|
||||
----------
|
||||
vectorstore_instance
|
||||
Instance of the Chroma vector store to be used by the local search tool.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AgentExecutor
|
||||
Configured agent ready for use.
|
||||
"""
|
||||
tools = [search_local_kb, web_search]
|
||||
# Make the vectorstore available to the tool via the module global
|
||||
import tools
|
||||
tools.vectorstore = vectorstore_instance
|
||||
|
||||
# Define the tools
|
||||
tools_list: List[Tool] = [
|
||||
Tool(
|
||||
name="search_local_kb",
|
||||
func=search_local_kb,
|
||||
description="Search the local knowledge base (ChromaDB). Use when the answer is likely contained in the local documents.",
|
||||
),
|
||||
Tool(
|
||||
name="web_search",
|
||||
func=web_search,
|
||||
description="Search the web via Tavily. Use when the answer requires up‑to‑date information.",
|
||||
),
|
||||
]
|
||||
|
||||
# LLM for the agent
|
||||
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)
|
||||
|
||||
# System prompt guiding the agent
|
||||
system_prompt = (
|
||||
"You are an assistant that answers user questions. "
|
||||
"If the answer can be found in the local knowledge base, use the tool "
|
||||
"`search_local_kb`. If the question asks for recent or current information, "
|
||||
"use the tool `web_search`. After obtaining the information, provide a "
|
||||
"concise answer and state the source (`chromadb` or `tavily`)."
|
||||
)
|
||||
|
||||
# Create the agent using the function calling approach
|
||||
agent = create_openai_functions_agent(llm=llm, tools=tools_list, system_message=system_prompt)
|
||||
|
||||
# Wrap in an executor for easy use
|
||||
return AgentExecutor(agent=agent, tools=tools_list, 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:
|
||||
try:
|
||||
question = input("\nЗапрос: ")
|
||||
except EOFError:
|
||||
break
|
||||
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
|
||||
# End of module
|
||||
# ---------------------------------------------------------------------------
|
||||
Reference in New Issue
Block a user