From 156661da05c05b01ac2fb8d81d006727a8366eca Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=D0=9A=D0=B8=D1=80=D0=B8=D0=BB=D0=BB=20=D0=A0=D0=BE=D0=BC?= =?UTF-8?q?=D0=B0=D0=BD=D0=BE=D0=B2?= Date: Thu, 4 Jun 2026 19:18:51 +0000 Subject: [PATCH] Update agent.py --- agent.py | 146 +++++++++++++++---------------------------------------- 1 file changed, 38 insertions(+), 108 deletions(-) diff --git a/agent.py b/agent.py index 8e2c086..0767dea 100644 --- a/agent.py +++ b/agent.py @@ -1,121 +1,51 @@ -"""Deep Agent implementation based on LangGraph and LangChain. +""" +Agent creation using LangChain 1.x `create_agent` API. -The module exposes two public functions: - -* ``create_agent`` – returns a :class:`langgraph.graph.StateGraph` that can be - executed. -* ``create_agent_executor`` – returns an :class:`langchain.agents.AgentExecutor` - that can be used directly. - -Both functions lazily import the heavy LangChain/LangGraph dependencies so -that the module can be imported even if those packages are not installed. +This module exposes a single helper `create_agent_executor` that builds an +AgentExecutor configured with the custom tools defined in :mod:`tools`. """ -from __future__ import annotations +from langchain_ollama import ChatOllama +from langchain.agents import create_agent, AgentExecutor +from langchain.tools import BaseTool +from typing import List -from typing import Any +from .tools import search, write_file -# Lightweight imports – can be imported eagerly -from virtual_fs import VirtualFileSystem -from tools import WebSearch, CreateVirtualFile, ExportVirtualFiles +# Define the list of tools that the agent can use +TOOLS: List[BaseTool] = [search, write_file] -# Shared virtual file system instance used by all tools -vfs = VirtualFileSystem() +# LLM configuration – Ollama local model +LLM = ChatOllama(model="llama3.1:latest", temperature=0.0) -# Define the tools -web_search_tool = WebSearch() -create_file_tool = CreateVirtualFile(vfs) -export_files_tool = ExportVirtualFiles(vfs) +# Create the agent executor using the new LangChain 1.x API +# The `create_agent` function returns an AgentExecutor instance +# that can be called like a normal function. -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- +def create_agent_executor() -> AgentExecutor: + """Instantiate and return an AgentExecutor. -def create_agent() -> Any: - """Return a LangGraph graph. - - The function performs a lazy import of :mod:`langgraph` and related - dependencies. If the imports fail, a clear ``ImportError`` is raised. + The executor is configured with: + * The Ollama Chat model. + * The custom tools defined in :mod:`tools`. + * The default agent type "openai-retrieval-qa" is not used – we rely on + the automatically selected agent type by `create_agent`. """ - try: - from langgraph.graph import StateGraph - from langgraph.prebuilt import create_react_agent as create_agent - from langchain_ollama import ChatOllama - from langchain_core.prompts import ChatPromptTemplate - from langchain_core.output_parsers import StrOutputParser - except Exception as exc: # pragma: no cover - raise ImportError( - "Failed to import LangGraph/LangChain dependencies. Ensure that the\n" - "required packages are installed and the environment is correctly\n" - "configured." - ) from exc - - # LLM and prompt - llm = ChatOllama(model="llama3.1") - prompt = ChatPromptTemplate.from_messages( - [ - ("system", "You are a helpful assistant that can search the web, create virtual files, and export them to disk."), - ("human", "{input}"), - ] - ) - parser = StrOutputParser() - - # Build a simple agent using LangGraph prebuilt - react_agent = create_agent( - llm=llm, - tools=[web_search_tool, create_file_tool, export_files_tool], - prompt=prompt, - output_parser=parser, - ) - - # Create the graph - graph = StateGraph() - graph.add_node("react_agent", react_agent) - graph.set_entry_point("react_agent") - graph.set_finish_point("react_agent") - return graph - - -def create_agent_executor() -> Any: - """Return an :class:`langchain.agents.AgentExecutor`. - - The function lazily imports the required classes. It is useful for - quick experimentation and for environments where the full graph is - unnecessary. - """ - try: - from langchain.agents import AgentExecutor - from langgraph.prebuilt import create_react_agent as create_agent - from langchain_ollama import ChatOllama - from langchain_core.prompts import ChatPromptTemplate - from langchain_core.output_parsers import StrOutputParser - except Exception as exc: # pragma: no cover - raise ImportError( - "Failed to import LangChain dependencies for the AgentExecutor." - ) from exc - - llm = ChatOllama(model="llama3.1") - prompt = ChatPromptTemplate.from_messages( - [ - ("system", "You are a helpful assistant that can search the web, create virtual files, and export them to disk."), - ("human", "{input}"), - ] - ) - parser = StrOutputParser() - - # Build the agent - react_agent = create_agent( - llm=llm, - tools=[web_search_tool, create_file_tool, export_files_tool], - prompt=prompt, - output_parser=parser, - ) - - return AgentExecutor.from_agent_and_tools( - agent=react_agent, - tools=[web_search_tool, create_file_tool, export_files_tool], + agent = create_agent( + llm=LLM, + tools=TOOLS, verbose=True, ) + # The returned object is already an AgentExecutor + return agent -# Expose public names for tests -__all__ = ["create_agent", "create_agent_executor", "vfs"] +# Expose the executor for external use +agent_executor = create_agent_executor() + +# For convenience, a small helper that runs a single prompt + +def run_prompt(prompt: str) -> str: + """Run the prompt through the agent and return the final answer.""" + result = agent_executor.invoke({"input": prompt}) + # The result is a dict with keys: "output" and possibly others + return result.get("output", "") \ No newline at end of file