Update agent.py

agent.py

@@ -1,150 +1,187 @@
-"""
+"""Deep Agent implementation based on LangGraph.
-Deep agent that can search the web, create virtual files and finally dump them to disk.
 
-The agent is built on top of LangChain 0.2+ and uses the "deep agents from scratch"
+This agent can:
-approach described in the course.  It is intentionally minimal but fully functional.
+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 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.
 """
 
-from __future__ import annotations
+import pathlib
+from typing import Dict, List, Any, Optional
 
-import os
+# LangChain imports
-from pathlib import Path
+from langchain_ollama import ChatOllama, OllamaEmbeddings
-from typing import Any, Dict
-
-from langchain_core.prompts import ChatPromptTemplate
-from langchain_core.output_parsers import StrOutputParser
-from langchain_core.runnables import Runnable
-from langchain_ollama import ChatOllama
 from langchain_text_splitters import RecursiveCharacterTextSplitter
 from langchain_chroma import Chroma
 
-# Local modules
+# LangGraph imports
-from virtual_fs import virtual_fs
+from langgraph.graph import Graph, State
 
-# ---------------------------------------------------------------------------
+# Requests for web search
-# 1. Tools
-# ---------------------------------------------------------------------------
-
-# 1.1 Web search tool – simple HTTP GET + title extraction
 import requests
-from bs4 import BeautifulSoup
 
+# -----------------------------
+# 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)
 
-def web_search(query: str) -> str:
+# -----------------------------
-    """Return the title and first paragraph of the first search result.
+# Helper functions
+# -----------------------------
 
-    This is a very small wrapper around a Google search.  For a production
+def search_web(query: str) -> str:
-    system you would use a real search API.
+    """Return a short summary of the search results using DuckDuckGo."""
-    """
+    params = {
-    # Simple Bing search URL – works without API key for a few requests
+        "q": query,
-    url = f"https://www.bing.com/search?q={requests.utils.quote(query)}"
+        "format": "json",
-    resp = requests.get(url, timeout=10)
+        "no_html": 1,
-    resp.raise_for_status()
+        "skip_disambig": 1,
-    soup = BeautifulSoup(resp.text, "html.parser")
+    }
-    results = soup.select("li.b_algo")
-    if not results:
-        return "No results found."
-    first = results[0]
-    title = first.select_one("h2").get_text(strip=True)
-    snippet = first.select_one("p").get_text(strip=True)
-    return f"Title: {title}\nSnippet: {snippet}"
-
-# 1.2 Write file tool
-
-def write_file_tool(path: str, content: str) -> str:
-    virtual_fs.write(path, content)
-    return f"File written to {path}."
-
-# 1.3 Read file tool
-
-def read_file_tool(path: str) -> str:
     try:
-        return virtual_fs.read(path)
+        r = requests.get(SEARCH_URL, params=params, timeout=10)
-    except KeyError:
+        r.raise_for_status()
-        return f"File {path} does not exist in virtual FS."
+        data = r.json()
-
+        abstract = data.get("AbstractText")
-# 1.4 Dump virtual FS to disk
+        if abstract:
-
+            return abstract
-def dump_virtual_fs_tool(output_dir: str = "output") -> str:
+        topics = data.get("RelatedTopics", [])
-    root = Path(output_dir)
+        snippets = [t.get("Text", "") for t in topics if "Text" in t]
-    virtual_fs.dump_to_disk(root)
+        return "\n".join(snippets[:5])
-    return f"Virtual FS dumped to {root.resolve()}"
-
-# ---------------------------------------------------------------------------
-# 2. Agent definition – deep agent style
-# ---------------------------------------------------------------------------
-
-# 2.1 LLM
-llm = ChatOllama(model="llama3.1", temperature=0.7)
-
-# 2.2 Prompt template – instruct the agent how to use tools
-prompt = ChatPromptTemplate.from_messages([
-    ("system", "You are a helpful assistant that can search the web, write files, read files and dump virtual files to disk.")
-])
-
-# 2.3 Tool mapping
-from langchain.tools import tool
-
-# Wrap tools with langchain Tool objects
-from langchain.tools import Tool
-
-search_tool = Tool(
-    name="WebSearch",
-    func=web_search,
-    description="Use this to search the web for information. Input should be a natural language query.",
-)
-write_tool = Tool(
-    name="WriteFile",
-    func=write_file_tool,
-    description="Write content to a file in the virtual file system. Input: path and content.",
-)
-read_tool = Tool(
-    name="ReadFile",
-    func=read_file_tool,
-    description="Read a file from the virtual file system. Input: path.",
-)
-dump_tool = Tool(
-    name="DumpVirtualFS",
-    func=dump_virtual_fs_tool,
-    description="Dump all virtual files to the real file system. Input: output directory (optional).",
-)
-
-# 2.4 Agent chain – simple chain that lets the LLM decide which tool to call
-from langchain.agents import AgentExecutor, ZeroShotAgent
-
-# Define the tool names and descriptions for the prompt
-tool_names = [search_tool.name, write_tool.name, read_tool.name, dump_tool.name]
-tool_descriptions = [t.description for t in [search_tool, write_tool, read_tool, dump_tool]]
-
-# Build the agent
-agent = ZeroShotAgent.from_llm_and_tools(
-    llm=llm,
-    tools=[search_tool, write_tool, read_tool, dump_tool],
-    prefix="You are a helpful assistant. Use the following tools when needed.",
-    suffix="When you are finished, output the final answer.",
-    tool_prompt="You can use the following tools: {tool_names}. {tool_descriptions}"
-)
-
-# Executor
-executor = AgentExecutor.from_agent_and_tools(
-    agent=agent,
-    tools=[search_tool, write_tool, read_tool, dump_tool],
-    verbose=True,
-)
-
-# ---------------------------------------------------------------------------
-# 3. Demo / entry point
-# ---------------------------------------------------------------------------
-
-if __name__ == "__main__":
-    print("Deep Agent Demo – type your question. Type 'exit' to quit.")
-    while True:
-        user_input = input("> ")
-        if user_input.lower() in {"exit", "quit"}:
-            print("Exiting…")
-            break
-        try:
-            result = executor.invoke({"input": user_input})
-            print("\nResult:\n", result)
     except Exception as e:
-            print("Error:", e)
+        return f"Error during search: {e}"
+
+# -----------------------------
+# State definition
+# -----------------------------
+class AgentState(State):
+    query: str
+    virtual_files: Dict[str, str]
+    history: List[Dict[str, str]]
+    answer: Optional[str] = None
+
+# -----------------------------
+# 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")
+
+    return graph
+
+# -----------------------------
+# Executor helper
+# -----------------------------
+
+def create_agent_executor() -> Graph:
+    return create_agent()
+
+# -----------------------------
+# 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()
+
+    graph = create_agent_executor()
+    initial_state = AgentState(query=args.query, virtual_files={}, history=[{"role": "user", "content": args.query}])
+
+    final_state = graph.invoke(initial_state)
+
+    print("\n=== Final Answer ===\n")
+    print(final_state.answer if final_state.answer else "No answer produced.")
+
+    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}")