import os
from langchain_community.chat_models import ChatOllama
from langchain.tools import tool
from deepagents import create_deep_agent
from deepagents.backends import FilesystemBackend, LocalShellBackend, CompositeBackend
from langchain_core.messages import HumanMessage

from tools import search_local_kb, web_search

llm = ChatOllama(
    model="llama3",
    base_url=os.getenv("OLLAMA_HOST", "http://localhost:11434"),
)

backend = CompositeBackend(
    [
        LocalShellBackend(workspace_dir="./workspace"),
        FilesystemBackend(),
    ]
)

agent = create_deep_agent(
    model=llm,
    tools=[search_local_kb, web_search],
    backend=backend,
    system_prompt=(
        "You are a helpful agent. "
        "Use search_local_kb for local knowledge and web_search for up-to-date info. "
        "Prefix answers with [Local KB] or [Web Search] and indicate source."
    ),
)

async def run_agent(query: str) -> str:
    result = await agent.ainvoke(
        {"messages": [HumanMessage(content=query)]},
        {"configurable": {"thread_id": "session-1"}},
    )
    return result["messages"][-1].content