|
| 1 | +from __future__ import annotations |
| 2 | + |
| 3 | +import os |
| 4 | +from pathlib import Path |
| 5 | +from typing import Any |
| 6 | + |
| 7 | +from deepagents import create_deep_agent |
| 8 | +from langchain_mcp_adapters.client import MultiServerMCPClient |
| 9 | +from langchain_mcp_adapters.tools import load_mcp_tools |
| 10 | + |
| 11 | + |
| 12 | +REPOSITORY_ROOT = Path(__file__).resolve().parents[4] |
| 13 | +CLI_PATH = REPOSITORY_ROOT / "packages/stagehand-codemode/dist/cli.js" |
| 14 | +SKILL_PATH = REPOSITORY_ROOT / "packages/stagehand-codemode/STAGEHAND_CODEMODE_SKILL.md" |
| 15 | +STAGEHAND_CODEMODE_SKILL = SKILL_PATH.read_text().strip() |
| 16 | + |
| 17 | + |
| 18 | +async def run_stagehand_agent(prompt: str, model: str | Any = "openai:gpt-5-mini") -> dict[str, Any]: |
| 19 | + client = MultiServerMCPClient( |
| 20 | + { |
| 21 | + "stagehand": { |
| 22 | + "transport": "stdio", |
| 23 | + "command": "node", |
| 24 | + "args": [str(CLI_PATH)], |
| 25 | + "env": dict(os.environ), |
| 26 | + } |
| 27 | + }, |
| 28 | + tool_name_prefix=False, |
| 29 | + handle_tool_errors=True, |
| 30 | + ) |
| 31 | + |
| 32 | + # An explicit session is required. client.get_tools() would start a fresh stdio |
| 33 | + # process for each tool call and lose the browser between calls. |
| 34 | + async with client.session("stagehand") as session: |
| 35 | + tools = await load_mcp_tools(session) |
| 36 | + code_tools = [tool for tool in tools if tool.name == "code_execute"] |
| 37 | + if len(code_tools) != 1: |
| 38 | + raise RuntimeError(f"Expected one code_execute tool, got {[tool.name for tool in tools]}") |
| 39 | + agent = create_deep_agent( |
| 40 | + model=model, |
| 41 | + tools=code_tools, |
| 42 | + system_prompt=STAGEHAND_CODEMODE_SKILL, |
| 43 | + ) |
| 44 | + return await agent.ainvoke( |
| 45 | + {"messages": [{"role": "user", "content": prompt}]}, |
| 46 | + config={"recursion_limit": 20}, |
| 47 | + ) |
0 commit comments