beaver_project/app-instance/backend/beaver/tasks/router.py

"""LLM-based routing between simple chat and internal Task mode."""

from __future__ import annotations

import asyncio
import json
from typing import Any

from .models import MainAgentDecision, TaskRecord


class MainAgentRouter:
    """Semantic router for deciding whether a message belongs to a Task."""

    async def classify(
        self,
        message: str,
        *,
        active_task: TaskRecord | None = None,
        provider: Any | None = None,
        model: str | None = None,
        recent_messages: list[dict[str, Any]] | None = None,
        thinking_enabled: bool | None = None,
        timeout_seconds: float = 8.0,
    ) -> MainAgentDecision:
        if provider is None:
            return self._fallback(active_task=active_task, reason="router_provider_unavailable")
        try:
            chat_kwargs: dict[str, Any] = {
                "messages": [
                    {
                        "role": "system",
                        "content": (
                            "You route user messages for Beaver's internal Task mode. "
                            "Return only compact JSON. Do not explain."
                        ),
                    },
                    {
                        "role": "user",
                        "content": self._prompt(
                            message=message,
                            active_task=active_task,
                            recent_messages=recent_messages or [],
                        ),
                    },
                ],
                "tools": None,
                "model": model,
                "max_tokens": 256,
                "temperature": 0.0,
            }
            if thinking_enabled is not None:
                chat_kwargs["thinking_enabled"] = thinking_enabled
            response = await asyncio.wait_for(provider.chat(**chat_kwargs), timeout=timeout_seconds)
            return self.from_json(response.content or "", active_task=active_task)
        except Exception as exc:
            return self._fallback(active_task=active_task, reason=f"router_failed: {exc}")

    def from_json(self, text: str, *, active_task: TaskRecord | None = None) -> MainAgentDecision:
        payload = self._parse_json_object(text)
        raw_action = str(payload.get("action") or payload.get("mode") or "").strip().lower()
        reason = str(payload.get("reason") or raw_action or "llm_router")
        short_title = _clean_short_title(payload.get("short_title") or payload.get("title"))

        if raw_action in {"continue_task", "continue", "task"}:
            return MainAgentDecision(mode="task", reason=reason, short_title=short_title)
        if raw_action in {"new_task", "new"}:
            return MainAgentDecision(mode="task", reason=reason, starts_new_task=True, short_title=short_title)
        if raw_action in {"close_task", "close", "done", "finish"}:
            return MainAgentDecision(mode="simple", reason=reason, closes_task=active_task is not None, short_title=short_title)
        if raw_action in {"abandon_task", "abandon", "cancel_task"}:
            return MainAgentDecision(mode="simple", reason=reason, abandons_task=active_task is not None, short_title=short_title)
        return MainAgentDecision(mode="simple", reason=reason or "simple_chat", short_title=short_title)

    def _fallback(self, *, active_task: TaskRecord | None, reason: str) -> MainAgentDecision:
        if active_task is not None:
            return MainAgentDecision(mode="task", reason=reason)
        return MainAgentDecision(mode="simple", reason=reason)

    @staticmethod
    def _prompt(
        *,
        message: str,
        active_task: TaskRecord | None,
        recent_messages: list[dict[str, Any]],
    ) -> str:
        active_task_payload = None
        if active_task is not None:
            active_task_payload = {
                "task_id": active_task.task_id,
                "description": active_task.description,
                "goal": active_task.goal,
                "status": active_task.status,
                "short_title": active_task.metadata.get("short_title"),
            }
        recent = [
            {"role": item.get("role"), "content": str(item.get("content") or "")[:500]}
            for item in recent_messages[-8:]
            if item.get("role") in {"user", "assistant"}
        ]
        return (
            "Decide how to route the current user message.\n\n"
            "Actions:\n"
            "- simple_chat: no Task should be created or continued.\n"
            "- continue_task: keep the user in the active Task.\n"
            "- new_task: start a separate new Task.\n"
            "- close_task: user explicitly says the active Task is done/satisfactory/finished.\n"
            "- abandon_task: user explicitly says to stop, cancel, abandon, or no longer do the active Task.\n\n"
            "Critical policy:\n"
            "- If there is an active Task, choose continue_task unless the user's topic is completely unrelated "
            "to that Task or the user explicitly closes/abandons it.\n"
            "- Follow-up questions, corrections, partial changes, extra constraints, and result discussion stay in continue_task.\n"
            "- Use new_task only when the user clearly asks to start a different task.\n"
            "- If there is no active Task, choose new_task only for work that requires execution, iteration, tools, files, "
            "implementation, validation, or multi-step completion. Otherwise choose simple_chat.\n"
            "- short_title must be 5-15 Chinese characters or a similarly short English phrase when a Task is involved.\n\n"
            "Return JSON only with keys: action, reason, short_title.\n\n"
            f"Active task:\n{json.dumps(active_task_payload, ensure_ascii=False)}\n\n"
            f"Recent conversation:\n{json.dumps(recent, ensure_ascii=False)}\n\n"
            f"Current user message:\n{message}"
        )

    @staticmethod
    def _parse_json_object(text: str) -> dict[str, Any]:
        cleaned = text.strip()
        if cleaned.startswith("```"):
            cleaned = cleaned.strip("`")
            if cleaned.lower().startswith("json"):
                cleaned = cleaned[4:].strip()
        start = cleaned.find("{")
        end = cleaned.rfind("}")
        if start >= 0 and end >= start:
            cleaned = cleaned[start : end + 1]
        payload = json.loads(cleaned)
        if not isinstance(payload, dict):
            raise ValueError("router response must be a JSON object")
        return payload


def _clean_short_title(value: Any) -> str | None:
    if value in (None, ""):
        return None
    title = " ".join(str(value).strip().split())
    return title[:40] or None