atlas-iac/services/hermes/scripts/classifier_broker.py

#!/usr/bin/env python3
"""Bound Switchyard classifier context before forwarding it to local Ollama."""

from __future__ import annotations

import copy
import json
import os
import threading
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from typing import Any, Final

import httpx

try:
    from voice_route_preflight import (
        COORDINATOR as VOICE_PREFLIGHT,
        MAX_BODY_BYTES as VOICE_PREFLIGHT_MAX_BODY_BYTES,
        validate_request as validate_voice_preflight_request,
    )
except ModuleNotFoundError:  # Test imports use the repository package path.
    from services.hermes.scripts.voice_route_preflight import (
        COORDINATOR as VOICE_PREFLIGHT,
        MAX_BODY_BYTES as VOICE_PREFLIGHT_MAX_BODY_BYTES,
        validate_request as validate_voice_preflight_request,
    )


HOST: Final = os.environ.get("HERMES_CLASSIFIER_BROKER_HOST", "0.0.0.0")
PORT: Final = int(os.environ.get("HERMES_CLASSIFIER_BROKER_PORT", "9008"))
VOICE_PORT: Final = int(os.environ.get("HERMES_VOICE_PREFLIGHT_PORT", "9009"))
UPSTREAM: Final = os.environ.get(
    "HERMES_CLASSIFIER_BROKER_UPSTREAM",
    "http://ollama.ai.svc.cluster.local:11434",
).rstrip("/")
MAX_BODY_BYTES: Final = int(
    os.environ.get("HERMES_CLASSIFIER_BROKER_MAX_BODY", str(32 << 20))
)
MAX_SYSTEM_CHARS: Final = int(
    os.environ.get("HERMES_CLASSIFIER_MAX_SYSTEM_CHARS", "6500")
)
MAX_CONTEXT_CHARS: Final = int(
    os.environ.get("HERMES_CLASSIFIER_MAX_CONTEXT_CHARS", "7000")
)
READ_TIMEOUT_SECONDS: Final = float(
    os.environ.get("HERMES_CLASSIFIER_BROKER_READ_TIMEOUT", "60")
)
ALLOWED_PATHS: Final = {"/v1/chat/completions", "/v1/models"}
VOICE_PREFLIGHT_PATH: Final = "/voice/route-preflight"


def _bounded_text(value: str, limit: int) -> str:
    """Keep both ends of text because intent and current status often sit apart."""
    if len(value) <= limit:
        return value
    if limit < 80:
        return value[:limit]
    marker = "\n...[classifier context compacted]...\n"
    remaining = limit - len(marker)
    head = remaining // 2
    return f"{value[:head]}{marker}{value[-(remaining - head):]}"


def _compact_content(content: Any, limit: int) -> Any:
    """Remove binary/multimodal payloads and bound text sent to the judge."""
    if isinstance(content, str):
        return _bounded_text(content, limit)
    if not isinstance(content, list):
        return content

    compacted: list[Any] = []
    remaining = limit
    for block in content:
        if not isinstance(block, dict):
            continue
        kind = str(block.get("type") or "")
        if kind in {"text", "input_text", "output_text"}:
            key = "text"
            text = str(block.get(key) or "")
            if not text or remaining <= 0:
                continue
            text = _bounded_text(text, remaining)
            compacted.append({**block, key: text})
            remaining -= len(text)
        elif kind in {"image", "image_url", "input_image"}:
            marker = "[image attachment available to the selected worker]"
            if remaining >= len(marker):
                compacted.append({"type": "text", "text": marker})
                remaining -= len(marker)
    return compacted


def _compact_message(message: dict[str, Any], limit: int) -> dict[str, Any]:
    """Copy routing-relevant message metadata while bounding large values."""
    result = copy.deepcopy(message)
    if "content" in result:
        result["content"] = _compact_content(result["content"], limit)
    # The local judge only classifies the current boundary. OpenAI-compatible
    # Ollama validates assistant tool-call JSON and tool/result pairing before
    # inference, so a bounded or selectively retained transcript can become an
    # invalid conversation even though its text is sufficient for routing.
    # Preserve tool evidence as plain user text and remove protocol metadata.
    had_tool_calls = bool(result.pop("tool_calls", None))
    result.pop("tool_call_id", None)
    result.pop("name", None)
    if result.get("role") == "tool":
        result["role"] = "user"
        content = result.get("content")
        result["content"] = f"[tool evidence]\n{content or ''}"
    elif had_tool_calls and not result.get("content"):
        result["content"] = "[assistant requested an external tool]"
    return result


def _compact_messages(messages: list[Any]) -> list[Any]:
    """Keep the routing contract, opening task, and latest decision context."""
    valid = [message for message in messages if isinstance(message, dict)]
    system_indices = [
        index
        for index, message in enumerate(valid)
        if message.get("role") in {"system", "developer"}
    ]
    non_system = [index for index in range(len(valid)) if index not in system_indices]
    user_indices = [index for index in non_system if valid[index].get("role") == "user"]

    opening_user = user_indices[0] if user_indices else None
    latest_user = user_indices[-1] if user_indices else None
    selected = set(system_indices)
    selected.update(non_system[-4:])
    if opening_user is not None:
        selected.add(opening_user)
    if latest_user is not None:
        selected.add(latest_user)

    other_indices = [
        index
        for index in selected
        if index not in system_indices and index not in {opening_user, latest_user}
    ]
    latest_budget = min(3500, MAX_CONTEXT_CHARS)
    opening_budget = min(1200, max(0, MAX_CONTEXT_CHARS - latest_budget))
    other_budget = max(0, MAX_CONTEXT_CHARS - latest_budget - opening_budget)
    other_limit = min(1200, other_budget // max(1, len(other_indices)))
    system_limit = MAX_SYSTEM_CHARS // max(1, len(system_indices))

    result: list[Any] = []
    for index, message in enumerate(valid):
        if index not in selected:
            continue
        role = str(message.get("role") or "")
        if role in {"system", "developer"}:
            compacted = _compact_message(message, system_limit)
        elif index == latest_user:
            compacted = _compact_message(message, latest_budget)
        elif index == opening_user:
            compacted = _compact_message(message, opening_budget)
        else:
            if other_limit <= 0:
                continue
            compacted = _compact_message(message, other_limit)
        result.append(compacted)
    return result


def compact_payload(payload: dict[str, Any]) -> dict[str, Any]:
    """Return the same OpenAI request with only classifier input compacted."""
    result = copy.deepcopy(payload)
    messages = result.get("messages")
    if isinstance(messages, list):
        result["messages"] = _compact_messages(messages)
    # The judge never needs tools or binary inputs. Switchyard supplies a
    # response schema separately, and that contract must remain untouched.
    result.pop("tools", None)
    result.pop("tool_choice", None)
    result.pop("parallel_tool_calls", None)
    return result


class Handler(BaseHTTPRequestHandler):
    """Proxy only Switchyard's local judge calls with strict input bounds."""

    server_version = "HermesClassifierBroker/1"

    def log_message(self, format: str, *args: object) -> None:
        return

    def _json(self, status: int, value: dict[str, object]) -> None:
        body = json.dumps(value, separators=(",", ":")).encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(body)

    def _empty(self, status: int) -> None:
        self.send_response(status)
        self.send_header("Content-Length", "0")
        self.send_header("Cache-Control", "no-store")
        self.end_headers()

    def _text(self, status: int, body: str, content_type: str) -> None:
        encoded = body.encode("utf-8")
        self.send_response(status)
        self.send_header("Content-Type", content_type)
        self.send_header("Content-Length", str(len(encoded)))
        self.send_header("Cache-Control", "no-store")
        self.end_headers()
        self.wfile.write(encoded)

    def do_GET(self) -> None:  # noqa: N802
        if self.path == "/health":
            self._json(200, {"ok": True, "upstream": "ollama"})
            return
        if self.path == "/metrics":
            self._text(200, VOICE_PREFLIGHT.metrics(), "text/plain; version=0.0.4")
            return
        if self.path not in ALLOWED_PATHS:
            self._json(404, {"error": "not found"})
            return
        self._proxy(None)

    def do_POST(self) -> None:  # noqa: N802
        if self.path not in ALLOWED_PATHS:
            self._json(404, {"error": "not found"})
            return
        authoritative = self.path == "/v1/chat/completions"
        if authoritative:
            # Mark priority as soon as the authoritative path arrives, before
            # even validating or reading its local body, so an
            # in-flight disposable decode is closed at authoritative arrival.
            VOICE_PREFLIGHT.begin_authoritative()
        try:
            try:
                length = int(self.headers.get("Content-Length", "0"))
            except ValueError:
                length = -1
            if length <= 0 or length > MAX_BODY_BYTES:
                self._json(413, {"error": "request too large"})
                return
            try:
                payload = json.loads(self.rfile.read(length))
                if not isinstance(payload, dict):
                    raise ValueError("request must be a JSON object")
                body = json.dumps(
                    compact_payload(payload), separators=(",", ":")
                ).encode("utf-8")
            except (ValueError, TypeError, json.JSONDecodeError) as exc:
                self._json(400, {"error": str(exc)})
                return
            print(
                f"classifier-broker request_bytes={length} compacted_bytes={len(body)}",
                flush=True,
            )
            if authoritative:
                self._proxy(body)
                return
            self._proxy(body)
        finally:
            if authoritative:
                VOICE_PREFLIGHT.end_authoritative()

    def _voice_preflight(self) -> None:
        """Return one disposable local tier without starting a hosted turn."""
        try:
            length = int(self.headers.get("Content-Length", "0"))
        except ValueError:
            length = -1
        if length <= 0 or length > VOICE_PREFLIGHT_MAX_BODY_BYTES:
            self._json(413, {"error": "request too large"})
            return
        try:
            payload = json.loads(self.rfile.read(length))
            if not isinstance(payload, dict):
                raise ValueError("request must be a JSON object")
            turn_id, revision, transcript = validate_voice_preflight_request(payload)
        except (ValueError, TypeError, json.JSONDecodeError) as exc:
            self._json(400, {"error": str(exc)})
            return

        cancel = VOICE_PREFLIGHT.begin(turn_id)
        if cancel is None:
            self._empty(204)
            return
        try:
            tier = VOICE_PREFLIGHT.stream(transcript, cancel)
        finally:
            VOICE_PREFLIGHT.end()
        if not tier:
            self._empty(204)
            return
        self._json(
            200,
            {
                "turn_id": turn_id,
                "revision": revision,
                "tier": tier,
                "target": f"atlas/auto/{tier}",
                "advisory": True,
            },
        )

    def _proxy(self, body: bytes | None) -> None:
        try:
            timeout = httpx.Timeout(10.0, read=READ_TIMEOUT_SECONDS)
            with httpx.Client(timeout=timeout) as client:
                response = client.request(
                    self.command,
                    f"{UPSTREAM}{self.path}",
                    headers={"Content-Type": "application/json"},
                    content=body,
                )
            self.send_response(response.status_code)
            self.send_header(
                "Content-Type", response.headers.get("Content-Type", "application/json")
            )
            self.send_header("Content-Length", str(len(response.content)))
            self.send_header("Cache-Control", "no-store")
            self.end_headers()
            self.wfile.write(response.content)
        except (httpx.HTTPError, OSError) as exc:
            self._json(503, {"error": f"classifier unavailable: {exc}"})


class VoiceHandler(Handler):
    """Expose only advisory, health and telemetry routes to chat tenants."""

    server_version = "HermesVoicePreflight/1"

    def do_GET(self) -> None:  # noqa: N802
        if self.path == "/health":
            self._json(200, {"ok": True, "upstream": "ollama"})
            return
        if self.path == "/metrics":
            self._text(200, VOICE_PREFLIGHT.metrics(), "text/plain; version=0.0.4")
            return
        self._json(404, {"error": "not found"})

    def do_POST(self) -> None:  # noqa: N802
        if self.path == VOICE_PREFLIGHT_PATH:
            self._voice_preflight()
            return
        self._json(404, {"error": "not found"})


def main() -> None:
    voice_server = ThreadingHTTPServer((HOST, VOICE_PORT), VoiceHandler)
    voice_thread = threading.Thread(target=voice_server.serve_forever, daemon=True)
    voice_thread.start()
    server = ThreadingHTTPServer((HOST, PORT), Handler)
    try:
        server.serve_forever()
    finally:
        voice_server.shutdown()
        voice_server.server_close()


if __name__ == "__main__":
    main()