96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
import os
|
|
from dotenv import load_dotenv
|
|
from langchain_ollama import ChatOllama
|
|
from langchain.tools import tool
|
|
from langchain.agents import initialize_agent, AgentType
|
|
from langchain.schema import Document
|
|
from langchain.vectorstores import Chroma
|
|
|
|
# Global variable to hold the vector store for tool access
|
|
_vectorstore: Chroma | None = None
|
|
|
|
def set_vectorstore(vs: Chroma) -> None:
|
|
global _vectorstore
|
|
_vectorstore = vs
|
|
|
|
@tool
|
|
def search_local_kb(query: str, top_k: int = 3) -> str:
|
|
"""
|
|
Search the local knowledge base (ChromaDB) for relevant documents.
|
|
|
|
Args:
|
|
query (str): The search query.
|
|
top_k (int): Number of top documents to return.
|
|
|
|
Returns:
|
|
str: Concatenated content of the top documents or a not-found message.
|
|
"""
|
|
if _vectorstore is None:
|
|
return "Vector store not initialized."
|
|
|
|
retriever = _vectorstore.as_retriever(search_kwargs={"k": top_k})
|
|
docs = retriever.get_relevant_documents(query)
|
|
|
|
if not docs:
|
|
return "No relevant documents found in local knowledge base."
|
|
|
|
return "\n---\n".join(
|
|
[f"Document {i+1}:\n{doc.page_content}" for i, doc in enumerate(docs)]
|
|
)
|
|
|
|
@tool
|
|
def web_search(query: str) -> str:
|
|
"""
|
|
Perform a web search using Tavily.
|
|
|
|
Args:
|
|
query (str): The search query.
|
|
|
|
Returns:
|
|
str: Concatenated content of the top search results or a not-found message.
|
|
"""
|
|
from tavily import TavilyClient
|
|
|
|
api_key = os.getenv("TAVILY_API_KEY")
|
|
if not api_key:
|
|
return "TAVILY_API_KEY not set in environment."
|
|
|
|
client = TavilyClient(api_key=api_key)
|
|
try:
|
|
results = client.search(query, max_results=3)
|
|
except Exception as e:
|
|
return f"Web search failed: {e}"
|
|
|
|
if not results:
|
|
return "No results found on the web."
|
|
|
|
return "\n---\n".join(
|
|
[f"Result {i+1}:\n{res.get('content', '')}" for i, res in enumerate(results)]
|
|
)
|
|
|
|
def create_agent(vectorstore: Chroma):
|
|
"""
|
|
Create a LangChain agent that can route queries to either the local KB or the web.
|
|
|
|
Args:
|
|
vectorstore (Chroma): The vector store to use for local search.
|
|
|
|
Returns:
|
|
AgentExecutor: The configured agent.
|
|
"""
|
|
set_vectorstore(vectorstore)
|
|
|
|
llm = ChatOllama(model="llama3")
|
|
|
|
tools = [search_local_kb, web_search]
|
|
|
|
system_prompt = """
|
|
You are an AI assistant. When answering a question, decide whether the answer can be found in the local knowledge base or requires up-to-date information from the web. Use the tool search_local_kb if the answer is in the local knowledge base. Use web_search if the answer requires recent information. Always include the source of the information in your answer, either 'chromadb' or 'tavily'. Do not call both tools unless necessary. If you call a tool, the tool will return the content. Use that content to answer the question. Do not mention the tool usage in your answer. Just provide the answer and the source."""
|
|
agent = initialize_agent(
|
|
tools=tools,
|
|
llm=llm,
|
|
agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
|
|
verbose=False,
|
|
agent_kwargs={"system_message": system_prompt},
|
|
)
|
|
return agent |