File size: 3,273 Bytes
c63bc31 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 | """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 = []
|