Update agent.py
This commit is contained in:
@@ -1,128 +1,69 @@
|
||||
"""Agent construction for RAG with local KB and web search.
|
||||
"""Agent implementation using LangChain 1.x.
|
||||
|
||||
The agent uses LangChain's `create_openai_functions_agent` style with a custom
|
||||
`AgentExecutor` that routes queries to either the local knowledge base or the
|
||||
web search based on a simple heuristic: if the query contains words that
|
||||
suggest recent events ("news", "today", "now", "latest"), we use the web
|
||||
search; otherwise we use the local KB.
|
||||
The agent uses a simple tool‑based architecture. The LLM is a local
|
||||
`ChatOllama` model (llama3). Two tools are available:
|
||||
|
||||
The agent returns the answer along with the source used.
|
||||
* ``search_local_kb`` – semantic search in Qdrant.
|
||||
* ``web_search`` – web search via Tavily.
|
||||
|
||||
The system prompt instructs the LLM to decide which tool to use based on the
|
||||
question. The response always contains a marker indicating the source
|
||||
(`chromadb`/`tavily`). The marker is added by the LLM itself.
|
||||
"""
|
||||
|
||||
from typing import Dict, Any
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import List
|
||||
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_ollama import ChatOllama
|
||||
from langchain.agents import create_openai_functions_agent, AgentExecutor
|
||||
from langchain_core.prompts import ChatPromptTemplate
|
||||
from langchain_core.runnables import Runnable
|
||||
from langchain_core.tools import BaseTool
|
||||
|
||||
from rag_tools import search_local_kb, web_search
|
||||
from vectorstore import create_vectorstore
|
||||
# Import tools without relative import to allow top‑level import
|
||||
import tools
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helper: decide whether to use web search
|
||||
# Helper: create a tool list
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
WEB_KEYWORDS = {"news", "today", "now", "latest", "recent", "current"}
|
||||
|
||||
|
||||
def should_use_web(query: str) -> bool:
|
||||
words = set(query.lower().split())
|
||||
return bool(words & WEB_KEYWORDS)
|
||||
def get_tools() -> List[BaseTool]:
|
||||
"""Return the list of tools used by the agent."""
|
||||
return [tools.search_local_kb, tools.web_search]
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# System prompt
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
SYSTEM_PROMPT = "You are an assistant that answers questions. Use the local knowledge base when the question is about stored documents; otherwise use web search. Respond with the answer and indicate the source (chromadb or tavily)."
|
||||
SYSTEM_PROMPT = (
|
||||
"You are an AI assistant that can answer questions using two sources: "
|
||||
"1) a local knowledge base (Qdrant) and 2) the web via Tavily. "
|
||||
"If the answer can be found in the local KB, use the `search_local_kb` tool. "
|
||||
"If the answer requires up‑to‑date information, use the `web_search` tool. "
|
||||
"Return the answer followed by a source marker on a new line: "
|
||||
"`Source: chromadb` or `Source: tavily`."
|
||||
)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Agent construction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def create_agent(vectorstore) -> AgentExecutor:
|
||||
"""Create an AgentExecutor that routes queries to the appropriate tool.
|
||||
def create_agent() -> Runnable:
|
||||
"""Create a runnable agent.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
vectorstore: Chroma
|
||||
The vector store used by the local KB search tool.
|
||||
|
||||
Returns
|
||||
-------
|
||||
AgentExecutor
|
||||
Configured agent.
|
||||
The agent is a simple chain: system prompt → user message → tool calls → LLM
|
||||
response. It uses the default tool‑calling behaviour of LangChain 1.x.
|
||||
"""
|
||||
# Attach the vectorstore to the local search tool via closure
|
||||
def search_local_kb_wrapper(query: str, top_k: int = 3) -> str:
|
||||
return search_local_kb(query, top_k=top_k, vectorstore=vectorstore)
|
||||
llm = ChatOllama(model="llama3", temperature=0.0)
|
||||
|
||||
# Wrap the web search tool (no extra params needed)
|
||||
def web_search_wrapper(query: str, top_k: int = 3) -> str:
|
||||
return web_search(query, top_k=top_k)
|
||||
# AgentExecutor can accept a system message via agent_kwargs
|
||||
from langchain.agents import AgentExecutor
|
||||
|
||||
# Define tools list with updated references
|
||||
tools = [search_local_kb_wrapper, web_search_wrapper]
|
||||
agent = AgentExecutor.from_llm_and_tools(
|
||||
llm=llm,
|
||||
tools=get_tools(),
|
||||
verbose=True,
|
||||
# Provide the system prompt so the LLM knows how to behave
|
||||
agent_kwargs={"system_message": SYSTEM_PROMPT},
|
||||
)
|
||||
|
||||
# LLM
|
||||
llm = ChatOllama(model="llama3")
|
||||
|
||||
# Prompt template
|
||||
prompt = ChatPromptTemplate.from_messages([
|
||||
("system", SYSTEM_PROMPT),
|
||||
("human", "{input}"),
|
||||
])
|
||||
|
||||
# Create agent
|
||||
agent = create_openai_functions_agent(llm=llm, tools=tools, prompt=prompt)
|
||||
|
||||
# Wrap with AgentExecutor
|
||||
return AgentExecutor(agent=agent, tools=tools, verbose=True)
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Simple executor for tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
def create_agent_executor(vectorstore) -> Any:
|
||||
"""Return a simple callable that mimics the agent for testing.
|
||||
|
||||
The returned callable takes a dictionary with keys ``input`` and optionally
|
||||
``tool_choice``. It selects the appropriate tool based on the query and
|
||||
returns a dictionary with an ``output`` key containing the result.
|
||||
"""
|
||||
agent = create_agent(vectorstore)
|
||||
|
||||
def executor(inputs: Dict[str, Any]) -> Dict[str, Any]:
|
||||
query = inputs.get("input", "")
|
||||
tool_choice = inputs.get("tool_choice")
|
||||
if tool_choice == "web_search":
|
||||
result = web_search(query)
|
||||
elif tool_choice == "search_local_kb":
|
||||
result = search_local_kb(query, vectorstore=vectorstore)
|
||||
else:
|
||||
# Default routing
|
||||
if should_use_web(query):
|
||||
result = web_search(query)
|
||||
else:
|
||||
result = search_local_kb(query, vectorstore=vectorstore)
|
||||
return {"output": result}
|
||||
|
||||
return executor
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Example usage
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
if __name__ == "__main__":
|
||||
store = create_vectorstore()
|
||||
agent = create_agent(store)
|
||||
while True:
|
||||
q = input("Query> ")
|
||||
if q.lower() in {"exit", "quit", "q"}:
|
||||
break
|
||||
# Simple routing: if query contains web keywords, use web tool
|
||||
if should_use_web(q):
|
||||
result = agent.invoke({"input": q, "tool_choice": "web_search"})
|
||||
else:
|
||||
result = agent.invoke({"input": q, "tool_choice": "search_local_kb"})
|
||||
print("Answer:", result.get("output", ""))
|
||||
""
|
||||
return agent
|
||||
Reference in New Issue
Block a user