core/task_detector.py

"""Task detection - analyze requests to determine required capabilities."""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import Any

from loguru import logger

from core.anthropic.content import get_block_attr


# Keywords that indicate specific task types
CODING_KEYWORDS = {
    "python",
    "javascript",
    "typescript",
    "java",
    "c++",
    "cpp",
    "golang",
    "rust",
    "ruby",
    "php",
    "swift",
    "kotlin",
    "sql",
    "html",
    "css",
    "react",
    "vue",
    "angular",
    "node",
    "django",
    "flask",
    "fastapi",
    "spring",
    "function",
    "class",
    "method",
    "api",
    "endpoint",
    "database",
    "query",
    "algorithm",
    "debug",
    "error",
    "fix",
    "implement",
    "create",
    "write",
    "code",
    "programming",
    "script",
    "module",
    "import",
    "export",
    "def ",
    "const ",
    "let ",
    "var ",
    "function ",
    "async ",
    "await ",
}

REASONING_KEYWORDS = {
    "analyze",
    "analysis",
    "reason",
    "why",
    "how",
    "explain",
    "compare",
    "contrast",
    "evaluate",
    "assess",
    "conclude",
    "deduce",
    "infer",
    "logic",
    "proof",
    "theorem",
    "hypothesis",
    "synthesize",
    "strategy",
    "think",
    "solve",
    "derive",
    "calculate",
    "compute",
    "math",
    "equation",
    "formula",
    "solution",
    "optimal",
    "best",
    "improve",
    "optimize",
    "design",
    "architecture",
    "system",
    "plan",
    "decision",
    "recommend",
}

VISION_KEYWORDS = {
    "image",
    "picture",
    "photo",
    "screenshot",
    "diagram",
    "chart",
    "graph",
    "visual",
    "see",
    "look at",
    "describe what",
    "what's in",
    "identify",
    "recognize",
    "detect",
    "object",
    "scene",
    "face",
    "text in image",
}


@dataclass(frozen=True, slots=True)
class TaskRequirements:
    """Detected requirements for a request."""

    requires_vision: bool = False
    requires_coding: bool = False
    requires_reasoning: bool = False
    requires_general_text: bool = True
    confidence: float = 0.0  # 0-1 confidence in detection

    @property
    def required_capabilities(self) -> set[str]:
        caps = set()
        if self.requires_vision:
            caps.add("vision")
        if self.requires_coding:
            caps.add("coding")
        if self.requires_reasoning:
            caps.add("reasoning")
        if self.requires_general_text:
            caps.add("general_text")
        return caps


class TaskDetector:
    """Analyze request messages to detect required capabilities."""

    def detect_requirements(self, messages: list[Any]) -> TaskRequirements:
        """Analyze messages and return required capabilities."""
        has_vision = False
        has_coding = False
        has_reasoning = False
        total_text = ""

        for msg in messages:
            # Handle both dict and object message formats
            if isinstance(msg, dict):
                content = msg.get("content")
            elif hasattr(msg, "content"):
                content = msg.content
            else:
                continue

            if isinstance(content, str):
                total_text += content.lower() + " "
            elif isinstance(content, list):
                for block in content:
                    b_type = get_block_attr(block, "type") or ""

                    # Check for image content
                    if b_type == "image":
                        has_vision = True
                        logger.debug("TaskDetector: Found image in message")

                    # Get text content
                    if b_type == "text":
                        text = get_block_attr(block, "text", "") or ""
                        total_text += text.lower() + " "

        # Analyze text for keywords
        if total_text:
            has_coding = self._detect_coding(total_text)
            has_reasoning = self._detect_reasoning(total_text)

        # Calculate confidence
        confidence = self._calculate_confidence(
            has_vision, has_coding, has_reasoning, total_text
        )

        # Default to general text if nothing detected
        if not has_vision and not has_coding and not has_reasoning:
            has_general = True

        result = TaskRequirements(
            requires_vision=has_vision,
            requires_coding=has_coding,
            requires_reasoning=has_reasoning,
            requires_general_text=True,
            confidence=confidence,
        )

        logger.info(
            "TaskDetector: detected caps={} confidence={:.2f}",
            result.required_capabilities,
            confidence,
        )

        return result

    def _detect_coding(self, text: str) -> bool:
        """Detect if request requires coding capabilities."""
        # Check exact word matches first
        words = set(re.findall(r"\b\w+\b", text))
        coding_matches = words & CODING_KEYWORDS
        if len(coding_matches) >= 2:
            return True

        # Also check for substring matches (e.g., "python" in "write python code")
        for keyword in CODING_KEYWORDS:
            if keyword in text:
                # Found one keyword as substring, check for another
                remaining = text.replace(keyword, "")
                for kw2 in CODING_KEYWORDS:
                    if kw2 in remaining and kw2 != keyword:
                        return True
                # Also check for programming patterns
                if any(
                    pat in text
                    for pat in [
                        "def ",
                        "function ",
                        "class ",
                        "import ",
                        "const ",
                        "let ",
                        "var ",
                        "()",
                        "=>",
                    ]
                ):
                    return True

        return False

    def _detect_reasoning(self, text: str) -> bool:
        """Detect if request requires reasoning capabilities."""
        words = set(re.findall(r"\b\w+\b", text))
        reasoning_matches = words & REASONING_KEYWORDS
        if len(reasoning_matches) >= 1:
            return True
        # Also check substring
        for keyword in REASONING_KEYWORDS:
            if keyword in text:
                return True
        return False

    def _calculate_confidence(
        self,
        has_vision: bool,
        has_coding: bool,
        has_reasoning: bool,
        text: str,
    ) -> float:
        """Calculate confidence in the detection."""
        if has_vision:
            return 0.95  # Image detection is reliable
        if has_coding or has_reasoning:
            # More text = more confident
            word_count = len(text.split())
            base = 0.7
            if word_count > 50:
                base = 0.8
            if word_count > 100:
                base = 0.85
            return base
        return 0.5  # Default confidence for general text

    def get_priority_hint(self, requirements: TaskRequirements) -> str:
        """Get a hint for model priority based on requirements."""
        if requirements.requires_vision:
            return "vision"
        if requirements.requires_coding:
            return "coding"
        if requirements.requires_reasoning:
            return "reasoning"
        return "balanced"