Agent Zero
192 lines
8.2 KiB
Python
192 lines
8.2 KiB
Python
"""AI Tools for the Tool Registry — task tools for the AI Assistant.
|
|
|
|
Each tool is an async function (arguments: dict, context: dict) -> str
|
|
that delegates to the tasks service layer and returns a JSON-string result.
|
|
All tools require the ``tasks:write`` permission.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import logging
|
|
import uuid
|
|
from typing import Any
|
|
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.plugins.builtins.tasks import services
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
def _get_db_and_tenant(context: dict[str, Any]) -> tuple[AsyncSession, uuid.UUID, uuid.UUID]:
|
|
"""Extract tenant_id, user_id from context and return the DB session."""
|
|
tenant_id = uuid.UUID(context.get("tenant_id", ""))
|
|
user_id = uuid.UUID(context.get("user_id", ""))
|
|
db = context.get("db")
|
|
if db is not None and isinstance(db, AsyncSession):
|
|
return db, tenant_id, user_id
|
|
raise RuntimeError("No db session in context")
|
|
|
|
|
|
async def create_task_handler(arguments: dict[str, Any], context: dict[str, Any]) -> str:
|
|
"""Create a new task (supports polymorphic assignee/entity/creator)."""
|
|
try:
|
|
db, tenant_id, user_id = _get_db_and_tenant(context)
|
|
data = dict(arguments)
|
|
result = await services.create_task(db, tenant_id, user_id, data)
|
|
return json.dumps(result, default=str)
|
|
except Exception as e:
|
|
logger.exception("create_task_handler failed")
|
|
return json.dumps({"error": str(e)})
|
|
|
|
|
|
async def assign_task_handler(arguments: dict[str, Any], context: dict[str, Any]) -> str:
|
|
"""Assign a task to a user, agent, or group (polymorphic)."""
|
|
try:
|
|
db, tenant_id, _ = _get_db_and_tenant(context)
|
|
task_id = uuid.UUID(arguments["task_id"])
|
|
assignee_type = arguments.get("assignee_type", "user")
|
|
assignee_id = uuid.UUID(arguments["assignee_id"]) if arguments.get("assignee_id") else None
|
|
result = await services.assign_task(
|
|
db, tenant_id, task_id, None,
|
|
assignee_type=assignee_type, assignee_id=assignee_id,
|
|
)
|
|
if result is None:
|
|
return json.dumps({"error": "Task not found"})
|
|
return json.dumps(result, default=str)
|
|
except Exception as e:
|
|
logger.exception("assign_task_handler failed")
|
|
return json.dumps({"error": str(e)})
|
|
|
|
|
|
async def update_task_status_handler(arguments: dict[str, Any], context: dict[str, Any]) -> str:
|
|
"""Update a task's lifecycle status."""
|
|
try:
|
|
db, tenant_id, _ = _get_db_and_tenant(context)
|
|
task_id = uuid.UUID(arguments["task_id"])
|
|
status = arguments["status"]
|
|
result = await services.update_task_status(db, tenant_id, task_id, status)
|
|
if result is None:
|
|
return json.dumps({"error": "Task not found"})
|
|
return json.dumps(result, default=str)
|
|
except Exception as e:
|
|
logger.exception("update_task_status_handler failed")
|
|
return json.dumps({"error": str(e)})
|
|
|
|
|
|
async def decompose_goal_handler(arguments: dict[str, Any], context: dict[str, Any]) -> str:
|
|
"""Decompose a goal into subtasks (milestones/todos)."""
|
|
try:
|
|
db, tenant_id, user_id = _get_db_and_tenant(context)
|
|
goal_id = uuid.UUID(arguments["goal_id"])
|
|
subtasks = arguments.get("subtasks") or []
|
|
result = await services.decompose_goal(db, tenant_id, user_id, goal_id, subtasks)
|
|
if result is None:
|
|
return json.dumps({"error": "Goal not found"})
|
|
return json.dumps(result, default=str)
|
|
except Exception as e:
|
|
logger.exception("decompose_goal_handler failed")
|
|
return json.dumps({"error": str(e)})
|
|
|
|
|
|
# ─── Registration ───
|
|
|
|
|
|
def register_task_tools(registry) -> None:
|
|
"""Register task AI tools into the existing Tool Registry."""
|
|
registry.register(
|
|
name="create_task",
|
|
description="Neue Aufgabe erstellen (Unified Task System). Unterstützt polymorphe Zuweisung (user/agent/group), Entity-Verknüpfung (contact/company/...), Subtasks, Dependencies, Goals/Milestones und Agent-Subtasks.",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"title": {"type": "string", "description": "Aufgabentitel"},
|
|
"description": {"type": "string", "description": "Beschreibung"},
|
|
"status": {"type": "string", "enum": ["open", "in_progress", "review", "blocked", "done", "cancelled"], "default": "open"},
|
|
"priority": {"type": "string", "enum": ["low", "medium", "high", "urgent"], "default": "medium"},
|
|
"due_date": {"type": "string", "description": "ISO Datum"},
|
|
"assignee_type": {"type": "string", "enum": ["user", "agent", "group"], "default": "user"},
|
|
"assignee_id": {"type": "string", "description": "UUID des Assignees"},
|
|
"entity_type": {"type": "string", "description": "z.B. contact, company"},
|
|
"entity_id": {"type": "string", "description": "UUID der Entity"},
|
|
"parent_task_id": {"type": "string", "description": "UUID des Parent-Tasks (Subtask)"},
|
|
"depends_on": {"type": "array", "items": {"type": "string"}, "description": "Task-IDs von denen diese Aufgabe abhängt"},
|
|
"task_type": {"type": "string", "enum": ["todo", "approval", "follow_up", "review", "goal", "milestone", "agent_subtask"], "default": "todo"},
|
|
"success_criteria": {"type": "object", "description": "Strukturierte Erfolgskriterien für Goals"},
|
|
"target_date": {"type": "string", "description": "Deadline für Goal/Milestone"},
|
|
},
|
|
"required": ["title"],
|
|
},
|
|
handler=create_task_handler,
|
|
plugin_name="tasks",
|
|
required_permission="tasks:write",
|
|
category="tasks",
|
|
)
|
|
|
|
registry.register(
|
|
name="assign_task",
|
|
description="Aufgabe einem User, Agent oder einer Gruppe zuweisen (polymorph).",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"task_id": {"type": "string", "description": "UUID der Aufgabe"},
|
|
"assignee_type": {"type": "string", "enum": ["user", "agent", "group"], "default": "user"},
|
|
"assignee_id": {"type": "string", "description": "UUID des Assignees"},
|
|
},
|
|
"required": ["task_id", "assignee_id"],
|
|
},
|
|
handler=assign_task_handler,
|
|
plugin_name="tasks",
|
|
required_permission="tasks:write",
|
|
category="tasks",
|
|
)
|
|
|
|
registry.register(
|
|
name="update_task_status",
|
|
description="Status einer Aufgabe aktualisieren (open/in_progress/review/blocked/done/cancelled).",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"task_id": {"type": "string", "description": "UUID der Aufgabe"},
|
|
"status": {"type": "string", "enum": ["open", "in_progress", "review", "blocked", "done", "cancelled"]},
|
|
},
|
|
"required": ["task_id", "status"],
|
|
},
|
|
handler=update_task_status_handler,
|
|
plugin_name="tasks",
|
|
required_permission="tasks:write",
|
|
category="tasks",
|
|
)
|
|
|
|
registry.register(
|
|
name="decompose_goal",
|
|
description="Ein Goal in Subtasks (Milestones/Todos) zerlegen.",
|
|
parameters={
|
|
"type": "object",
|
|
"properties": {
|
|
"goal_id": {"type": "string", "description": "UUID des Goals"},
|
|
"subtasks": {
|
|
"type": "array",
|
|
"items": {
|
|
"type": "object",
|
|
"properties": {
|
|
"title": {"type": "string"},
|
|
"description": {"type": "string"},
|
|
"milestone": {"type": "boolean", "default": False},
|
|
"assignee_type": {"type": "string", "enum": ["user", "agent", "group"]},
|
|
"assignee_id": {"type": "string"},
|
|
},
|
|
"required": ["title"],
|
|
},
|
|
},
|
|
},
|
|
"required": ["goal_id", "subtasks"],
|
|
},
|
|
handler=decompose_goal_handler,
|
|
plugin_name="tasks",
|
|
required_permission="tasks:write",
|
|
category="tasks",
|
|
)
|
|
logger.info("Task AI tools registered")
|