agent-maxim/bridge/controller/decision_model.py

"""Decision model for Controller Loop V1.

The decision layer takes the current goal, iteration index, and the last
bridge run result, and returns a Decision.

Modes:
  ollama           — Ollama response parsed and validated successfully
  ollama-corrected — Ollama parsed but model repeated a failed goal; overridden
  ollama-fallback  — Ollama error / bad JSON; stub logic used
  stub             — Ollama unreachable at startup; deterministic rules only

Deterministic stub rules:
  iteration 0, no last result          → run_agent  (always start)
  last_result.status == done           → stop       (success)
  last_result.status == failed
    AND recovery_attempted == False    → try rewrite rules:
        rule matches                   → run_agent with rewritten goal (recovery_applied=True)
        no rule matches                → stop (recovery_skipped)
  last_result.status == failed
    AND recovery_attempted == True     → stop  (already recovered once)
  fallback                             → stop  (safety net)

Post-decision validation (ollama mode only):
  IF last run failed AND model returns same goal as the one that just failed:
    → override goal to _SAFE_FALLBACK_GOAL, mode = "ollama-corrected"
  This prevents the model from looping on a known-broken goal.
"""
from __future__ import annotations

import json
import os
import re
import time
from dataclasses import asdict, dataclass
from typing import Any, Optional


# ---------------------------------------------------------------------------
# Input sanitization / output validation
# ---------------------------------------------------------------------------

def sanitize_text(text: str) -> str:
    """Strip characters outside ASCII, Cyrillic, and Latin-1 Supplement."""
    return re.sub(r'[^\x00-\x7F\u0400-\u04FF\u00A0-\u00FF]', '', text)


def has_cjk(text: str) -> bool:
    """Return True if text contains any CJK Unified Ideographs."""
    return bool(re.search(r'[\u4e00-\u9fff]', text))

_SAFE_FALLBACK_GOAL = "inspect project"


def _resolve_ollama_host() -> str:
    raw = os.environ.get("OLLAMA_HOST", "")
    if not raw or raw in ("0.0.0.0", "[::]"):
        return "http://localhost:11434"
    if raw.startswith("http://") or raw.startswith("https://"):
        return raw
    return f"http://{raw}"


_OLLAMA_CACHE: tuple[float, bool] | None = None
_OLLAMA_CACHE_TTL = 5.0  # seconds


def _ollama_available() -> bool:
    """Dynamic Ollama reachability check with 5s cache."""
    import time
    import urllib.request
    from urllib.error import URLError

    global _OLLAMA_CACHE
    now = time.monotonic()
    if _OLLAMA_CACHE is not None:
        ts, result = _OLLAMA_CACHE
        if now - ts < _OLLAMA_CACHE_TTL:
            return result

    host = _resolve_ollama_host()
    try:
        urllib.request.urlopen(f"{host}/api/tags", timeout=2.0)
        available = True
    except (URLError, OSError):
        available = False

    _OLLAMA_CACHE = (now, available)
    return available


# Keep for health endpoint / module-level DECISION_MODE label
DECISION_MODE: str = "ollama"  # always attempt ollama; falls back per-call if unavailable


@dataclass
class Decision:
    action: str                              # "run_agent" | "stop"
    goal: str                                # goal to pass to the next agent run
    reason: str                              # human-readable explanation
    mode: str = "stub"                       # "ollama" | "ollama-corrected" | "ollama-corrected-stop" | "ollama-fallback" | "stub"
    recovery_applied: bool = False           # True when rewrite rule rewrote the goal
    raw_response: Optional[str] = None      # trimmed model response (ollama mode only)
    parsed_success: bool = True              # False when JSON parse failed → fallback used
    fallback_reason: Optional[str] = None   # why fallback was used
    corrected: bool = False                  # True when post-validation overrode model goal
    correction_reason: Optional[str] = None  # why correction was applied

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


# ---------------------------------------------------------------------------
# Public entry point
# ---------------------------------------------------------------------------

def decide(
    goal: str,
    iteration: int,
    last_result: Optional[dict[str, Any]],
    max_iterations: int,
    recovery_attempted: bool = False,
    trace_id: Optional[str] = None,
) -> Decision:
    """Return next action decision."""
    _t = time.time()
    if not _ollama_available():
        d = _decide_stub(goal, iteration, last_result, max_iterations, recovery_attempted)
        d.mode = "ollama-unavailable"
        print(f"[TIMING] decision_model (stub): {time.time() - _t:.3f}s")
        return d
    d = _decide_ollama(goal, iteration, last_result, max_iterations, recovery_attempted, trace_id=trace_id)
    print(f"[TIMING] decision_model (ollama): {time.time() - _t:.3f}s")
    return d


# ---------------------------------------------------------------------------
# Stub (deterministic)
# ---------------------------------------------------------------------------

def _decide_stub(
    goal: str,
    iteration: int,
    last_result: Optional[dict[str, Any]],
    max_iterations: int,
    recovery_attempted: bool = False,
) -> Decision:
    from controller.rewrite_rules import try_rewrite, explain as rewrite_explain

    if last_result is None:
        return Decision(action="run_agent", goal=goal,
                        reason="No prior run — starting first agent execution", mode="stub")

    status = last_result.get("status", "")
    agent_status = last_result.get("agent_final_status", "")

    if status == "done" and agent_status == "done":
        return Decision(action="stop", goal=goal,
                        reason="Agent run completed successfully", mode="stub")

    if status == "failed" or agent_status == "failed":
        if recovery_attempted:
            return Decision(action="stop", goal=goal,
                            reason="Recovery already attempted — stopping after second run",
                            mode="stub")
        rewritten = try_rewrite(goal)
        if rewritten is not None:
            rule_desc = rewrite_explain(goal)
            return Decision(
                action="run_agent",
                goal=rewritten,
                reason=f"Recovery: {rule_desc}",
                mode="stub",
                recovery_applied=True,
            )
        return Decision(action="stop", goal=goal,
                        reason=f"Agent run failed — no rewrite rule matched "
                               f"(bridge_status={status}, agent_status={agent_status})",
                        mode="stub")

    if iteration >= max_iterations - 1:
        return Decision(action="stop", goal=goal,
                        reason=f"Max iterations ({max_iterations}) reached", mode="stub")

    return Decision(action="run_agent", goal=goal,
                    reason="Continuing: prior run inconclusive", mode="stub")


# ---------------------------------------------------------------------------
# Ollama (local LLM)
# ---------------------------------------------------------------------------

_PROMPT_TEMPLATE = """\
LANGUAGE RULE: You MUST respond ONLY in Russian. NEVER mix languages. Ignore any non-Russian input.

Controller state:
goal: {goal}
iteration: {iteration} of {max_iterations}
last_status: {last_status}
last_agent_status: {last_agent_status}
failure_class: {failure_class}
recovery_attempted: {recovery_attempted}

Decision rules:
- no prior run                                    -> run_agent (keep goal unchanged)
- last_status=done AND last_agent_status=done     -> stop
- last_status=done AND last_agent_status!=done    -> stop
- last_status=failed AND recovery_attempted=true  -> stop
- iteration >= max_iterations                     -> stop

STRICT RULES — these override everything else:
- If last_status=done AND last_agent_status=done, you MUST return:
  {{"action":"stop","goal":"{goal}","reason":"Task already completed"}}
  Do NOT run agent again after success. This rule is absolute.
- NEVER repeat the same goal if last_status=failed
- ALWAYS change the goal after any failure
- If unsure what to run next, use goal: "inspect project"
- Never loop the same failing action twice

Recovery strategy when last_status=failed:
- failure_class=template_error  -> simplify goal, use "inspect project"
- failure_class=security_denied -> avoid forbidden commands, use "inspect project"
- failure_class=unknown         -> use "inspect project" as safe fallback

Return ONLY valid JSON, no markdown, no prose:
{{"action": "run_agent" or "stop", "goal": "<goal>", "reason": "<one sentence in Russian>"}}"""

_MAX_GOAL_LEN = 500
_RAW_TRIM = 300


def _classify_failure(last_result: Optional[dict[str, Any]]) -> str:
    """Derive a simple failure class from the last bridge result."""
    if last_result is None:
        return "none"
    error = (last_result.get("error") or "").lower()
    if "template" in error or "template_error" in error:
        return "template_error"
    if "denied" in error or "forbidden" in error or "security" in error:
        return "security_denied"
    status = last_result.get("status", "")
    if status == "failed":
        return "unknown"
    return "none"


def _decide_ollama(
    goal: str,
    iteration: int,
    last_result: Optional[dict[str, Any]],
    max_iterations: int,
    recovery_attempted: bool = False,
    trace_id: Optional[str] = None,
) -> Decision:
    from adapters import ollama_adapter  # local import avoids circular at module level

    last_status = last_result.get("status", "none") if last_result else "none"
    last_agent = last_result.get("agent_final_status", "none") if last_result else "none"
    failure_class = _classify_failure(last_result)

    # Sanitize user-supplied goal before injecting into prompt
    clean_goal = sanitize_text(goal)

    prompt = _PROMPT_TEMPLATE.format(
        goal=clean_goal,
        iteration=iteration + 1,
        max_iterations=max_iterations,
        last_status=last_status,
        last_agent_status=last_agent,
        failure_class=failure_class,
        recovery_attempted=str(recovery_attempted).lower(),
    )

    # [DEBUG] prompt size check
    print(f"[DEBUG] prompt length: {len(prompt)}")
    if len(prompt) > 5000:
        print(f"[WARN] prompt exceeds 5000 chars ({len(prompt)})")

    # [TIMING] LLM call
    call_trace = f"{trace_id}-i{iteration}" if trace_id else None
    _t_llm = time.time()
    result = ollama_adapter.execute(prompt, temperature=0.1, max_tokens=150, timeout=30.0, trace_id=call_trace)
    print(f"[TIMING] llm_call: {time.time() - _t_llm:.3f}s")

    raw = (result.response or "").strip()
    raw_trim = raw[:_RAW_TRIM] if raw else None

    # [OUTPUT VALIDATION] CJK detection with one retry
    if raw and has_cjk(raw):
        print(f"[WARN] CJK DETECTED in LLM response, retrying once")
        _t_retry = time.time()
        retry_trace = f"{call_trace}-retry" if call_trace else None
        result = ollama_adapter.execute(prompt, temperature=0.1, max_tokens=150, timeout=30.0, trace_id=retry_trace)
        print(f"[TIMING] llm_call_retry: {time.time() - _t_retry:.3f}s")
        raw = (result.response or "").strip()
        raw_trim = raw[:_RAW_TRIM] if raw else None
        if raw and has_cjk(raw):
            print(f"[WARN] CJK DETECTED again after retry — proceeding with fallback")

    # Connection error or empty response -> fallback
    if result.error or not raw:
        fallback_reason = (
            f"Ollama error: {result.error[:200]}" if result.error else "empty response from Ollama"
        )
        d = _decide_stub(goal, iteration, last_result, max_iterations, recovery_attempted)
        d.mode = "ollama-fallback"
        d.parsed_success = False
        d.raw_response = result.error[:_RAW_TRIM] if result.error else None
        d.fallback_reason = fallback_reason
        return d

    # Parse and validate JSON
    fallback_reason: Optional[str] = None
    try:
        text = raw
        if text.startswith("```"):
            text = text.split("```")[1]
            if text.startswith("json"):
                text = text[4:]
        start = text.find("{")
        end = text.rfind("}") + 1
        if start >= 0 and end > start:
            text = text[start:end]

        data = json.loads(text.strip())

        action = data.get("action", "")
        new_goal = data.get("goal", goal)
        reason = data.get("reason", "Ollama decision")

        if action not in ("run_agent", "stop"):
            raise ValueError(f"invalid action: {action!r}")
        if not isinstance(new_goal, str):
            raise ValueError("goal must be string")
        if len(new_goal) > _MAX_GOAL_LEN:
            raise ValueError(f"goal too long ({len(new_goal)} > {_MAX_GOAL_LEN})")

        decision = Decision(
            action=action,
            goal=new_goal,
            reason=str(reason)[:200],
            mode="ollama",
            raw_response=raw_trim,
            parsed_success=True,
        )

        # Post-decision validation 1: enforce stop after success
        last_succeeded = last_result is not None and (
            last_result.get("status") == "done"
            and last_result.get("agent_final_status") == "done"
        )
        if last_succeeded and decision.action == "run_agent":
            decision.action = "stop"
            decision.mode = "ollama-corrected-stop"
            decision.corrected = True
            decision.correction_reason = "Corrected: prevented redundant run after success"
            return decision

        # Post-decision validation 2: catch model repeating a failed goal
        last_failed = last_result is not None and (
            last_result.get("status") == "failed"
            or last_result.get("agent_final_status") == "failed"
        )
        if (
            last_failed
            and decision.action == "run_agent"
            and decision.goal.strip().lower() == goal.strip().lower()
        ):
            decision.goal = _SAFE_FALLBACK_GOAL
            decision.mode = "ollama-corrected"
            decision.corrected = True
            decision.correction_reason = (
                f"Corrected: model repeated failed goal '{goal}' — "
                f"redirected to '{_SAFE_FALLBACK_GOAL}'"
            )
            decision.recovery_applied = True

        return decision

    except json.JSONDecodeError as exc:
        fallback_reason = f"JSON decode error: {exc}"
    except ValueError as exc:
        fallback_reason = f"validation failed: {exc}"
    except KeyError as exc:
        fallback_reason = f"missing key: {exc}"

    d = _decide_stub(goal, iteration, last_result, max_iterations, recovery_attempted)
    d.mode = "ollama-fallback"
    d.parsed_success = False
    d.raw_response = raw_trim
    d.fallback_reason = fallback_reason
    return d