72 lines
2.5 KiB
Python
72 lines
2.5 KiB
Python
"""
|
||
Agent creation for the RAG system.
|
||
|
||
Provides a function ``create_agent`` that returns an ``AgentExecutor`` capable of
|
||
choosing between the local KB search and the Tavily web search.
|
||
"""
|
||
|
||
from typing import List
|
||
|
||
from langchain_ollama import ChatOllama
|
||
from langchain.agents import AgentExecutor, create_openai_functions_agent
|
||
from langchain.tools import Tool
|
||
|
||
# Import the tools defined in tools.py
|
||
from tools import search_local_kb, web_search
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Agent creation
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def create_agent(vectorstore_instance) -> AgentExecutor:
|
||
"""Create an agent that can decide between local KB and web search.
|
||
|
||
Parameters
|
||
----------
|
||
vectorstore_instance
|
||
Instance of the Chroma vector store to be used by the local search tool.
|
||
|
||
Returns
|
||
-------
|
||
AgentExecutor
|
||
Configured agent ready for use.
|
||
"""
|
||
# 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)
|
||
|
||
# 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)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# End of module
|
||
# --------------------------------------------------------------------------- |