ducklm/duck_core/runtime_loop.py

import json
from dataclasses import dataclass
from typing import Any

from duck_core.approvals.service import ApprovalService
from duck_core.context_builder import ContextBuilder
from duck_core.events.store import EventStore
from duck_core.model_client import ModelClient
from duck_core.tasks.store import TaskStore
from duck_core.tools.gateway import ToolGateway


@dataclass
class ChatResult:
    task_id: str
    status: str
    final_response: str
    reasoning_content: str | None = None


class RuntimeLoop:
    def __init__(
        self,
        task_store: TaskStore,
        event_store: EventStore,
        model_client: ModelClient | None = None,
        context_builder: ContextBuilder | None = None,
        approval_service: ApprovalService | None = None,
    ):
        self.task_store = task_store
        self.event_store = event_store
        self.model_client = model_client or ModelClient()
        self.context_builder = context_builder or ContextBuilder()
        self.approval_service = approval_service

    async def run_chat(
        self, message: str, workspace: str | None = None, debug: bool = False
    ) -> ChatResult:
        task = await self.task_store.create_task(message, workspace, debug)
        await self.event_store.append(
            task.task_id,
            "task_created",
            {"message": message, "workspace": workspace, "debug": debug},
        )
        try:
            messages = self.context_builder.build_basic_messages(task)
            tool_observations = await self._run_action_tools(task.task_id, messages, workspace)
            if any(observation.get("requires_approval") for observation in tool_observations):
                await self.task_store.waiting_for_approval(task.task_id)
                await self.event_store.append(
                    task.task_id,
                    "task_waiting_for_approval",
                    {"observations": tool_observations},
                )
                return ChatResult(
                    task_id=task.task_id,
                    status="waiting_for_approval",
                    final_response="Waiting for approval.",
                    reasoning_content=None,
                )
            if tool_observations:
                messages = [
                    *messages,
                    {
                        "role": "user",
                        "content": "tool_observations:\n"
                        + json.dumps(tool_observations, ensure_ascii=False, indent=2),
                    },
                ]
            await self.event_store.append(
                task.task_id, "model_call_started", {"role": "thinker"}
            )
            response = await self.model_client.chat("thinker", messages)
            await self.event_store.append(
                task.task_id,
                "cognition_response",
                {
                    "role": response.role,
                    "content": response.content,
                    "reasoning_content": response.reasoning_content,
                },
            )
            await self.event_store.append(
                task.task_id,
                "model_call_finished",
                {
                    "role": response.role,
                    "model": response.model,
                    "latency_ms": response.latency_ms,
                    "prompt_tokens": response.prompt_tokens,
                    "completion_tokens": response.completion_tokens,
                    "total_tokens": response.total_tokens,
                },
            )
            await self.task_store.complete_task(task.task_id, response.content)
            await self.event_store.append(
                task.task_id,
                "task_completed",
                {
                    "final_response": response.content,
                    "reasoning_content": response.reasoning_content,
                },
            )
            return ChatResult(
                task_id=task.task_id,
                status="completed",
                final_response=response.content,
                reasoning_content=response.reasoning_content,
            )
        except Exception as exc:
            await self.task_store.fail_task(task.task_id, str(exc))
            await self.event_store.append(
                task.task_id, "task_failed", {"error": str(exc)}
            )
            return ChatResult(
                task_id=task.task_id,
                status="failed",
                final_response=str(exc),
                reasoning_content=None,
            )

    async def _run_action_tools(
        self, task_id: str, messages: list[dict[str, str]], workspace: str | None
    ) -> list[dict[str, Any]]:
        try:
            await self.event_store.append(task_id, "model_call_started", {"role": "action"})
            response = await self.model_client.chat("action", messages)
            directive = json.loads(response.content)
        except Exception as exc:
            await self.event_store.append(
                task_id,
                "action_directive_failed",
                {"error": str(exc)},
            )
            return []

        await self.event_store.append(task_id, "action_directive", directive)
        actions = directive.get("actions") or []
        if not isinstance(actions, list) or not actions:
            return []

        gateway = ToolGateway.default(workspace or ".")
        observations: list[dict[str, Any]] = []
        for index, action in enumerate(actions, start=1):
            if not isinstance(action, dict):
                observations.append(
                    {"index": index, "ok": False, "error": "Action must be an object"}
                )
                continue
            tool_name = str(action.get("tool", ""))
            await self.event_store.append(
                task_id,
                "tool_call_started",
                {"index": index, "tool": tool_name, "args": action.get("args") or {}},
            )
            result = await gateway.run_action(action)
            result_payload = result.model_dump()
            if result.metadata.get("requires_approval"):
                approval = None
                if self.approval_service is not None:
                    approval = await self.approval_service.create_pending(task_id, action)
                await self.event_store.append(
                    task_id,
                    "tool_approval_requested",
                    {
                        "index": index,
                        "tool": tool_name,
                        "action": action,
                        "approval_id": approval.approval_id if approval else None,
                        "reason": result.error,
                    },
                )
                observations.append(
                    {
                        "index": index,
                        "tool": tool_name,
                        "reason": action.get("reason"),
                        "requires_approval": True,
                        "approval_id": approval.approval_id if approval else None,
                        "result": result_payload,
                    }
                )
                break
            await self.event_store.append(
                task_id,
                "tool_call_finished",
                {"index": index, "tool": tool_name, "result": result_payload},
            )
            observations.append(
                {
                    "index": index,
                    "tool": tool_name,
                    "reason": action.get("reason"),
                    "result": result_payload,
                }
            )
        return observations