cortex_ai/client.py

"""Minimal Python client for the CORTEX AI API.

    from cortex_ai.client import CortexClient

    client = CortexClient("http://localhost:8000")
    print(client.ask("Combien font 12 * 8 ?").content)
"""

from __future__ import annotations

import json
import urllib.error
import urllib.request
from dataclasses import dataclass, field
from typing import Any


@dataclass
class Reply:
    """One assistant reply, with the reasoning trace and any tool calls."""

    content: str
    reasoning: str = ""
    tool_calls: list[dict[str, Any]] = field(default_factory=list)
    usage: dict[str, int] = field(default_factory=dict)
    raw: dict[str, Any] = field(default_factory=dict)


class CortexError(Exception):
    """Raised when the API returns an error."""


class CortexClient:
    """Talks to a running CORTEX AI server over HTTP."""

    def __init__(self, base_url: str = "http://localhost:8000", api_key: str = "") -> None:
        self.base_url = base_url.rstrip("/")
        self.api_key = api_key
        self._history: list[dict[str, str]] = []

    def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:
        request = urllib.request.Request(
            self.base_url + path,
            data=json.dumps(payload).encode("utf-8"),
            headers={"Content-Type": "application/json"},
            method="POST",
        )
        if self.api_key:
            request.add_header("Authorization", f"Bearer {self.api_key}")
        try:
            with urllib.request.urlopen(request, timeout=600) as response:
                return json.loads(response.read().decode("utf-8"))
        except urllib.error.HTTPError as exc:
            raise CortexError(f"HTTP {exc.code}: {exc.read().decode('utf-8', 'replace')}") from None
        except urllib.error.URLError as exc:
            raise CortexError(f"cannot reach {self.base_url}: {exc.reason}") from None

    def _get(self, path: str) -> dict[str, Any]:
        request = urllib.request.Request(self.base_url + path)
        if self.api_key:
            request.add_header("Authorization", f"Bearer {self.api_key}")
        with urllib.request.urlopen(request, timeout=60) as response:
            return json.loads(response.read().decode("utf-8"))

    def health(self) -> dict[str, Any]:
        return self._get("/health")

    def models(self) -> list[str]:
        return [m["id"] for m in self._get("/v1/models")["data"]]

    def ask(self, message: str, *, keep_history: bool = False) -> Reply:
        """Send one user message and return the reply."""
        messages = [*self._history, {"role": "user", "content": message}]
        body = self._post("/v1/chat/completions", {"messages": messages})
        choice = body["choices"][0]["message"]
        reply = Reply(
            content=choice["content"],
            reasoning=body.get("reasoning_content", ""),
            tool_calls=body.get("tool_calls", []),
            usage=body.get("usage", {}),
            raw=body,
        )
        if keep_history:
            self._history = [
                *messages,
                {"role": "assistant", "content": reply.content},
            ]
        return reply

    def reset(self) -> None:
        """Forget the conversation history."""
        self._history = []