"""
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,
)