myagents/myagents/ollama_adapter.py

"""Anthropic -> Ollama normalization adapter for Claude Code.

Claude Code injects some system content (e.g. the agent-types list) as
``role: "system"`` messages inside ``messages[]``, mid-conversation. Qwen-family
chat templates require a leading system message, so Ollama's ``/v1/messages``
responds HTTP 500 ("System message must be at the beginning"). This adapter
moves any in-messages system content into the top-level Anthropic ``system``
field before forwarding to Ollama.

Runs standalone:

    python -m myagents.ollama_adapter          # 127.0.0.1:8199

Stdlib only (no fastapi/uvicorn) to keep myagents a lightweight CLI package.
"""

from __future__ import annotations

import contextlib
import json
import logging
import os
from http.client import HTTPConnection
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from urllib.parse import urlsplit

DEFAULT_HOST = os.environ.get("XIAOHE_OLLAMA_ADAPTER_HOST", "127.0.0.1")
DEFAULT_PORT = int(os.environ.get("XIAOHE_OLLAMA_ADAPTER_PORT", "8199"))
# Respect OLLAMA_HOST just like the ollama CLI.
OLLAMA_BASE = os.environ.get("OLLAMA_HOST", "http://127.0.0.1:11434")

_log = logging.getLogger("myagents.ollama_adapter")


def normalize_system(body: dict) -> dict:
    """Move ``role: "system"`` messages into the top-level ``system`` field.

    Mutates and returns *body*. Non-system messages keep their order.
    """
    raw_sys = body.get("system")
    if isinstance(raw_sys, list):
        sys_blocks = [b for b in raw_sys if isinstance(b, dict)]
    elif isinstance(raw_sys, str) and raw_sys:
        sys_blocks = [{"type": "text", "text": raw_sys}]
    else:
        sys_blocks = []

    kept: list = []
    for message in body.get("messages") or []:
        if isinstance(message, dict) and message.get("role") == "system":
            content = message.get("content")
            if isinstance(content, list):
                for block in content:
                    if (
                        isinstance(block, dict)
                        and block.get("type") == "text"
                        and block.get("text")
                    ):
                        sys_blocks.append(
                            {"type": "text", "text": block["text"]}
                        )
        else:
            kept.append(message)

    body["system"] = sys_blocks
    body["messages"] = kept
    return body


class _Handler(BaseHTTPRequestHandler):
    protocol_version = "HTTP/1.1"

    # -- plumbing -----------------------------------------------------------

    def _json(self, status: int, data: dict) -> None:
        body = json.dumps(data).encode()
        self.send_response(status)
        self.send_header("content-type", "application/json")
        self.send_header("content-length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def _forward(self, raw: bytes, path: str) -> None:
        """POST to Ollama and stream the response through unchanged."""
        parts = urlsplit(OLLAMA_BASE)
        headers = {
            "x-api-key": "ollama",
            "content-type": "application/json",
            "anthropic-version": self.headers.get(
                "anthropic-version", "2023-06-01"
            ),
            "content-length": str(len(raw)),
        }
        conn = HTTPConnection(
            parts.hostname or "127.0.0.1", parts.port or 80, timeout=600
        )
        try:
            conn.request("POST", path, body=raw, headers=headers)
            resp = conn.getresponse()
        except OSError as exc:
            _log.warning("upstream %s failed: %s", OLLAMA_BASE, exc)
            self._json(502, {"error": {"message": f"upstream error: {exc}"}})
            return
        self.send_response(resp.status)
        content_type = resp.getheader("content-type")
        if content_type:
            self.send_header("content-type", content_type)
        self.send_header("cache-control", "no-cache")
        self.end_headers()
        while True:
            chunk = resp.read(65536)
            if not chunk:
                break
            try:
                self.wfile.write(chunk)
                self.wfile.flush()
            except (BrokenPipeError, ConnectionResetError):
                break
        conn.close()

    # -- routes -------------------------------------------------------------

    def do_GET(self) -> None:
        if self.path == "/health":
            self._json(200, {"status": "ok", "upstream": OLLAMA_BASE})
            return
        if self.path == "/v1/models":
            # Anthropic-style model list so /model switching works offline.
            try:
                from urllib.request import urlopen

                with urlopen(OLLAMA_BASE + "/api/tags", timeout=5) as resp:
                    data = json.loads(resp.read().decode())
                names = [m["name"] for m in data.get("models", [])]
            except OSError:
                names = []
            self._json(
                200,
                {"data": [{"id": n, "object": "model"} for n in names]},
            )
            return
        self._json(404, {"error": {"message": "not found"}})

    def do_POST(self) -> None:
        if self.path.split("?", 1)[0] != "/v1/messages":
            self._json(404, {"error": {"message": "not found"}})
            return
        length = int(self.headers.get("content-length") or 0)
        raw = self.rfile.read(length) if length else b""
        try:
            body = json.loads(raw)
        except json.JSONDecodeError:
            body = None
        if isinstance(body, dict):
            normalize_system(body)
            raw = json.dumps(body).encode()
        self._forward(raw, self.path)

    def log_message(self, format: str, *args: object) -> None:  # noqa: A002
        if args:
            format = format % args
        _log.info("%s %s", self.command, format)


def main() -> None:
    logging.basicConfig(level=logging.INFO)
    server = ThreadingHTTPServer((DEFAULT_HOST, DEFAULT_PORT), _Handler)
    server.daemon_threads = True
    print(
        f"ollama-adapter: http://{DEFAULT_HOST}:{DEFAULT_PORT} -> {OLLAMA_BASE}",
        flush=True,
    )
    with contextlib.suppress(KeyboardInterrupt):
        server.serve_forever()


if __name__ == "__main__":
    main()