quickstart.py

"""最小验证脚本：跑通 LangGraph + 当前 RUNNER_MODE + LangSmith trace 链路。

按 .env 里的 RUNNER_MODE 自动选 runner，节点内部委托给 runner.run_task。
任务：用 trimesh 写脚本生成长方体，存到 workspace/quickstart/box.py 并跑通。
"""

import os
from pathlib import Path
from typing import TypedDict

from dotenv import load_dotenv

load_dotenv()

from langgraph.graph import StateGraph, START, END  # noqa: E402

from assembleagent.core.config import get_runner  # noqa: E402

WORKSPACE = Path(__file__).resolve().parent / "workspace" / "quickstart"

QUICKSTART_TASK = """\
请用 Python + trimesh 库做这件事：
1. 写一个 box.py 到当前工作目录
2. 脚本里 import trimesh 并生成一个 100x60x40（毫米）的长方体
3. export 为 box.stl
4. 用 run_python（或 Bash）执行 box.py，确认产物存在并打印 BOX_OK
5. 简要总结你做了什么
"""


class QState(TypedDict, total=False):
    runner_summary: str
    artifacts: list[str]
    success: bool
    cost_usd: float
    duration_s: float
    model_usage: dict


def call_runner_node(state: QState) -> QState:
    runner = get_runner()
    print(f"[quickstart] 当前 runner = {runner.name}")
    result = runner.run_task(task=QUICKSTART_TASK, workspace=WORKSPACE, timeout=600)
    return {
        "runner_summary": (result.get("summary") or "")[:600],
        "artifacts": result.get("artifacts", []),
        "success": result.get("success", False),
        "cost_usd": result.get("cost_usd", 0),
        "duration_s": result.get("duration_s", 0),
        "model_usage": result.get("model_usage", {}),
    }


def build_graph():
    g = StateGraph(QState)
    g.add_node("call_runner", call_runner_node)
    g.add_edge(START, "call_runner")
    g.add_edge("call_runner", END)
    return g.compile()


def main() -> None:
    required = ("LANGSMITH_API_KEY", "RUNNER_MODE")
    missing = [k for k in required if not os.getenv(k)]
    if missing:
        raise SystemExit(f"缺少环境变量: {missing}，请检查 .env")

    print("=" * 60)
    print(f"RUNNER_MODE = {os.getenv('RUNNER_MODE')}")
    print(f"workspace   = {WORKSPACE}")
    print("=" * 60)

    app = build_graph()
    result = app.invoke({})

    print("=" * 60)
    print("执行结果:")
    print(f"  success:     {result.get('success')}")
    print(f"  duration_s:  {result.get('duration_s'):.2f}")
    print(f"  cost_usd:    {result.get('cost_usd')}")
    print(f"  artifacts:   {result.get('artifacts')}")
    if result.get("model_usage"):
        print(f"  model_usage:")
        for model, usage in result["model_usage"].items():
            inp = usage.get("inputTokens", 0)
            out = usage.get("outputTokens", 0)
            cost = usage.get("costUSD", 0)
            print(f"    - {model}: {inp} in / {out} out  (${cost:.4f})")
    print(f"  summary(尾):\n{result.get('runner_summary')}")
    print("=" * 60)
    print(f"产物目录: {WORKSPACE}")
    print("LangSmith → Assembleagent 项目应当有本次 trace。")


if __name__ == "__main__":
    main()