symbiont/symbiont/router.py

"""
Router: Classifies tasks and picks the cheapest capable model.

The router itself runs on Haiku (cheapest) to minimize overhead.
It examines the task and assigns a model tier based on complexity.
"""

import json
import logging
from typing import Optional
from .dispatcher import ModelTier, dispatch

logger = logging.getLogger(__name__)

# Classification prompt — this runs on Haiku to keep costs minimal
CLASSIFIER_SYSTEM_PROMPT = """You are a task classifier for an AI orchestration system.
Your job is to assess a task and decide which model tier should handle it.

Tiers:
- HAIKU: Simple tasks. Template filling, reformatting, extraction from structured data,
  classification, boilerplate generation, simple Q&A. ~70% of tasks.
- SONNET: Medium tasks. Content writing, code generation, moderate analysis,
  customer-facing communication, summarization of complex documents. ~25% of tasks.
- OPUS: Complex tasks. Strategic decisions, novel problem-solving, multi-step reasoning,
  quality review of important outputs, creative/nuanced work. ~5% of tasks.

Respond with ONLY a JSON object:
{"tier": "HAIKU"|"SONNET"|"OPUS", "confidence": 0.0-1.0, "reason": "brief explanation"}
"""


def classify_task(task_description: str) -> tuple[ModelTier, float, str]:
    """
    Use Haiku to classify which tier should handle this task.
    Returns (tier, confidence, reason).

    If Haiku is unavailable (rate-limited), falls back to a simple
    heuristic classifier.
    """
    result = dispatch(
        prompt=f"Classify this task:\n\n{task_description}",
        tier=ModelTier.HAIKU,
        system_prompt=CLASSIFIER_SYSTEM_PROMPT,
    )

    if result.success:
        try:
            # Try to parse the JSON response
            text = result.output.strip()
            # Handle markdown code blocks
            if text.startswith("```"):
                text = text.split("\n", 1)[1].rsplit("```", 1)[0].strip()

            classification = json.loads(text)
            tier_name = classification.get("tier", "SONNET").upper()
            confidence = float(classification.get("confidence", 0.5))
            reason = classification.get("reason", "no reason given")

            tier = {
                "HAIKU": ModelTier.HAIKU,
                "SONNET": ModelTier.SONNET,
                "OPUS": ModelTier.OPUS,
            }.get(tier_name, ModelTier.SONNET)

            logger.info(f"Classified as {tier.value} (confidence={confidence}): {reason}")
            return tier, confidence, reason

        except (json.JSONDecodeError, KeyError, ValueError) as e:
            logger.warning(f"Failed to parse classification: {e}, falling back to heuristic")

    # Fallback: simple heuristic classifier (free, no LLM needed)
    return _heuristic_classify(task_description)


def _heuristic_classify(task: str) -> tuple[ModelTier, float, str]:
    """
    Dead-simple keyword heuristic. Used when Haiku is unavailable.
    Costs zero tokens — pure Python.
    """
    task_lower = task.lower()

    # OPUS indicators
    opus_signals = [
        "design", "architect", "strategy", "analyze complex", "review",
        "creative", "novel", "brainstorm", "plan", "evaluate",
        "multi-step", "reasoning", "trade-off", "nuanced",
    ]
    if any(signal in task_lower for signal in opus_signals):
        return ModelTier.OPUS, 0.4, "heuristic: complexity keywords detected"

    # HAIKU indicators
    haiku_signals = [
        "extract", "format", "convert", "classify", "list",
        "template", "simple", "boilerplate", "parse", "json",
        "reformat", "count", "sort", "filter", "summarize short",
    ]
    if any(signal in task_lower for signal in haiku_signals):
        return ModelTier.HAIKU, 0.5, "heuristic: simplicity keywords detected"

    # Default to SONNET
    return ModelTier.SONNET, 0.3, "heuristic: defaulting to sonnet"


def route_task(
    task: str,
    system_prompt: Optional[str] = None,
    force_tier: Optional[ModelTier] = None,
) -> dict:
    """
    Full routing pipeline: classify the task, dispatch to the right model,
    return the result with metadata.
    """
    if force_tier:
        tier = force_tier
        confidence = 1.0
        reason = "forced by caller"
    else:
        tier, confidence, reason = classify_task(task)

    from .dispatcher import dispatch_with_fallback

    result = dispatch_with_fallback(
        prompt=task,
        preferred_tier=tier,
        system_prompt=system_prompt,
    )

    return {
        "output": result.output,
        "success": result.success,
        "model_used": result.model.value,
        "model_requested": tier.value,
        "classification_confidence": confidence,
        "classification_reason": reason,
        "elapsed_seconds": result.elapsed_seconds,
        "estimated_cost_usd": result.estimated_cost_usd,
        "rate_limited": result.rate_limited,
        "error": result.error,
    }