wuyang
346 lines
11 KiB
Python
346 lines
11 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
import httpx
|
|
|
|
from agent_platform.auth import UserIdentity, issue_internal_identity
|
|
from agent_platform.config import Settings
|
|
from agent_platform.runtime.schemas import PlanItem
|
|
from agent_platform.store import RuntimeStore
|
|
from agent_platform.text import truncate_middle
|
|
|
|
|
|
@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 complete UTF-8 file contents. Use actual line breaks between source lines, not literal \\\\n text.",
|
|
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 complete non-interactive shell command. Name the script, pass -c code, or invoke a test; "
|
|
"bare python, node, or shells are rejected. For Python dependencies, create or reuse "
|
|
"/workspace/.venv and install with .venv/bin/python -m pip; /tmp is not executable and the "
|
|
"read-only system Python must not be modified.",
|
|
object_schema(
|
|
{
|
|
"command": {"type": "string"},
|
|
"cwd": {"type": "string", "default": "."},
|
|
"timeout_seconds": {"type": "integer", "minimum": 1, "maximum": 1800, "default": 900},
|
|
},
|
|
["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 Agent tasks.",
|
|
object_schema({"content": {"type": "string"}}, ["content"]),
|
|
False,
|
|
False,
|
|
),
|
|
"recall_memory": ToolMetadata(
|
|
"recall_memory",
|
|
"Search durable Agent 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 Agent 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
|
|
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": issue_internal_identity(
|
|
context.identity,
|
|
self.settings.openwebui_forward_jwt_secret,
|
|
),
|
|
},
|
|
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)
|
|
return truncate_middle(text, limit)
|