122 lines
4.1 KiB
Python
122 lines
4.1 KiB
Python
"""Deep Agent implementation based on LangGraph and LangChain.
|
||
|
||
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.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
from typing import Any
|
||
|
||
# Lightweight imports – can be imported eagerly
|
||
from virtual_fs import VirtualFileSystem
|
||
from tools import WebSearch, CreateVirtualFile, ExportVirtualFiles
|
||
|
||
# Shared virtual file system instance used by all tools
|
||
vfs = VirtualFileSystem()
|
||
|
||
# Define the tools
|
||
web_search_tool = WebSearch()
|
||
create_file_tool = CreateVirtualFile(vfs)
|
||
export_files_tool = ExportVirtualFiles(vfs)
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Public API
|
||
# ---------------------------------------------------------------------------
|
||
|
||
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.
|
||
"""
|
||
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],
|
||
verbose=True,
|
||
)
|
||
|
||
# Expose public names for tests
|
||
__all__ = ["create_agent", "create_agent_executor", "vfs"]
|