Frankenstein-Labs/Cortex-ai

1"""Minimal Python client for the CORTEX AI API.2 3    from cortex_ai.client import CortexClient4 5    client = CortexClient("http://localhost:8000")6    print(client.ask("Combien font 12 * 8 ?").content)7"""8 9from __future__ import annotations10 11import json12import urllib.error13import urllib.request14from dataclasses import dataclass, field15from typing import Any16 17 18@dataclass19class Reply:20    """One assistant reply, with the reasoning trace and any tool calls."""21 22    content: str23    reasoning: str = ""24    tool_calls: list[dict[str, Any]] = field(default_factory=list)25    usage: dict[str, int] = field(default_factory=dict)26    raw: dict[str, Any] = field(default_factory=dict)27 28 29class CortexError(Exception):30    """Raised when the API returns an error."""31 32 33class CortexClient:34    """Talks to a running CORTEX AI server over HTTP."""35 36    def __init__(self, base_url: str = "http://localhost:8000", api_key: str = "") -> None:37        self.base_url = base_url.rstrip("/")38        self.api_key = api_key39        self._history: list[dict[str, str]] = []40 41    def _post(self, path: str, payload: dict[str, Any]) -> dict[str, Any]:42        request = urllib.request.Request(43            self.base_url + path,44            data=json.dumps(payload).encode("utf-8"),45            headers={"Content-Type": "application/json"},46            method="POST",47        )48        if self.api_key:49            request.add_header("Authorization", f"Bearer {self.api_key}")50        try:51            with urllib.request.urlopen(request, timeout=600) as response:52                return json.loads(response.read().decode("utf-8"))53        except urllib.error.HTTPError as exc:54            raise CortexError(f"HTTP {exc.code}: {exc.read().decode('utf-8', 'replace')}") from None55        except urllib.error.URLError as exc:56            raise CortexError(f"cannot reach {self.base_url}: {exc.reason}") from None57 58    def _get(self, path: str) -> dict[str, Any]:59        request = urllib.request.Request(self.base_url + path)60        if self.api_key:61            request.add_header("Authorization", f"Bearer {self.api_key}")62        with urllib.request.urlopen(request, timeout=60) as response:63            return json.loads(response.read().decode("utf-8"))64 65    def health(self) -> dict[str, Any]:66        return self._get("/health")67 68    def models(self) -> list[str]:69        return [m["id"] for m in self._get("/v1/models")["data"]]70 71    def ask(self, message: str, *, keep_history: bool = False) -> Reply:72        """Send one user message and return the reply."""73        messages = [*self._history, {"role": "user", "content": message}]74        body = self._post("/v1/chat/completions", {"messages": messages})75        choice = body["choices"][0]["message"]76        reply = Reply(77            content=choice["content"],78            reasoning=body.get("reasoning_content", ""),79            tool_calls=body.get("tool_calls", []),80            usage=body.get("usage", {}),81            raw=body,82        )83        if keep_history:84            self._history = [85                *messages,86                {"role": "assistant", "content": reply.content},87            ]88        return reply89 90    def reset(self) -> None:91        """Forget the conversation history."""92        self._history = []93