ai_router.py

from __future__ import annotations

from typing import Any

import httpx

from .config import OLLARMA_BASE_URL, OLLARMA_TIMEOUT_SECONDS


GENERIC_HINTS = (
    "i don't have access",
    "i do not have access",
    "i can't access",
    "i cannot access",
    "[blocked by guardrail",
)


def _extract_response_text(payload: dict[str, Any]) -> str:
    for key in ("final_response", "response", "message", "content"):
        value = payload.get(key)
        if isinstance(value, str) and value.strip():
            return value.strip()
    return ""


def _generic_response(text: str) -> bool:
    lowered = text.lower()
    return any(hint in lowered for hint in GENERIC_HINTS)


def _extract_reason_code(payload: dict[str, Any], status_code: int) -> str:
    reason_code = payload.get("reason_code")
    if isinstance(reason_code, str) and reason_code.strip():
        return reason_code.strip()

    error = payload.get("error")
    if isinstance(error, str) and ":" in error:
        candidate = error.split(":", 1)[0].strip()
        if candidate and candidate.upper() == candidate:
            return candidate

    return f"OLLARMA_HTTP_{status_code}"


def query_local_ai(message: str, *, project: str | None = None) -> dict[str, Any]:
    endpoint = "/route" if project else "/chat"
    payload: dict[str, Any] = {"project": project, "prompt": message} if project else {"message": message}
    attempted_lane = "ollarma.route" if project else "ollarma.chat"
    try:
        with httpx.Client(base_url=OLLARMA_BASE_URL, timeout=OLLARMA_TIMEOUT_SECONDS) as client:
            response = client.post(endpoint, json=payload)
        body = response.json()
    except httpx.TimeoutException:
        return {
            "attempted_lane": attempted_lane,
            "lane_used": None,
            "fallback_recommended": True,
            "reason_code": "OLLARMA_TIMEOUT",
            "response_text": "",
            "raw": None,
        }
    except httpx.HTTPError as exc:
        return {
            "attempted_lane": attempted_lane,
            "lane_used": None,
            "fallback_recommended": True,
            "reason_code": "OLLARMA_UNAVAILABLE",
            "response_text": "",
            "raw": {"error": str(exc)},
        }

    text = _extract_response_text(body)
    if response.status_code >= 400:
        return {
            "attempted_lane": attempted_lane,
            "lane_used": body.get("lane") or attempted_lane,
            "fallback_recommended": True,
            "reason_code": _extract_reason_code(body, response.status_code),
            "response_text": text,
            "raw": body,
        }
    route_reason = body.get("reason_code") if project else None
    route_blocked = bool(route_reason) or body.get("kb_status") == "blocked" or body.get("lane") == "frontier_or_human"
    if route_blocked:
        return {
            "attempted_lane": attempted_lane,
            "lane_used": body.get("lane") or attempted_lane,
            "fallback_recommended": True,
            "reason_code": route_reason or "OLLARMA_ROUTE_BLOCKED",
            "response_text": text,
            "raw": body,
        }
    if _generic_response(text):
        return {
            "attempted_lane": attempted_lane,
            "lane_used": attempted_lane,
            "fallback_recommended": True,
            "reason_code": "LOW_CONFIDENCE_GENERIC",
            "response_text": text,
            "raw": body,
        }
    return {
        "attempted_lane": attempted_lane,
        "lane_used": body.get("lane") or attempted_lane,
        "fallback_recommended": False,
        "reason_code": None,
        "response_text": text,
        "raw": body,
    }


def submit_deterministic_workflow(project: str, manifest_ref: str, step_id: str) -> dict[str, Any]:
    try:
        with httpx.Client(base_url=OLLARMA_BASE_URL, timeout=OLLARMA_TIMEOUT_SECONDS) as client:
            response = client.post(
                "/workflow",
                json={
                    "project": project,
                    "manifest_ref": manifest_ref,
                    "step_id": step_id,
                },
            )
        body = response.json()
    except httpx.TimeoutException:
        return {
            "attempted_lane": "ollarma.workflow",
            "lane_used": None,
            "accepted": False,
            "status_code": 504,
            "reason_code": "OLLARMA_TIMEOUT",
            "raw": None,
        }
    except httpx.HTTPError as exc:
        return {
            "attempted_lane": "ollarma.workflow",
            "lane_used": None,
            "accepted": False,
            "status_code": 503,
            "reason_code": "OLLARMA_UNAVAILABLE",
            "raw": {"error": str(exc)},
        }
    return {
        "attempted_lane": "ollarma.workflow",
        "lane_used": "ollarma.workflow",
        "accepted": response.status_code in {200, 202},
        "status_code": response.status_code,
        "reason_code": body.get("reason_code"),
        "raw": body,
    }