99 lines
3.0 KiB
Python
99 lines
3.0 KiB
Python
import os
|
|
from typing import Any, Dict
|
|
|
|
from langchain_ollama import Ollama
|
|
from langchain.agents import Tool, AgentExecutor, initialize_agent, AgentType
|
|
from langchain.schema import AgentAction, AgentFinish
|
|
from langchain.tools import tool
|
|
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
|
from langchain_core.output_parsers import StrOutputParser
|
|
|
|
from vectorstore import QdrantVectorStore
|
|
from tools import search_local_kb, web_search
|
|
|
|
|
|
def create_agent(
|
|
vectorstore: QdrantVectorStore,
|
|
tavily_api_key: str,
|
|
llm_model: str = "llama3",
|
|
) -> AgentExecutor:
|
|
"""
|
|
Create a LangChain agent that routes queries to either the local KB or the web.
|
|
|
|
Parameters
|
|
----------
|
|
vectorstore : QdrantVectorStore
|
|
The vector store for local knowledge base.
|
|
tavily_api_key : str
|
|
Tavily API key.
|
|
llm_model : str
|
|
Ollama model name for LLM.
|
|
|
|
Returns
|
|
-------
|
|
AgentExecutor
|
|
Configured agent executor.
|
|
"""
|
|
# LLM
|
|
llm = Ollama(model=llm_model)
|
|
|
|
# Define tools with partial application of required arguments
|
|
tools = [
|
|
Tool(
|
|
name="search_local_kb",
|
|
func=lambda query, top_k=3: search_local_kb(
|
|
query=query, top_k=top_k, vectorstore=vectorstore
|
|
),
|
|
description=(
|
|
"Use this tool to search the local knowledge base. "
|
|
"Return the most relevant snippets."
|
|
),
|
|
),
|
|
Tool(
|
|
name="web_search",
|
|
func=lambda query, max_results=3: web_search(
|
|
query=query, tavily_api_key=tavily_api_key, max_results=max_results
|
|
),
|
|
description=(
|
|
"Use this tool to search the web via Tavily. "
|
|
"Return the most relevant snippets."
|
|
),
|
|
),
|
|
]
|
|
|
|
# Prompt template for the agent
|
|
system_prompt = (
|
|
"You are an AI assistant that can answer questions using either "
|
|
"the local knowledge base or the web. If the question is about "
|
|
"recent events, news, or requires up-to-date information, "
|
|
"use the web_search tool. If the question is about "
|
|
"information contained in the local documents, use the "
|
|
"search_local_kb tool. After retrieving the information, "
|
|
"provide a concise answer and state the source (chromadb or tavily)."
|
|
)
|
|
|
|
prompt = ChatPromptTemplate.from_messages(
|
|
[
|
|
("system", system_prompt),
|
|
MessagesPlaceholder("history"),
|
|
("human", "{input}"),
|
|
MessagesPlaceholder("agent_scratchpad"),
|
|
]
|
|
)
|
|
|
|
# Output parser
|
|
output_parser = StrOutputParser()
|
|
|
|
# Agent
|
|
agent = initialize_agent(
|
|
tools=tools,
|
|
llm=llm,
|
|
agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
|
|
verbose=False,
|
|
handle_parsing_errors=True,
|
|
max_iterations=5,
|
|
early_stopping_method="generate",
|
|
system_message=system_prompt,
|
|
)
|
|
|
|
return agent |