Update agent.py

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