|
| 1 | +"""Assemble the deepagents graph for `assembly code`. |
| 2 | +
|
| 3 | +Wires the gateway model to deepagents' built-in coding toolset (filesystem + shell, |
| 4 | +rooted at the working directory via a `LocalShellBackend`), plus the custom `assembly` |
| 5 | +CLI tool and any MCP/docs tools, the installed-skills middleware, and human-in-the-loop |
| 6 | +approval on the mutating tools. The compiled graph is driven turn-by-turn from |
| 7 | +`session.py`; an `InMemorySaver` checkpointer gives both conversation memory and the |
| 8 | +interrupt/resume the approval flow needs. |
| 9 | +""" |
| 10 | + |
| 11 | +from __future__ import annotations |
| 12 | + |
| 13 | +from collections.abc import Mapping, Sequence |
| 14 | +from pathlib import Path |
| 15 | +from typing import TYPE_CHECKING, Protocol |
| 16 | + |
| 17 | +from aai_cli.code_agent.cli_tool import CLI_TOOL_NAME |
| 18 | +from aai_cli.code_agent.fetch_tool import FETCH_TOOL_NAME |
| 19 | +from aai_cli.code_agent.prompt import build_system_prompt |
| 20 | + |
| 21 | +if TYPE_CHECKING: |
| 22 | + from langchain.agents.middleware import AgentMiddleware |
| 23 | + from langchain_core.language_models.chat_models import BaseChatModel |
| 24 | + from langchain_core.tools import BaseTool |
| 25 | + from langgraph.checkpoint.base import BaseCheckpointSaver |
| 26 | + |
| 27 | +# The tools whose effects reach outside the model — file writes, edits, arbitrary |
| 28 | +# shell, the AssemblyAI CLI (which can spend account credits), and URL fetches (which |
| 29 | +# can reach internal/SSRF targets). Each is gated behind human approval unless the |
| 30 | +# session opts into --auto. |
| 31 | +MUTATING_TOOLS = ("write_file", "edit_file", "execute", CLI_TOOL_NAME, FETCH_TOOL_NAME) |
| 32 | + |
| 33 | + |
| 34 | +class CompiledAgent(Protocol): |
| 35 | + """The slice of the compiled langgraph graph the session drives. |
| 36 | +
|
| 37 | + A structural type so we needn't name langgraph's deeply-generic |
| 38 | + ``CompiledStateGraph`` (and don't drag its type params through our code). |
| 39 | + """ |
| 40 | + |
| 41 | + def invoke( |
| 42 | + self, input: object, config: Mapping[str, object] | None = None |
| 43 | + ) -> dict[str, object]: |
| 44 | + """Run one step of the graph, returning the updated state (incl. messages).""" |
| 45 | + |
| 46 | + |
| 47 | +def _interrupt_config(*, auto_approve: bool) -> dict[str, bool] | None: |
| 48 | + """The ``interrupt_on`` map: approve every mutating tool, or ``None`` under --auto.""" |
| 49 | + if auto_approve: |
| 50 | + return None |
| 51 | + return dict.fromkeys(MUTATING_TOOLS, True) |
| 52 | + |
| 53 | + |
| 54 | +def build_agent( |
| 55 | + *, |
| 56 | + model: BaseChatModel, |
| 57 | + root_dir: Path, |
| 58 | + tools: Sequence[BaseTool] = (), |
| 59 | + middlewares: Sequence[AgentMiddleware] = (), |
| 60 | + checkpointer: BaseCheckpointSaver | None = None, |
| 61 | + auto_approve: bool = False, |
| 62 | +) -> CompiledAgent: |
| 63 | + """Compile the coding agent over ``root_dir`` with ``tools`` and ``middlewares``. |
| 64 | +
|
| 65 | + ``model`` is the only network seam — tests pass a fake chat model so the real |
| 66 | + deepagents graph (filesystem + shell tools, approval, checkpointing) runs offline. |
| 67 | + ``checkpointer`` defaults to an in-memory saver (one ephemeral session); the command |
| 68 | + passes a SQLite saver for persistent, resumable sessions. |
| 69 | + """ |
| 70 | + from deepagents import create_deep_agent |
| 71 | + from deepagents.backends import LocalShellBackend |
| 72 | + from langgraph.checkpoint.memory import InMemorySaver |
| 73 | + |
| 74 | + # virtual_mode=True maps the model's "/"-rooted paths under root_dir and blocks |
| 75 | + # traversal escapes, so file ops and shell stay inside the working directory. |
| 76 | + backend = LocalShellBackend(root_dir=str(root_dir), virtual_mode=True) |
| 77 | + |
| 78 | + return create_deep_agent( |
| 79 | + model=model, |
| 80 | + backend=backend, |
| 81 | + system_prompt=build_system_prompt(str(root_dir)), |
| 82 | + tools=list(tools), |
| 83 | + middleware=list(middlewares), |
| 84 | + interrupt_on=_interrupt_config(auto_approve=auto_approve), |
| 85 | + checkpointer=checkpointer if checkpointer is not None else InMemorySaver(), |
| 86 | + ) |
0 commit comments