From 1ae78b521f2d89aefa5bb8e127b3bf44dfad9ec2 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 16:15:21 +0000 Subject: [PATCH] Update agent.py --- agent.py | 264 +++++++++++++++++++++---------------------------------- 1 file changed, 99 insertions(+), 165 deletions(-) diff --git a/agent.py b/agent.py index 56a3df6..db5c40b 100644 --- a/agent.py +++ b/agent.py @@ -1,187 +1,121 @@ -"""Deep Agent implementation based on LangGraph. +"""Deep Agent implementation based on LangGraph and LangChain. -This agent can: -1. Search the web for information using DuckDuckGo API. -2. Create virtual files in memory. -3. At the end of the run, persist virtual files to the real file system. +The module exposes two public functions: -The code follows the requirements: -- Uses langchain>=1.0.0 and langgraph>=1.0.0. -- Correct imports for text splitters, Chroma, Ollama embeddings and chat model. -- Implements `create_agent` and `create_agent_executor` functions. -- Uses the `write_file` tool to write virtual files. +* ``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. """ -import pathlib -from typing import Dict, List, Any, Optional +from __future__ import annotations -# LangChain imports -from langchain_ollama import ChatOllama, OllamaEmbeddings -from langchain_text_splitters import RecursiveCharacterTextSplitter -from langchain_chroma import Chroma +from typing import Any -# LangGraph imports -from langgraph.graph import Graph, State +# Lightweight imports – can be imported eagerly +from virtual_fs import VirtualFileSystem +from tools import WebSearch, CreateVirtualFile, ExportVirtualFiles -# Requests for web search -import requests +# Shared virtual file system instance used by all tools +vfs = VirtualFileSystem() -# ----------------------------- -# Configuration -# ----------------------------- -OLLAMA_MODEL = "llama3" -EMBEDDINGS_MODEL = "llama3" -SEARCH_URL = "https://api.duckduckgo.com/" -OUTPUT_DIR = pathlib.Path("output_files") -OUTPUT_DIR.mkdir(exist_ok=True) +# Define the tools +web_search_tool = WebSearch() +create_file_tool = CreateVirtualFile(vfs) +export_files_tool = ExportVirtualFiles(vfs) -# ----------------------------- -# Helper functions -# ----------------------------- +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- -def search_web(query: str) -> str: - """Return a short summary of the search results using DuckDuckGo.""" - params = { - "q": query, - "format": "json", - "no_html": 1, - "skip_disambig": 1, - } +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: - r = requests.get(SEARCH_URL, params=params, timeout=10) - r.raise_for_status() - data = r.json() - abstract = data.get("AbstractText") - if abstract: - return abstract - topics = data.get("RelatedTopics", []) - snippets = [t.get("Text", "") for t in topics if "Text" in t] - return "\n".join(snippets[:5]) - except Exception as e: - return f"Error during search: {e}" + from langgraph.graph import StateGraph + from langgraph.prebuilt import create_react_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 -# ----------------------------- -# State definition -# ----------------------------- -class AgentState(State): - query: str - virtual_files: Dict[str, str] - history: List[Dict[str, str]] - answer: Optional[str] = None + # 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() -# ----------------------------- -# Tool: write_file -# ----------------------------- - -def write_file_tool(state: AgentState, file_name: str, content: str) -> AgentState: - new_files = state.virtual_files.copy() - new_files[file_name] = content - return state.copy(update={"virtual_files": new_files}) - -# ----------------------------- -# Tool: search -# ----------------------------- - -def search_tool(state: AgentState, query: str) -> AgentState: - result = search_web(query) - new_history = state.history + [{"role": "tool", "name": "search", "content": result}] - return state.copy(update={"history": new_history}) - -# ----------------------------- -# Agent logic -# ----------------------------- - -def create_agent() -> Graph: - llm = ChatOllama(model=OLLAMA_MODEL, temperature=0.7) - embeddings = OllamaEmbeddings(model=EMBEDDINGS_MODEL) - splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50) - - def agent(state: AgentState) -> Dict[str, Any]: - # Build prompt from history - messages = [] - for msg in state.history: - if msg["role"] == "user": - messages.append({"role": "user", "content": msg["content"]}) - elif msg["role"] == "assistant": - messages.append({"role": "assistant", "content": msg["content"]}) - elif msg["role"] == "tool": - messages.append({"role": "assistant", "content": f"[Tool: {msg['name']}] {msg['content']}"}) - messages.append({"role": "assistant", "content": f"User query: {state.query}"}) - - response = llm.invoke(messages) - text = response.content.strip() - - if text.upper().startswith("ANSWER:"): - answer = text[7:].strip() - new_history = state.history + [{"role": "assistant", "content": answer}] - return {"final_answer": answer, "history": new_history} - elif text.upper().startswith("SEARCH:"): - query = text[7:].strip() - return {"search_query": query} - elif text.upper().startswith("WRITE:"): - try: - rest = text[6:].strip() - file_name, content = rest.split("|", 1) - return {"write_file": {"file_name": file_name.strip(), "content": content.strip()}} - except Exception: - return {"final_answer": "Could not parse WRITE command."} - else: - answer = text - new_history = state.history + [{"role": "assistant", "content": answer}] - return {"final_answer": answer, "history": new_history} - - graph = Graph() - graph.add_node("agent", agent) - graph.add_node("search", search_tool) - graph.add_node("write_file", write_file_tool) - - def final_answer_node(state: AgentState) -> AgentState: - for msg in reversed(state.history): - if msg["role"] == "assistant": - state = state.copy(update={"answer": msg["content"]}) - break - return state - - graph.add_node("final_answer", final_answer_node) - - graph.add_edge("agent", "search", condition=lambda out: "search_query" in out) - graph.add_edge("agent", "write_file", condition=lambda out: "write_file" in out) - graph.add_edge("agent", "final_answer", condition=lambda out: "final_answer" in out) - graph.add_edge("search", "agent", condition=lambda out: True) - graph.add_edge("write_file", "agent", condition=lambda out: True) - - graph.set_start("agent") - graph.set_end("final_answer") + # Build a simple react agent using LangGraph prebuilt + react_agent = create_react_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 -# ----------------------------- -# Executor helper -# ----------------------------- -def create_agent_executor() -> Graph: - return create_agent() +def create_agent_executor() -> Any: + """Return an :class:`langchain.agents.AgentExecutor`. -# ----------------------------- -# Main execution -# ----------------------------- -if __name__ == "__main__": - import argparse - parser = argparse.ArgumentParser(description="Run the deep agent.") - parser.add_argument("query", type=str, help="User query to process") - args = parser.parse_args() + 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 + 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 - graph = create_agent_executor() - initial_state = AgentState(query=args.query, virtual_files={}, history=[{"role": "user", "content": args.query}]) + 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() - final_state = graph.invoke(initial_state) + # Build the react agent + react_agent = create_react_agent( + llm=llm, + tools=[web_search_tool, create_file_tool, export_files_tool], + prompt=prompt, + output_parser=parser, + ) - print("\n=== Final Answer ===\n") - print(final_state.answer if final_state.answer else "No answer produced.") + return AgentExecutor.from_agent_and_tools( + agent=react_agent, + tools=[web_search_tool, create_file_tool, export_files_tool], + verbose=True, + ) - for fname, content in final_state.virtual_files.items(): - out_path = OUTPUT_DIR / fname - out_path.parent.mkdir(parents=True, exist_ok=True) - with open(out_path, "w", encoding="utf-8") as f: - f.write(content) - print(f"Virtual file written to {out_path}") +# Expose public names for tests +__all__ = ["create_agent", "create_agent_executor", "vfs"]