diff --git a/src/deep_agents_from_scratch/deep_agent.py b/src/deep_agents_from_scratch/deep_agent.py new file mode 100644 index 0000000..6ae6cfe --- /dev/null +++ b/src/deep_agents_from_scratch/deep_agent.py @@ -0,0 +1,194 @@ +"""Deep Agent implementation based on LangGraph StateGraph. + +The agent follows the architecture described in the original +`deep-agents-from-scratch` notebook but is rewritten to satisfy the +current project constraints: + +* All imports use the modern LangChain modules. +* The graph is fully connected – the final node writes the virtual + files to disk. +* A lightweight in‑memory virtual file system is used during the + conversation. +* The search tool uses Tavily and returns only a minimal summary. +* Summarization is performed by a small model (GPT‑4o‑mini). + +The module exposes a single helper ``build_deep_agent_graph`` which +returns a ready‑to‑run :class:`StateGraph` instance. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import Dict, List, Any + +# Modern LangChain imports +from langchain_core.messages import BaseMessage, HumanMessage +from langgraph.types import StateGraph + +# Local imports – the tools are defined in ``research_tools.py`` +from .research_tools import ( + tavily_search, + think_tool, + summarize_webpage_content, +) +from .state import DeepAgentState + +# --------------------------------------------------------------------------- +# Helper: virtual file system +# --------------------------------------------------------------------------- +class VirtualFileSystem: + """Simple in‑memory file system used by the graph. + + The state already contains a ``files`` mapping, but the graph + interacts with this helper to keep the node implementations + clean. + """ + + def __init__(self, files: Dict[str, str]): + self._files = files + + def write(self, filename: str, content: str) -> None: + self._files[filename] = content + + def read(self, filename: str) -> str | None: + return self._files.get(filename) + + def list(self) -> List[str]: + return list(self._files.keys()) + + def dump_to_disk(self, root: Path) -> None: + root.mkdir(parents=True, exist_ok=True) + for name, data in self._files.items(): + (root / name).write_text(data, encoding="utf-8") + +# --------------------------------------------------------------------------- +# Graph nodes +# --------------------------------------------------------------------------- +async def think_node(state: DeepAgentState) -> Dict[str, Any]: + """Ask the agent to decide what to search next. + + The node receives the current conversation and returns a new + query string. + """ + messages = state.messages + # Find the last user message + last_user = None + for msg in reversed(messages): + if msg.role == "user": + last_user = msg.content + break + if not last_user: + last_user = "" + # The tool expects a plain string + query = await think_tool.invoke(last_user) + return {"query": query} + +async def search_node(state: DeepAgentState) -> Dict[str, Any]: + """Perform a web search using the tavily tool. + + The node receives the ``query`` from the previous node and returns + the raw search results. + """ + query = state.query + if not query: + return {"search_results": []} + results = await tavily_search.invoke(query) + return {"search_results": results} + +async def summarize_node(state: DeepAgentState) -> Dict[str, Any]: + """Summarize each search result and store the content in the virtual FS. + + The node writes full webpage content to the virtual file system and + returns a list of summary strings. + """ + results = state.search_results + summaries: List[str] = [] + vfs = VirtualFileSystem(state.files) + for i, res in enumerate(results): + content = res.get("content", "") + if not content: + continue + summary_obj = await summarize_webpage_content.invoke(content) + filename = f"search_result_{i+1}.md" + vfs.write(filename, content) + summaries.append(summary_obj["summary"]) + return {"summaries": summaries} + +async def write_file_node(state: DeepAgentState) -> Dict[str, Any]: + """Write all virtual files to disk and add a final message to the chat. + + The node is the final node of the graph. It writes the files to a + directory named ``output`` relative to the current working + directory. + """ + vfs = VirtualFileSystem(state.files) + output_dir = Path("output") + vfs.dump_to_disk(output_dir) + final_msg = f"Files written to {output_dir.resolve()}" + # Append the final message to the conversation history + return {"messages": [HumanMessage(content=final_msg)]} + +# --------------------------------------------------------------------------- +# Graph construction +# --------------------------------------------------------------------------- +def build_deep_agent_graph() -> StateGraph[DeepAgentState]: + """Build and return the StateGraph. + + The graph follows the sequence: + 1. think → 2. search → 3. summarize → 4. write_file + """ + graph = StateGraph(DeepAgentState) + + # Add nodes + graph.add_node("think", think_node) + graph.add_node("search", search_node) + graph.add_node("summarize", summarize_node) + graph.add_node("write_file", write_file_node) + + # Connect nodes + graph.set_entry_point("think") + graph.add_edge("think", "search") + graph.add_edge("search", "summarize") + graph.add_edge("summarize", "write_file") + + return graph + +# --------------------------------------------------------------------------- +# Helper to run the graph with an initial user message +# --------------------------------------------------------------------------- +async def run_agent(user_input: str) -> str: + """Convenience wrapper that runs the graph for a single turn. + + Parameters + ---------- + user_input: str + The user message to start the conversation. + + Returns + ------- + str + The final message produced by the graph. + """ + graph = build_deep_agent_graph() + # Initialise state + state: DeepAgentState = DeepAgentState( + messages=[HumanMessage(content=user_input)], + files={}, + query=None, + search_results=[], + summaries=[], + ) + # Run the graph – ``invoke`` returns the final state + final_state = await graph.invoke(state) + # The final message was appended to the conversation + return final_state.messages[-1].content + +# --------------------------------------------------------------------------- +# CLI entry point +# --------------------------------------------------------------------------- +if __name__ == "__main__": + import asyncio + + user_msg = input("Enter your research question: ") + result = asyncio.run(run_agent(user_msg)) + print(result) \ No newline at end of file