33 lines
1.0 KiB
Python
33 lines
1.0 KiB
Python
"""
|
||
Create a LangChain agent that chooses between local KB and web search.
|
||
"""
|
||
|
||
from typing import List, Dict
|
||
|
||
from langchain.agents import create_agent
|
||
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
|
||
from langchain.schema.document import Document
|
||
|
||
# Import tools
|
||
from tools import search_local_kb, web_search
|
||
|
||
SYSTEM_PROMPT = """
|
||
You are an assistant that can answer questions using either a local knowledge base or the web. Use the tool `search_local_kb` when the question is about content already in your documents. Use `web_search` for up‑to‑date facts.
|
||
When you provide an answer, include the source: either "chromadb" or "tavily".
|
||
"""
|
||
|
||
# Prompt template with tool messages
|
||
prompt = ChatPromptTemplate.from_messages([
|
||
("system", SYSTEM_PROMPT),
|
||
MessagesPlaceholder("history"),
|
||
("human", "{input}"),
|
||
])
|
||
|
||
# Create agent harness
|
||
agent = create_agent(
|
||
model="ollama:llama3",
|
||
tools=[search_local_kb, web_search],
|
||
system_prompt=SYSTEM_PROMPT,
|
||
prompt_template=prompt,
|
||
)
|