wuyang/zk-data-agent
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: rebuild as multi-user web agent

Browse Source
This commit is contained in:
wuyang
2026-07-26 06:52:12 +08:00
parent 45792c8fd5
commit 37f83eaac7
634 changed files with 5060 additions and 139619 deletions
Download Patch File Download Diff File Expand all files Collapse all files
agent_platform/runtime/tools.py
+341
View File
@@ -0,0 +1,341 @@
from __future__ import annotations
import json
from dataclasses import dataclass
from typing import Any
import httpx
from agent_platform.auth import UserIdentity
from agent_platform.config import Settings
from agent_platform.runtime.schemas import PlanItem
from agent_platform.store import RuntimeStore
@dataclass(frozen=True, slots=True)
class ToolMetadata:
name: str
description: str
schema: dict[str, Any]
parallel_safe: bool
read_only: bool
def openai_spec(self) -> dict[str, Any]:
return {
"type": "function",
"function": {
"name": self.name,
"description": self.description,
"parameters": self.schema,
},
}
def object_schema(properties: dict[str, Any], required: list[str] | None = None) -> dict[str, Any]:
return {
"type": "object",
"properties": properties,
"required": required or [],
"additionalProperties": False,
}
TOOL_METADATA: dict[str, ToolMetadata] = {
"workspace_status": ToolMetadata(
"workspace_status",
"Return the current user's isolated workspace status.",
object_schema({}),
True,
True,
),
"list_files": ToolMetadata(
"list_files",
"List files and directories in the isolated workspace.",
object_schema(
{
"path": {"type": "string", "default": "."},
"max_depth": {"type": "integer", "minimum": 1, "maximum": 12, "default": 4},
"limit": {"type": "integer", "minimum": 1, "maximum": 5000, "default": 500},
}
),
True,
True,
),
"read_file": ToolMetadata(
"read_file",
"Read a UTF-8 text file with line numbers.",
object_schema(
{
"path": {"type": "string"},
"start_line": {"type": "integer", "minimum": 1, "default": 1},
"max_lines": {"type": "integer", "minimum": 1, "maximum": 5000, "default": 1000},
},
["path"],
),
True,
True,
),
"search_files": ToolMetadata(
"search_files",
"Search workspace text using ripgrep.",
object_schema(
{
"query": {"type": "string"},
"path": {"type": "string", "default": "."},
"glob": {"type": ["string", "null"], "default": None},
"limit": {"type": "integer", "minimum": 1, "maximum": 2000, "default": 200},
},
["query"],
),
True,
True,
),
"write_file": ToolMetadata(
"write_file",
"Write the complete UTF-8 contents of a workspace file.",
object_schema(
{
"path": {"type": "string", "maxLength": 4096},
"content": {"type": "string", "maxLength": 2_000_000},
},
["path", "content"],
),
False,
False,
),
"apply_patch": ToolMetadata(
"apply_patch",
"Apply a unified diff in a workspace repository.",
object_schema(
{"patch": {"type": "string"}, "cwd": {"type": "string", "default": "."}},
["patch"],
),
False,
False,
),
"exec": ToolMetadata(
"exec",
"Run a shell command in the isolated workspace.",
object_schema(
{
"command": {"type": "string"},
"cwd": {"type": "string", "default": "."},
"timeout_seconds": {"type": "integer", "minimum": 1, "maximum": 1800, "default": 120},
},
["command"],
),
False,
False,
),
"git_status": ToolMetadata(
"git_status",
"Show concise Git status for a workspace repository.",
object_schema({"cwd": {"type": "string", "default": "."}}),
True,
True,
),
"git_diff": ToolMetadata(
"git_diff",
"Show the unstaged or staged Git diff.",
object_schema(
{
"cwd": {"type": "string", "default": "."},
"staged": {"type": "boolean", "default": False},
}
),
True,
True,
),
"start_process": ToolMetadata(
"start_process",
"Start a long-running background process and return a process id.",
object_schema(
{"command": {"type": "string"}, "cwd": {"type": "string", "default": "."}},
["command"],
),
False,
False,
),
"poll_process": ToolMetadata(
"poll_process",
"Poll a background process and return recent output.",
object_schema({"process_id": {"type": "string"}}, ["process_id"]),
True,
True,
),
"cancel_process": ToolMetadata(
"cancel_process",
"Stop a background process.",
object_schema({"process_id": {"type": "string"}}, ["process_id"]),
False,
False,
),
"update_plan": ToolMetadata(
"update_plan",
"Publish a concise execution plan. At most one step may be in progress.",
object_schema(
{
"explanation": {"type": ["string", "null"]},
"items": {
"type": "array",
"minItems": 1,
"maxItems": 50,
"items": {
"type": "object",
"properties": {
"step": {"type": "string"},
"status": {"type": "string", "enum": ["pending", "in_progress", "completed"]},
},
"required": ["step", "status"],
"additionalProperties": False,
},
},
},
["items"],
),
False,
False,
),
"remember": ToolMetadata(
"remember",
"Store a durable user preference or fact that will help future Work tasks.",
object_schema({"content": {"type": "string"}}, ["content"]),
False,
False,
),
"recall_memory": ToolMetadata(
"recall_memory",
"Search durable Work memory for this user.",
object_schema(
{
"query": {"type": "string", "default": ""},
"limit": {"type": "integer", "minimum": 1, "maximum": 50, "default": 8},
}
),
True,
True,
),
"forget_memory": ToolMetadata(
"forget_memory",
"Delete one durable Work memory by id.",
object_schema({"memory_id": {"type": "string"}}, ["memory_id"]),
False,
False,
),
"delegate_task": ToolMetadata(
"delegate_task",
"Delegate a bounded subtask to a child Agent. Read-only delegates may run in parallel.",
object_schema(
{
"task": {"type": "string"},
"role": {"type": "string", "default": "researcher"},
"allow_writes": {"type": "boolean", "default": False},
},
["task"],
),
True,
True,
),
}
GATEWAY_ENDPOINTS = {
name: f"/v1/tools/{name}"
for name in (
"workspace_status",
"list_files",
"read_file",
"search_files",
"write_file",
"apply_patch",
"exec",
"git_status",
"git_diff",
"start_process",
"poll_process",
"cancel_process",
)
}
READ_ONLY_TOOL_NAMES = {
name for name, metadata in TOOL_METADATA.items() if metadata.read_only and name != "delegate_task"
}
@dataclass(frozen=True, slots=True)
class ToolContext:
identity: UserIdentity
raw_user_jwt: str
chat_id: str
class ToolRegistry:
def __init__(
self,
settings: Settings,
store: RuntimeStore,
client: httpx.AsyncClient | None = None,
) -> None:
self.settings = settings
self.store = store
self.client = client or httpx.AsyncClient(timeout=httpx.Timeout(settings.tool_timeout_seconds))
self._owns_client = client is None
async def close(self) -> None:
if self._owns_client:
await self.client.aclose()
def specs(self, *, read_only: bool = False, allow_delegate: bool = True) -> list[dict[str, Any]]:
names = READ_ONLY_TOOL_NAMES if read_only else set(TOOL_METADATA)
if not allow_delegate:
names = names - {"delegate_task"}
return [TOOL_METADATA[name].openai_spec() for name in TOOL_METADATA if name in names]
async def execute(self, name: str, arguments: dict[str, Any], context: ToolContext) -> dict[str, Any]:
if name in GATEWAY_ENDPOINTS:
response = await self.client.post(
f"{self.settings.gateway_url}{GATEWAY_ENDPOINTS[name]}",
headers={
"Authorization": f"Bearer {self.settings.internal_gateway_key}",
"X-OpenWebUI-User-Jwt": context.raw_user_jwt,
},
json=arguments,
)
if response.status_code >= 400:
return {
"ok": False,
"status_code": response.status_code,
"error": response.text[:4000],
}
return response.json()
if name == "update_plan":
items = [PlanItem.model_validate(item).model_dump() for item in arguments.get("items", [])]
if sum(item["status"] == "in_progress" for item in items) > 1:
raise ValueError("At most one plan item may be in progress")
await self.store.update_plan(context.identity.user_id, context.chat_id, items)
return {"ok": True, "explanation": arguments.get("explanation"), "items": items}
if name == "remember":
content = str(arguments.get("content", "")).strip()
if not content:
raise ValueError("Memory content is required")
memory_id = await self.store.remember(context.identity.user_id, content[:20_000])
return {"ok": True, "memory_id": memory_id}
if name == "recall_memory":
return {
"ok": True,
"items": await self.store.recall(
context.identity.user_id,
str(arguments.get("query", "")),
int(arguments.get("limit", 8)),
),
}
if name == "forget_memory":
deleted = await self.store.forget(context.identity.user_id, str(arguments.get("memory_id", "")))
return {"ok": deleted}
raise ValueError(f"Unknown tool: {name}")
def tool_result_text(result: Any, limit: int) -> str:
text = json.dumps(result, ensure_ascii=False, separators=(",", ":"), default=str)
if len(text) <= limit:
return text
return text[:limit] + "…"