wuyang
258 lines
9.2 KiB
Python
Executable File
258 lines
9.2 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import json
|
|
import os
|
|
import tempfile
|
|
import time
|
|
import uuid
|
|
from dataclasses import replace
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from docker.errors import NotFound
|
|
|
|
from agent_platform.auth import UserIdentity
|
|
from agent_platform.config import Settings
|
|
from agent_platform.gateway.provider import DockerExecutionProvider, workspace_ref
|
|
from agent_platform.models import get_model_spec
|
|
from agent_platform.runtime.loop import AgentLoop
|
|
from agent_platform.runtime.provider import ModelProvider
|
|
from agent_platform.runtime.schemas import PlanItem
|
|
from agent_platform.runtime.tools import TOOL_METADATA, ToolContext
|
|
from agent_platform.store import RuntimeStore
|
|
|
|
EVAL_TOOL_NAMES = (
|
|
"workspace_status",
|
|
"list_files",
|
|
"read_file",
|
|
"search_files",
|
|
"write_file",
|
|
"apply_patch",
|
|
"exec",
|
|
"update_plan",
|
|
)
|
|
|
|
|
|
class DirectDockerRegistry:
|
|
def __init__(
|
|
self,
|
|
execution: DockerExecutionProvider,
|
|
store: RuntimeStore,
|
|
tool_timeout_seconds: int,
|
|
) -> None:
|
|
self.execution = execution
|
|
self.store = store
|
|
self.tool_timeout_seconds = tool_timeout_seconds
|
|
|
|
def specs(self, *, read_only: bool = False, allow_delegate: bool = True) -> list[dict[str, Any]]:
|
|
names = [name for name in EVAL_TOOL_NAMES if not read_only or TOOL_METADATA[name].read_only]
|
|
return [TOOL_METADATA[name].openai_spec() for name in names]
|
|
|
|
async def execute(
|
|
self,
|
|
name: str,
|
|
arguments: dict[str, Any],
|
|
context: ToolContext,
|
|
) -> dict[str, Any]:
|
|
user_id = context.identity.user_id
|
|
if name == "workspace_status":
|
|
result = await self.execution.status(user_id)
|
|
elif name == "list_files":
|
|
result = await self.execution.list_files(
|
|
user_id,
|
|
str(arguments.get("path", ".")),
|
|
int(arguments.get("max_depth", 4)),
|
|
int(arguments.get("limit", 500)),
|
|
)
|
|
elif name == "read_file":
|
|
result = await self.execution.read_file(
|
|
user_id,
|
|
str(arguments["path"]),
|
|
int(arguments.get("start_line", 1)),
|
|
int(arguments.get("max_lines", 1000)),
|
|
)
|
|
elif name == "search_files":
|
|
result = await self.execution.search_files(
|
|
user_id,
|
|
str(arguments["query"]),
|
|
str(arguments.get("path", ".")),
|
|
arguments.get("glob"),
|
|
int(arguments.get("limit", 200)),
|
|
)
|
|
elif name == "write_file":
|
|
result = await self.execution.write_file(
|
|
user_id,
|
|
str(arguments["path"]),
|
|
str(arguments["content"]),
|
|
)
|
|
elif name == "apply_patch":
|
|
result = await self.execution.apply_patch(
|
|
user_id,
|
|
str(arguments["patch"]),
|
|
str(arguments.get("cwd", ".")),
|
|
)
|
|
elif name == "exec":
|
|
result = await self.execution.exec(
|
|
user_id,
|
|
str(arguments["command"]),
|
|
str(arguments.get("cwd", ".")),
|
|
min(
|
|
int(arguments.get("timeout_seconds", 120)),
|
|
self.tool_timeout_seconds,
|
|
),
|
|
)
|
|
elif name == "update_plan":
|
|
items = [PlanItem.model_validate(item).model_dump() for item in arguments.get("items", [])]
|
|
await self.store.update_plan(user_id, context.chat_id, items)
|
|
return {
|
|
"ok": True,
|
|
"explanation": arguments.get("explanation"),
|
|
"items": items,
|
|
}
|
|
else:
|
|
raise ValueError(f"Unsupported evaluation tool: {name}")
|
|
return result.model_dump()
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Run the real Work loop against a live provider in a disposable Docker workspace."
|
|
)
|
|
parser.add_argument("--model", default="work-extreme")
|
|
parser.add_argument(
|
|
"--prompt",
|
|
default="写脚本对比一下几个常见排序算法，给一个报告",
|
|
)
|
|
parser.add_argument(
|
|
"--workspace-image",
|
|
default=os.getenv("WORKSPACE_IMAGE", "k1412-agent-workspace:test"),
|
|
)
|
|
parser.add_argument("--keep-workspace", action="store_true")
|
|
return parser.parse_args()
|
|
|
|
|
|
async def cleanup_workspace(execution: DockerExecutionProvider, user_id: str) -> None:
|
|
ref = workspace_ref(user_id)
|
|
try:
|
|
execution.client.containers.get(ref.container_name).remove(force=True)
|
|
except NotFound:
|
|
pass
|
|
try:
|
|
execution.client.volumes.get(ref.volume_name).remove(force=True)
|
|
except NotFound:
|
|
pass
|
|
try:
|
|
execution.client.networks.get(ref.network_name).remove()
|
|
except NotFound:
|
|
pass
|
|
|
|
|
|
async def main() -> None:
|
|
args = parse_args()
|
|
spec = get_model_spec(args.model)
|
|
if spec.mode != "work":
|
|
raise SystemExit("--model must select a Work tier")
|
|
if spec.provider == "deepseek" and not os.getenv("DEEPSEEK_API_KEY", "").strip():
|
|
raise SystemExit("DEEPSEEK_API_KEY is required")
|
|
|
|
user_id = f"live-eval-{uuid.uuid4().hex}"
|
|
chat_id = f"live-eval-{uuid.uuid4().hex}"
|
|
started = time.monotonic()
|
|
event_counts: dict[str, int] = {}
|
|
usage_totals = {
|
|
"prompt_tokens": 0,
|
|
"completion_tokens": 0,
|
|
"reasoning_tokens": 0,
|
|
"cached_tokens": 0,
|
|
"total_tokens": 0,
|
|
}
|
|
|
|
with tempfile.TemporaryDirectory(prefix="k1412-live-eval-") as temp_dir:
|
|
settings = replace(
|
|
Settings.from_env(),
|
|
database_url=f"sqlite+aiosqlite:///{Path(temp_dir) / 'runtime.sqlite3'}",
|
|
workspace_image=args.workspace_image,
|
|
workspace_network_enabled=False,
|
|
workspace_memory_limit="1g",
|
|
workspace_cpu_limit=1.0,
|
|
workspace_pids_limit=128,
|
|
tool_timeout_seconds=180,
|
|
)
|
|
store = RuntimeStore(settings.database_url)
|
|
await store.initialize()
|
|
provider = ModelProvider(settings)
|
|
execution = DockerExecutionProvider(settings)
|
|
registry = DirectDockerRegistry(execution, store, settings.tool_timeout_seconds)
|
|
|
|
async def record(event_type: str, payload: dict[str, Any]) -> None:
|
|
event_counts[event_type] = event_counts.get(event_type, 0) + 1
|
|
if event_type == "model.responded":
|
|
usage = payload.get("usage") or {}
|
|
prompt_details = usage.get("prompt_tokens_details") or {}
|
|
completion_details = usage.get("completion_tokens_details") or {}
|
|
usage_totals["prompt_tokens"] += int(usage.get("prompt_tokens", 0))
|
|
usage_totals["completion_tokens"] += int(usage.get("completion_tokens", 0))
|
|
usage_totals["reasoning_tokens"] += int(completion_details.get("reasoning_tokens", 0))
|
|
usage_totals["cached_tokens"] += int(
|
|
prompt_details.get("cached_tokens", 0) or usage.get("prompt_cache_hit_tokens", 0)
|
|
)
|
|
usage_totals["total_tokens"] += int(usage.get("total_tokens", 0))
|
|
if event_type == "tool.completed":
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"event": event_type,
|
|
"name": payload.get("name"),
|
|
"ok": payload.get("ok"),
|
|
"summary": str(payload.get("summary", ""))[:500],
|
|
},
|
|
ensure_ascii=False,
|
|
),
|
|
flush=True,
|
|
)
|
|
|
|
try:
|
|
answer = await AgentLoop(
|
|
provider,
|
|
registry,
|
|
store,
|
|
max_tool_output_chars=settings.max_tool_output_chars,
|
|
).run(
|
|
spec=spec,
|
|
messages=[{"role": "user", "content": args.prompt}],
|
|
identity=UserIdentity(user_id, "", "Live Eval", "user"),
|
|
raw_user_jwt="local-evaluation",
|
|
chat_id=chat_id,
|
|
callback=record,
|
|
)
|
|
listing = await execution.list_files(user_id, ".", 4, 500)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"model": spec.public_id,
|
|
"provider_model": spec.provider_model,
|
|
"elapsed_seconds": round(time.monotonic() - started, 3),
|
|
"events": event_counts,
|
|
"usage": usage_totals,
|
|
"answer": answer,
|
|
"workspace": listing.model_dump(),
|
|
"workspace_id": workspace_ref(user_id).workspace_id if args.keep_workspace else None,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
finally:
|
|
await provider.close()
|
|
await store.close()
|
|
if not args.keep_workspace:
|
|
await cleanup_workspace(execution, user_id)
|
|
execution.client.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
asyncio.run(main())
|