diff --git a/agent.py b/agent.py new file mode 100644 index 0000000..8f705ac --- /dev/null +++ b/agent.py @@ -0,0 +1,38 @@ +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 \ No newline at end of file