step_routing.py

Go to the documentation of this file.

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# oc_chat/core/step_routing.py
 
 
"""
Per-ReAct-iteration model, reasoning effort, and max_tokens routing (Phase 3).
 
Purpose
-------
Maps tool batches to profiles described in ``07 Token Efficiency…`` —
flash for cheap deterministic reads, pro+thinking for exploration and
writes, and the session model with strongest effort for planning and
final synthesis. Lets every iteration of the ReAct loop pick the
right model/effort/max_tokens triple without baking provider-
specific knowledge into ``main.py``.
 
Responsibilities
----------------
- :func:`classify_tool_batch` — bucket the most recent assistant
  message's tool calls into ``simple``, ``simple_thinking``,
  ``thinking``, ``thinking_write``, or ``complex``.
- :func:`resolve_iteration_routing` — pick model, effort, extra_body,
  and max_tokens for the next iteration, using the active provider's
  ``cheap`` and ``reasoning`` tiers when available.
- :func:`resolve_output_tokens` — derive a bucket's completion cap
  from the model's declared ``max_output_tokens``, the operator
  ceiling, and the ``/budget`` profile.
- :func:`scale_max_tokens_for_profile` — apply the user's ``/budget``
  profile (cheap / balanced / deep) to a base completion cap.
- :func:`last_assistant_tool_names` — extract tool names from the
  last assistant message that issued tool_calls.
 
Diagnostics
-----------
Domain: OC_CHAT.MAIN
Levels:
    L2 — lifecycle events (entry/exit boundaries, currently unused).
    L3 — per-bucket routing decisions in resolve_iteration_routing.
    L4 — classify_tool_batch result dumps.
 
Contracts
---------
- :class:`IterationRouting` is frozen — never mutate after return.
- Tier-preferred models are sanity-checked against the model
  manager's known list; unknown ids fall back to the session model.
- ``scale_max_tokens_for_profile`` clamps to documented bounds
  (``cheap`` ≥ 1024; every profile ≤ the ceiling passed in, which
  ``resolve_output_tokens`` sets to the tighter of the model's declared
  ``max_output_tokens`` and the operator ceiling).
- A completion cap is never 0: an undeclared ``max_output_tokens``
  resolves to :data:`~oc_chat.core.config.UNDECLARED_OUTPUT_TOKENS`,
  and every bucket has the :data:`_MIN_BUCKET_OUTPUT_TOKENS` floor.
- The cap bounds reasoning AND answer together on a thinking model.
  Sizing it for the answer alone is what let the model spend its whole
  budget on chain-of-thought and return no content at all.
 
Dependencies
------------
- oc_diagnostics (structured debug logging via _dbg)
- oc_chat.core.config (operator ceiling + undeclared-capability fallback)
- oc_chat.core.defaults (default - the per-bucket output shares and the minimum bucket output floor)
 
See Also
--------
oc_chat/docs/conventions/CONTRACTS.md §1.6 IterationRouting — owner of the per-
iteration API-parameters dataclass.
"""
from __future__ import annotations
 
from dataclasses import dataclass
from typing import Any
 
from oc_diagnostics import _dbg
 
from oc_chat.core.config import (
    UNDECLARED_OUTPUT_TOKENS,
    get_output_tokens_ceiling,
)
from oc_chat.core.defaults import default
 
#: Each bucket's share of the resolved output budget. Every iteration can be
#: the one that produces the final answer — including the cheap post-read hop —
#: so no bucket gets a share that could truncate an ordinary reply. The
#: gradient only reflects that a deterministic follow-up needs less room than
#: a planning or writing hop, not that it should be starved.
_BUCKET_OUTPUT_SHARE: dict[str, float] = {
    "planning":        default("bucket_share_planning"),
    "thinking":        default("bucket_share_thinking"),
    "thinking_write":  default("bucket_share_thinking_write"),
    "simple_thinking": default("bucket_share_simple_thinking"),
    "simple":          default("bucket_share_simple"),
}
 
#: Floor under every bucket, so a provider declaring a tiny output limit still
#: leaves room for a usable reply rather than a one-line stub.
_MIN_BUCKET_OUTPUT_TOKENS: int = default("min_bucket_output_tokens")
 
# Tier names asked of the active provider — provider.tier("cheap") /
# provider.tier("reasoning") return concrete model ids (DeepSeek maps
# these to deepseek-v4-flash / deepseek-v4-pro respectively). When a
# provider has no model for a given tier, _pick_model falls back to the
# session model.
_CHEAP_TIER = "cheap"
_REASONING_TIER = "reasoning"
 
# Deterministic / heavy I/O — flash, no thinking, low effort, tight completion cap.
SIMPLE_TOOLS = frozenset({"read_file", "run_command"})
# Non-deterministic listing / search — pro, thinking, low effort.
SIMPLE_THINKING_TOOLS = frozenset({"list_dir", "search_files"})
# Mutations, web, memory, planning — pro, thinking, elevated effort.
THINKING_TOOLS = frozenset({
    "write_file",
    "fetch_url",
    "web_search",
    "remember",
    "recall",
    "create_plan",
    "get_status",
    "mark_step",
    "stop_plan",
    "oct_new",
    # Mutating oct_git_* wrappers — read-only ones (status, check, commit_msg)
    # stay out because they're cheap and don't need the reasoning model.
    "oct_git_init",
    "oct_git_add",
    "oct_git_commit",
    "oct_git_branch",
    "oct_git_restore",
    "oct_git_ignore",
    "oct_git_changelog",
    "oct_git_hooks",
})
 
 
@dataclass(frozen=True)
class IterationRouting:
    """Resolved API parameters for one chat-completion call."""
 
    model: str
    effort: str | None
    extra_body: dict[str, Any] | None
    max_tokens: int
    # would_switch_model / cheap_model are VESTIGIAL pending the opt-in
    # cheap-model offload redesign: plan 132 stopped them triggering a
    # sub-loop, and plan 134 locks sub-loops to the session model. Retained so
    # the redesign has the signal; do not re-wire them to pick a model without
    # that redesign + a re-trigger guard.
    would_switch_model: bool = False
    cheap_model: str | None = None
 
 
def _message_text(message: dict) -> str:
    """Return message content as text, treating None/non-text as empty."""
    content = message.get("content")
    return content if isinstance(content, str) else ""
 
 
def _known(model_manager, name: str | None) -> bool:
    return bool(name) and name in model_manager.known_models()
 
 
def _pick_model(model_manager, preferred: str | None, session_model: str) -> tuple[str, bool, str | None]:
    """Return ``(model, would_switch, cheap_model)``.
 
    Always returns *session_model* for the model (preserving prefix cache).
    If *preferred* differs from *session_model*, ``would_switch=True`` and
    *cheap_model* is set.
 
    NOTE (plan 132 + 134): ``would_switch`` / ``cheap_model`` are NO LONGER a
    sub-loop trigger, and sub-loops no longer run on ``cheap_model`` — the
    interim model-lock runs them on the **session model** (no silent flash
    out-of-band usage). These fields are retained only for the future
    targeted, opt-in cheap-model offload redesign; nothing consumes
    ``cheap_model`` to pick a model today.
    """
    if preferred and _known(model_manager, preferred) and preferred != session_model:
        return session_model, True, preferred
    if _known(model_manager, session_model):
        return session_model, False, None
    keys = list(model_manager.known_models().keys())
    fallback = keys[0] if keys else session_model
    return fallback, False, None
 
 
def _caps(model_manager, name: str) -> dict:
    return model_manager.capabilities(name)
 
 
def _effort_for_model(model_manager, model: str, desired: str) -> str | None:
    c = _caps(model_manager, model)
    if not c.get("supports_reasoning"):
        return None
    return desired
 
 
def _thinking_body(model_manager, model: str, effort: str | None) -> dict[str, Any] | None:
    """Delegate to the active provider's ``build_extra_body`` hook.
 
    For DeepSeek this returns ``{"thinking": {"type": "enabled"}}`` when
    the model's capabilities support thinking. Other providers shape the
    payload however they need (or return ``None`` if they have no
    equivalent concept).
    """
    return model_manager.provider.build_extra_body(
        model, enable_thinking=True, effort=effort
    )
 
 
def classify_tool_batch(tool_names: list[str]) -> str:
    """Return a profile key: ``simple`` | ``simple_thinking`` | ``thinking`` | ``thinking_write``."""
    func = "classify_tool_batch"
    if not tool_names:
        result = "complex"
    else:
        names = set(tool_names)
        if names & SIMPLE_THINKING_TOOLS and not (names & THINKING_TOOLS) and names <= (SIMPLE_TOOLS | SIMPLE_THINKING_TOOLS):
            result = "simple_thinking"
        elif names <= SIMPLE_TOOLS:
            result = "simple"
        elif "write_file" in names:
            result = "thinking_write"
        elif names & THINKING_TOOLS:
            result = "thinking"
        else:
            result = "simple_thinking"
    _dbg("OC_CHAT.MAIN", func, f"tools={tool_names} -> {result}", 4)
    return result
 
 
def scale_max_tokens_for_profile(
    base: int, token_profile: str, *, ceiling: int | None = None
) -> int:
    """Scale completion ``max_tokens`` from ``/budget`` token profile.
 
    *ceiling* bounds the ``deep`` profile. It defaults to the operator
    ceiling; :func:`resolve_output_tokens` passes the tighter of that and the
    model's own declared limit, so ``deep`` can never ask a provider for more
    output than it says it can produce.
    """
    prof = (token_profile or "balanced").strip().lower()
    cap = ceiling if ceiling is not None else get_output_tokens_ceiling()
    if prof == "cheap":
        return max(1024, int(base * 0.62))
    if prof == "deep":
        return min(cap, int(base * 1.18))
    return min(cap, base)
 
 
def resolve_output_tokens(
    model_manager: Any, model: str, bucket: str, token_profile: str
) -> int:
    """Resolve this iteration's completion cap from the model's own capability.
 
    Replaces the former hardcoded per-bucket literals (12_288 / 8192 / 6144 /
    4096), which were sized for a small model and silently truncated a
    large one. The budget is the model's declared ``max_output_tokens``,
    bounded by the operator ceiling, then apportioned by bucket and scaled by
    the ``/budget`` profile.
    """
    func = "resolve_output_tokens"
    ceiling = get_output_tokens_ceiling()
    declared = int(_caps(model_manager, model).get("max_output_tokens") or 0)
    budget = min(declared if declared > 0 else UNDECLARED_OUTPUT_TOKENS, ceiling)
    share = _BUCKET_OUTPUT_SHARE.get(bucket, 1.0)
    base = max(_MIN_BUCKET_OUTPUT_TOKENS, int(budget * share))
    # ``budget`` is the ceiling passed downstream, so the floor lifts a starved
    # bucket but a model declaring less than the floor still caps at its own
    # limit — asking for more output than a provider says it can produce risks
    # a rejected request.
    out = scale_max_tokens_for_profile(base, token_profile, ceiling=budget)
    _dbg("OC_CHAT.MAIN", func,
         f"bucket={bucket} model={model} declared={declared} ceiling={ceiling} "
         f"budget={budget} share={share} profile={token_profile} -> {out}", 3)
    return out
 
 
def should_use_sub_loop(task: str, iteration_count: int, recent_messages: list[dict]) -> bool:
    """Determine if a task is complex enough to warrant sub-loop usage."""
    # Heuristics for sub-loop triggering
    complexity_indicators = [
        "multiple steps",
        "analyze", 
        "investigate",
        "complex",
        "multiple files",
        "search",
        "compare",
        "debug",
        "fix",
        "implement",
    ]
    
    # Check if task contains complexity indicators
    task_lower = task.lower()
    has_complexity = any(indicator in task_lower for indicator in complexity_indicators)
    
    # Consider iteration count - if we're stuck, use sub-loop (but not too early)
    is_stuck = iteration_count > 8  # Increased from 5 to prevent premature sub-loop
    
    # Consider recent message patterns - look for actual repetition patterns
    recent_user_messages = [m for m in recent_messages if m.get("role") == "user"]
    has_repetition = len(recent_user_messages) > 3 and any(
        "again" in _message_text(m).lower()
        for m in recent_user_messages[-3:]
    )
    
    # Prevent sub-loop triggering on first few iterations to allow main loop to work
    if iteration_count < 3:
        return False
    
    # Only trigger sub-loop if we have a clear reason and haven't used it recently
    sub_loop_used_recently = any(
        "[Sub-loop result]" in _message_text(m)
        for m in recent_messages[-5:]  # Check last 5 messages
    )
    
    if sub_loop_used_recently:
        return False
    
    # Additional safeguard: prevent sub-loop if we're already stuck in a loop
    if iteration_count > 15:  # Emergency cutoff before hitting 30
        return False
    
    # Check for repeated error patterns that indicate infinite loop
    recent_errors = [
        m for m in recent_messages[-10:] 
        if m.get("role") == "assistant" and 
        any(error in _message_text(m).lower() for error in ["error", "failed", "not found"])
    ]
    
    if len(recent_errors) > 3:  # Too many errors recently
        return False
    
    return has_complexity or is_stuck or has_repetition
 
 
def last_assistant_tool_names(messages: list[dict]) -> list[str]:
    """Tool function names from the most recent assistant message that issued tool_calls."""
    for m in reversed(messages):
        if m.get("role") != "assistant":
            continue
        tcs = m.get("tool_calls") or []
        out: list[str] = []
        for tc in tcs:
            fn = tc.get("function") or {}
            n = fn.get("name") or ""
            if n:
                out.append(n)
        if out:
            return out
    return []
 
 
# OCT-LINT: disable=no_hardcoded_secrets reason="Config defaults and tier names, not secrets"
def resolve_iteration_routing(
    *,
    session_model: str,
    session_reasoning_effort: str,
    model_manager,
    iteration: int,
    is_finalize_nudge: bool,
    messages: list[dict],
    token_profile: str = "balanced",
) -> IterationRouting:
    """Pick model / effort / thinking / max_tokens for this ReAct iteration."""
    func = "resolve_iteration_routing"
    tp = token_profile or "balanced"
    provider = model_manager.provider
    reasoning_model = provider.tier(_REASONING_TIER)
    cheap_model = provider.tier(_CHEAP_TIER)
 
    # Planning (first hop) and forced final answer: strongest profile on the
    # provider's reasoning tier (or session model if no reasoning tier).
    if iteration == 0 or is_finalize_nudge:
        model, would_switch, cheap_model = _pick_model(model_manager, reasoning_model, session_model)
        # Planning/finalize always uses the main model regardless of would_switch signal
        would_switch = False
        cheap_model = None
        effort = _effort_for_model(model_manager, model, "high")
        if is_finalize_nudge:
            effort = _effort_for_model(model_manager, model, "high") or effort
        else:
            # First iteration: allow session preference up to high for "planning".
            pref = session_reasoning_effort if session_reasoning_effort in ("low", "medium", "high") else "high"
            e2 = _effort_for_model(model_manager, model, pref)
            effort = e2 if e2 is not None else effort
        extra = _thinking_body(model_manager, model, effort)
        mt = resolve_output_tokens(model_manager, model, "planning", tp)
        _dbg("OC_CHAT.MAIN", func,
             f"planning/finalize: model={model} effort={effort} max_tokens={mt} finalize={is_finalize_nudge}", 3)
        return IterationRouting(model=model, effort=effort, extra_body=extra, max_tokens=mt, would_switch_model=would_switch, cheap_model=cheap_model)
 
    tool_names = last_assistant_tool_names(messages)
    bucket = classify_tool_batch(tool_names)
 
    if bucket == "simple":
        model, would_switch, cheap_model = _pick_model(model_manager, cheap_model, session_model)
        effort = _effort_for_model(model_manager, model, "low")
        mt = resolve_output_tokens(model_manager, model, "simple", tp)
        _dbg("OC_CHAT.MAIN", func,
             f"bucket=simple: model={model} effort={effort} max_tokens={mt}", 3)
        return IterationRouting(model=model, effort=effort, extra_body=None, max_tokens=mt, would_switch_model=would_switch, cheap_model=cheap_model)
 
    if bucket == "simple_thinking":
        model, would_switch, cheap_model = _pick_model(model_manager, reasoning_model, session_model)
        effort = _effort_for_model(model_manager, model, "low")
        extra = _thinking_body(model_manager, model, effort)
        mt = resolve_output_tokens(model_manager, model, "simple_thinking", tp)
        _dbg("OC_CHAT.MAIN", func,
             f"bucket=simple_thinking: model={model} effort={effort} max_tokens={mt}", 3)
        return IterationRouting(model=model, effort=effort, extra_body=extra, max_tokens=mt, would_switch_model=would_switch, cheap_model=cheap_model)
 
    if bucket == "thinking_write":
        model, would_switch, cheap_model = _pick_model(model_manager, reasoning_model, session_model)
        effort = _effort_for_model(model_manager, model, "high")
        extra = _thinking_body(model_manager, model, effort)
        mt = resolve_output_tokens(model_manager, model, "thinking_write", tp)
        _dbg("OC_CHAT.MAIN", func,
             f"bucket=thinking_write: model={model} effort={effort} max_tokens={mt}", 3)
        return IterationRouting(model=model, effort=effort, extra_body=extra, max_tokens=mt, would_switch_model=would_switch, cheap_model=cheap_model)
 
    # "thinking" — code/doc style tools without write_file singled out
    model, would_switch, cheap_model = _pick_model(model_manager, reasoning_model, session_model)
    effort = _effort_for_model(model_manager, model, "high")
    extra = _thinking_body(model_manager, model, effort)
    mt = resolve_output_tokens(model_manager, model, "thinking", tp)
    _dbg("OC_CHAT.MAIN", func,
         f"bucket=thinking: model={model} effort={effort} max_tokens={mt}", 3)
    return IterationRouting(model=model, effort=effort, extra_body=extra, max_tokens=mt, would_switch_model=would_switch, cheap_model=cheap_model)