Shane
d372adca7d
The 4B model was responding conversationally instead of calling memory_save/ memory_recall. Added imperative language (MUST call), concrete examples of trigger phrases, and explicit instructions to never skip the tool call. Verified: model now reliably generates tool_calls for save/recall/list.
127 lines
4.0 KiB
Python
127 lines
4.0 KiB
Python
"""In-memory task registry for background dispatch tasks.
|
|
|
|
One global instance per agent process. Workers add steps as they progress;
|
|
the agent subscribes to push updates to the room (UI) and to speak results.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import time
|
|
import uuid
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum
|
|
from typing import Any, Callable
|
|
|
|
|
|
class TaskStatus(str, Enum):
|
|
PENDING = "pending"
|
|
RUNNING = "running"
|
|
COMPLETED = "completed"
|
|
FAILED = "failed"
|
|
|
|
|
|
@dataclass
|
|
class TaskStep:
|
|
role: str # "thinking", "tool_call", "tool_result", "text"
|
|
content: str
|
|
timestamp: float = field(default_factory=time.time)
|
|
|
|
|
|
@dataclass
|
|
class Task:
|
|
id: str
|
|
description: str
|
|
status: TaskStatus = TaskStatus.PENDING
|
|
steps: list[TaskStep] = field(default_factory=list)
|
|
result: str | None = None
|
|
error: str | None = None
|
|
created_at: float = field(default_factory=time.time)
|
|
completed_at: float | None = None
|
|
|
|
|
|
class TaskRegistry:
|
|
"""Thread-safe registry of background tasks with pub/sub notifications."""
|
|
|
|
def __init__(self):
|
|
self._tasks: dict[str, Task] = {}
|
|
self._listeners: list[Callable[[str, dict], Any]] = []
|
|
self._lock = asyncio.Lock()
|
|
|
|
async def create(self, description: str) -> Task:
|
|
task = Task(id=uuid.uuid4().hex[:8], description=description)
|
|
async with self._lock:
|
|
self._tasks[task.id] = task
|
|
await self._notify(task.id, {"event": "created", "task": self._task_dict(task)})
|
|
return task
|
|
|
|
async def add_step(self, task_id: str, role: str, content: str):
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
task.steps.append(TaskStep(role=role, content=content))
|
|
if task.status == TaskStatus.PENDING:
|
|
task.status = TaskStatus.RUNNING
|
|
await self._notify(task_id, {"event": "step", "role": role, "content": content})
|
|
|
|
async def complete(self, task_id: str, result: str):
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
task.status = TaskStatus.COMPLETED
|
|
task.result = result
|
|
task.completed_at = time.time()
|
|
await self._notify(task_id, {"event": "completed", "result": result})
|
|
|
|
async def fail(self, task_id: str, error: str):
|
|
async with self._lock:
|
|
task = self._tasks.get(task_id)
|
|
if not task:
|
|
return
|
|
task.status = TaskStatus.FAILED
|
|
task.error = error
|
|
task.completed_at = time.time()
|
|
await self._notify(task_id, {"event": "failed", "error": error})
|
|
|
|
async def get(self, task_id: str) -> dict | None:
|
|
async with self._lock:
|
|
t = self._tasks.get(task_id)
|
|
return self._task_dict(t) if t else None
|
|
|
|
async def list_all(self) -> list[dict]:
|
|
async with self._lock:
|
|
return [self._task_dict(t) for t in self._tasks.values()]
|
|
|
|
def subscribe(self, callback: Callable[[str, dict], Any]):
|
|
self._listeners.append(callback)
|
|
|
|
async def _notify(self, task_id: str, event: dict):
|
|
for cb in self._listeners:
|
|
try:
|
|
result = cb(task_id, event)
|
|
if asyncio.iscoroutine(result):
|
|
await result
|
|
except Exception:
|
|
pass
|
|
|
|
@staticmethod
|
|
def _task_dict(task: Task) -> dict:
|
|
return {
|
|
"id": task.id,
|
|
"description": task.description,
|
|
"status": task.status.value,
|
|
"steps": [
|
|
{"role": s.role, "content": s.content, "ts": s.timestamp}
|
|
for s in task.steps
|
|
],
|
|
"result": task.result,
|
|
"error": task.error,
|
|
"created_at": task.created_at,
|
|
"completed_at": task.completed_at,
|
|
}
|
|
|
|
|
|
# Global singleton — one per agent process
|
|
registry = TaskRegistry()
|