saia-python 0.4.1__py3-none-any.whl

saia_python/chat.py

@@ -0,0 +1,72 @@
+"""Chat service — completions and streaming."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ._http import post_chat_completion
+from ._streaming import SSEStream
+
+if TYPE_CHECKING:
+    import requests
+
+
+class ChatService:
+    """Access the ``/chat/completions`` endpoint.
+
+    Args:
+        session: A :class:`requests.Session` with auth headers configured.
+        base_url: The SAIA API base URL.
+    """
+
+    def __init__(self, session: requests.Session, base_url: str):
+        self._session = session
+        self._base_url = base_url
+
+    def completions(
+        self,
+        model: str,
+        messages: list[dict],
+        *,
+        temperature: float | None = None,
+        top_p: float | None = None,
+        max_tokens: int | None = None,
+        stream: bool = False,
+        **kwargs,
+    ) -> dict | SSEStream:
+        """Send a chat completion request.
+
+        Args:
+            model: Model identifier (e.g. ``"meta-llama-3.1-8b-instruct"``).
+            messages: List of message dicts with ``"role"`` and ``"content"`` keys.
+            temperature: Sampling temperature (0–2).
+            top_p: Nucleus sampling parameter (0–1).
+            max_tokens: Maximum tokens to generate.
+            stream: If ``True``, return a generator yielding chunks.
+            **kwargs: Additional parameters forwarded to the API.
+
+        Returns:
+            When ``stream=False``: the API response dict, with an extra
+            ``"_rate_limits"`` key — a JSON-serializable dict of the current
+            rate-limit headers (see :class:`~saia_python.RateLimitInfo`).
+            When ``stream=True``: an ``SSEStream`` — iterate it for the
+            response chunks; its ``rate_limits`` attribute exposes the same
+            dict (available immediately, from the response headers).
+        """
+        body = {"model": model, "messages": messages, **kwargs}
+        if temperature is not None:
+            body["temperature"] = temperature
+        if top_p is not None:
+            body["top_p"] = top_p
+        if max_tokens is not None:
+            body["max_tokens"] = max_tokens
+
+        return post_chat_completion(
+            self._session,
+            f"{self._base_url}/chat/completions",
+            body,
+            stream=stream,
+        )
+
+    def __repr__(self):
+        return f"ChatService(base_url={self._base_url!r})"

saia_python/client.py

@@ -0,0 +1,239 @@
+"""Main SAIA client — composes all services."""
+
+from __future__ import annotations
+
+import requests as _requests
+
+from .arcana import ArcanaService
+from .auth import resolve_credentials
+from .chat import ChatService
+from .documents import DocumentService
+from .exceptions import raise_for_status
+from .models import ModelsService
+from .rate_limits import RateLimitInfo, parse_rate_limits
+from .voice import VoiceService
+
+
+class SAIAClient:
+    """High-level client for the GWDG SAIA platform.
+
+    Provides access to Chat, Voice AI, ARCANA, Documents, and model
+    listing through a shared, authenticated HTTP session. An OpenAI-
+    compatible client is available via the ``.openai`` property.
+
+    Args:
+        api_key: Your SAIA API key. If omitted, the key is resolved
+            automatically — see :func:`~saia_python.load_api_key` for the
+            resolution order.
+        base_url: Base URL for the SAIA API. Resolution order:
+            explicit parameter > ``[saia] base_url`` in ``config.toml`` >
+            hardcoded default (``https://chat-ai.academiccloud.de/v1``).
+        key_file: Explicit path to a ``.saia_api`` or ``.env`` file.
+            Ignored when ``api_key`` is provided.
+
+    Example::
+
+        # All settings resolved automatically
+        client = SAIAClient()
+
+        # Native services
+        client.chat.completions(model="...", messages=[...])
+
+        # OpenAI-compatible client (requires pip install saia-python[openai])
+        client.openai.chat.completions.create(model="...", messages=[...])
+    """
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str | None = None,
+        key_file: str | None = None,
+    ):
+        self._api_key, self._base_url = resolve_credentials(api_key, base_url, key_file)
+        self._session = _requests.Session()
+        self._session.headers.update(
+            {
+                "Authorization": f"Bearer {self._api_key}",
+                "Accept": "application/json",
+            }
+        )
+
+        self._chat: ChatService | None = None
+        self._voice: VoiceService | None = None
+        self._models: ModelsService | None = None
+        self._arcana: ArcanaService | None = None
+        self._documents: DocumentService | None = None
+        self._openai = None
+        self._openai_async = None
+
+    @property
+    def chat(self) -> ChatService:
+        """Chat completions service."""
+        if self._chat is None:
+            self._chat = ChatService(self._session, self._base_url)
+        return self._chat
+
+    @property
+    def voice(self) -> VoiceService:
+        """Voice AI (transcription/translation) service."""
+        if self._voice is None:
+            self._voice = VoiceService(self._session, self._base_url)
+        return self._voice
+
+    @property
+    def models(self) -> ModelsService:
+        """Model listing service."""
+        if self._models is None:
+            self._models = ModelsService(self._session, self._base_url)
+        return self._models
+
+    @property
+    def arcana(self) -> ArcanaService:
+        """ARCANA/RAG service."""
+        if self._arcana is None:
+            self._arcana = ArcanaService(self._session, self._base_url, self._api_key)
+        return self._arcana
+
+    @property
+    def documents(self) -> DocumentService:
+        """Document conversion (Docling) service."""
+        if self._documents is None:
+            self._documents = DocumentService(self._session, self._base_url)
+        return self._documents
+
+    @property
+    def openai(self):
+        """OpenAI-compatible synchronous client.
+
+        Returns an ``openai.OpenAI`` instance configured with the same
+        API key and base URL as this client. Requires
+        ``pip install saia-python[openai]``.
+
+        Example::
+
+            response = client.openai.chat.completions.create(
+                model="llama-3.3-70b-instruct",
+                messages=[{"role": "user", "content": "Hello!"}],
+            )
+        """
+        if self._openai is None:
+            from .openai_compat import create_openai_client
+
+            self._openai = create_openai_client(
+                api_key=self._api_key,
+                base_url=self._base_url,
+            )
+        return self._openai
+
+    @property
+    def openai_async(self):
+        """OpenAI-compatible asynchronous client.
+
+        Returns an ``openai.AsyncOpenAI`` instance. Requires
+        ``pip install saia-python[openai]``.
+        """
+        if self._openai_async is None:
+            from .openai_compat import create_openai_client
+
+            self._openai_async = create_openai_client(
+                api_key=self._api_key,
+                base_url=self._base_url,
+                async_client=True,
+            )
+        return self._openai_async
+
+    def get_rate_limits(self) -> RateLimitInfo:
+        """Fetch current rate-limit status by making a lightweight API call.
+
+        Uses a GET to ``/chat/completions`` which returns 400 but includes
+        rate-limit headers.
+
+        Returns:
+            Parsed :class:`RateLimitInfo`.
+
+        Raises:
+            AuthenticationError: If the API key is invalid or expired
+                (401/403). Other non-2xx statuses (notably the expected
+                400) are tolerated since they still carry the headers.
+        """
+        resp = self._session.get(f"{self._base_url}/chat/completions")
+        if resp.status_code in (401, 403):
+            # The probe is *expected* to 400 (missing request body) but still
+            # carries rate-limit headers. A 401/403 means the key is bad, so
+            # surface it instead of silently returning an empty RateLimitInfo.
+            raise_for_status(resp)
+        return parse_rate_limits(resp.headers)
+
+    def arcana_version(self) -> str:
+        """Return the ARCANA API version string.
+
+        Thin delegate to :meth:`ArcanaService.version`
+        (``client.arcana.version()``), which owns the ARCANA URL path and
+        auth scheme.
+
+        Returns:
+            The version string (e.g. ``"0.4.16"``).
+        """
+        return self.arcana.version()
+
+    def arcana_heartbeat(self) -> bool:
+        """Check if the ARCANA service is alive.
+
+        Thin delegate to :meth:`ArcanaService.heartbeat`
+        (``client.arcana.heartbeat()``). Returns ``True`` if the service
+        responds with 204, ``False`` otherwise.
+
+        Returns:
+            ``True`` if the service is reachable.
+        """
+        return self.arcana.heartbeat()
+
+    def health_check(self, *, verbose: bool = False) -> bool | dict:
+        """Verify that the client can reach the API and authenticate.
+
+        Combines two cheap GETs:
+
+        - ``GET /models`` (authenticated) — confirms the API key resolves
+          and the chat backend is reachable.
+        - ``GET /arcanas/api/v1/heartbeat`` (cheap 204) — confirms the
+          ARCANA backend is reachable.
+
+        Args:
+            verbose: If ``True``, return a diagnostic dict instead of a
+                bool. Useful in onboarding / setup scripts where you
+                want to surface *which* leg failed.
+
+        Returns:
+            ``True`` if both legs succeed, ``False`` otherwise. With
+            ``verbose=True``, a dict::
+
+                {
+                    "ok":            <bool>,
+                    "base_url":      <str>,
+                    "models_ok":     <bool>,
+                    "model_count":   <int>,    # 0 if models leg failed
+                    "arcana_ok":     <bool>,
+                    "error":         <str|None>,  # first leg that failed
+                }
+        """
+        details: dict = {
+            "base_url": self._base_url,
+            "models_ok": False,
+            "model_count": 0,
+            "arcana_ok": False,
+            "error": None,
+        }
+        try:
+            model_ids = self.models.list_ids()
+            details["models_ok"] = True
+            details["model_count"] = len(model_ids)
+        except Exception as exc:
+            details["error"] = f"models: {exc}"
+        details["arcana_ok"] = self.arcana_heartbeat()
+        if not details["arcana_ok"] and details["error"] is None:
+            details["error"] = "arcana heartbeat returned non-204"
+        details["ok"] = details["models_ok"] and details["arcana_ok"]
+        return details if verbose else bool(details["ok"])
+
+    def __repr__(self):
+        return f"SAIAClient(base_url={self._base_url!r})"

saia_python/documents.py

@@ -0,0 +1,145 @@
+"""Document conversion service (Docling) — convert PDFs and documents."""
+
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+from .exceptions import raise_for_status
+
+if TYPE_CHECKING:
+    import requests
+
+
+@dataclass
+class ConversionResult:
+    """Result of a document conversion via the Docling service.
+
+    Attributes:
+        filename: The original document filename.
+        response_type: The output format (``markdown``, ``html``, ``json``, or ``tokens``).
+        content: The converted content as a string.
+        images: A list of extracted image dicts, each with ``type``,
+            ``filename``, and ``data`` (base64-encoded) keys.
+    """
+
+    filename: str
+    response_type: str
+    content: str
+    images: list[dict] = field(default_factory=list)
+
+    def save(self, path: str | Path) -> Path:
+        """Save the converted content to a file.
+
+        Args:
+            path: Output file path.
+
+        Returns:
+            The path written to.
+        """
+        path = Path(path)
+        path.write_text(self.content, encoding="utf-8")
+        return path
+
+    def __str__(self):
+        n_img = len(self.images)
+        preview = self.content[:200]
+        suffix = "..." if len(self.content) > 200 else ""
+        return (
+            f"ConversionResult({self.filename!r}, {self.response_type}, "
+            f"{n_img} images, {len(self.content)} chars)\n{preview}{suffix}"
+        )
+
+
+class DocumentService:
+    """Access the ``/documents/convert`` endpoint (Docling).
+
+    Converts PDF and other document formats to markdown, HTML, JSON,
+    or token representations.
+
+    Args:
+        session: A :class:`requests.Session` with auth headers configured.
+        base_url: The SAIA API base URL.
+    """
+
+    def __init__(self, session: requests.Session, base_url: str):
+        self._session = session
+        self._base_url = base_url
+
+    def convert(
+        self,
+        file_path: str | Path,
+        *,
+        response_type: str = "markdown",
+        extract_tables_as_images: bool | None = None,
+        image_resolution_scale: int | None = None,
+    ) -> ConversionResult:
+        """Convert a document using the Docling service.
+
+        Args:
+            file_path: Path to the document file (PDF, etc.).
+            response_type: Output format — ``"markdown"`` (default),
+                ``"html"``, ``"json"``, or ``"tokens"``.
+            extract_tables_as_images: If ``True``, render tables as images
+                instead of structured text.
+            image_resolution_scale: Image quality multiplier (1–4).
+
+        Returns:
+            A :class:`ConversionResult` with the converted content and
+            any extracted images.
+        """
+        file_path = Path(file_path)
+        params: dict = {"response_type": response_type}
+        if extract_tables_as_images is not None:
+            params["extract_tables_as_images"] = str(extract_tables_as_images).lower()
+        if image_resolution_scale is not None:
+            params["image_resolution_scale"] = image_resolution_scale
+
+        with open(file_path, "rb") as f:
+            resp = self._session.post(
+                f"{self._base_url}/documents/convert",
+                params=params,
+                files={"document": (file_path.name, f)},
+            )
+        raise_for_status(resp)
+        data = resp.json()
+
+        content = data.get(response_type, data.get("markdown", ""))
+        if isinstance(content, (list, dict)):
+            content = json.dumps(content, indent=2)
+
+        return ConversionResult(
+            filename=data.get("filename", file_path.name),
+            response_type=data.get("response_type", response_type),
+            content=content,
+            images=data.get("images", []),
+        )
+
+    def convert_to_markdown(self, file_path: str | Path, **kwargs) -> str:
+        """Convert a document to markdown (convenience method).
+
+        Args:
+            file_path: Path to the document file.
+            **kwargs: Additional parameters passed to :meth:`convert`.
+
+        Returns:
+            The markdown content as a string.
+        """
+        return self.convert(file_path, response_type="markdown", **kwargs).content
+
+    def convert_to_html(self, file_path: str | Path, **kwargs) -> str:
+        """Convert a document to HTML (convenience method).
+
+        Args:
+            file_path: Path to the document file.
+            **kwargs: Additional parameters passed to :meth:`convert`.
+
+        Returns:
+            The HTML content as a string.
+        """
+        return self.convert(file_path, response_type="html", **kwargs).content
+
+    def __repr__(self):
+        return f"DocumentService(base_url={self._base_url!r})"

saia_python/exceptions.py

@@ -0,0 +1,68 @@
+"""Custom exceptions and shared HTTP error handling for the SAIA Python wrapper."""
+
+from __future__ import annotations
+
+import json as _json
+from typing import TYPE_CHECKING
+
+from .rate_limits import parse_rate_limits
+
+if TYPE_CHECKING:
+    import requests
+
+
+class SAIAError(Exception):
+    """Base exception for all SAIA API errors."""
+
+
+class AuthenticationError(SAIAError):
+    """Raised on 401/403 responses — invalid or missing API key."""
+
+
+class RateLimitError(SAIAError):
+    """Raised on 429 responses — rate limit exceeded."""
+
+    def __init__(self, message, rate_limits=None):
+        super().__init__(message)
+        self.rate_limits = rate_limits
+
+
+class APIError(SAIAError):
+    """Raised on unexpected HTTP errors."""
+
+    def __init__(self, message, status_code=None, response_body=None):
+        super().__init__(message)
+        self.status_code = status_code
+        self.response_body = response_body
+
+
+def _extract_detail(resp: requests.Response) -> str:
+    """Try to extract a human-readable message from a JSON error body.
+
+    The SAIA API typically returns ``{"detail": "..."}`` on errors.
+    Falls back to the raw response text.
+    """
+    try:
+        body = resp.json()
+        if isinstance(body, dict) and "detail" in body:
+            return body["detail"]
+    except (_json.JSONDecodeError, ValueError):
+        pass
+    return resp.text
+
+
+def raise_for_status(resp: requests.Response) -> None:
+    """Raise a typed SAIA exception for HTTP error responses.
+
+    Called by all service modules — this is the single implementation.
+    """
+    if resp.ok:
+        return
+    detail = _extract_detail(resp)
+    if resp.status_code in (401, 403):
+        raise AuthenticationError(detail)
+    if resp.status_code == 429:
+        raise RateLimitError(detail, rate_limits=parse_rate_limits(resp.headers))
+    # Any other non-2xx/3xx status. The early `return` above already handled
+    # the resp.ok case, so reaching here always means an error response.
+    raise APIError(detail, status_code=resp.status_code, response_body=resp.text)

saia_python/models.py

@@ -0,0 +1,146 @@
+"""Models service — list available SAIA models."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from ._util import progress_iter
+from .exceptions import raise_for_status
+
+if TYPE_CHECKING:
+    import requests
+
+
+class ModelsService:
+    """Access the ``/models`` endpoint.
+
+    Args:
+        session: A :class:`requests.Session` with auth headers configured.
+        base_url: The SAIA API base URL (e.g. ``https://chat-ai.academiccloud.de/v1``).
+    """
+
+    def __init__(self, session: requests.Session, base_url: str):
+        self._session = session
+        self._base_url = base_url
+
+    def list_raw(self) -> dict:
+        """Return the raw ``/models`` response envelope, as the API sent it.
+
+        Unlike :meth:`list`, this does **not** unwrap the OpenAI-style
+        ``{"object": "list", "data": [...]}`` envelope — it returns the
+        parsed JSON verbatim. Use it when you need the full
+        OpenAI-compatible payload, e.g. an adapter that re-serves SAIA's
+        models at its own ``GET /v1/models`` endpoint::
+
+            return client.models.list_raw()   # already the OpenAI envelope
+
+        Returns:
+            The parsed JSON response. For the SAIA / OpenAI-compatible API
+            this is a dict of the form
+            ``{"object": "list", "data": [...]}``.
+        """
+        resp = self._session.get(f"{self._base_url}/models")
+        raise_for_status(resp)
+        return resp.json()
+
+    def list(self) -> list[dict]:
+        """Return the full model list as returned by the API.
+
+        Returns:
+            A list of model dicts, each containing at least an ``"id"`` key.
+        """
+        data = self.list_raw()
+        if isinstance(data, list):
+            return data
+        if isinstance(data, dict) and "data" in data:
+            return data["data"]
+        return [data]
+
+    def list_ids(self) -> list[str]:
+        """Return a deduplicated list of model ID strings."""
+        models = self.list()
+        seen = set()
+        ids = []
+        for m in models:
+            mid = (
+                m.get("id")
+                or m.get("modelId")
+                or m.get("model_id")
+                or m.get("model")
+                or m.get("name")
+                or m.get("model_name")
+            )
+            if mid and mid not in seen:
+                seen.add(mid)
+                ids.append(mid)
+        return ids
+
+    def list_tool_capable(self, *, verbose: bool = False) -> list[str]:
+        """Identify models that support tool calling by probing each one.
+
+        Sends a minimal tool-calling request to each available model and
+        checks whether the response contains a ``tool_calls`` field. This
+        is a trial-and-error approach because the SAIA API does not expose
+        tool support as model metadata.
+
+        Args:
+            verbose: If ``True``, print per-model results during probing.
+
+        Returns:
+            A list of model ID strings that responded with a tool call.
+
+        Note:
+            This method consumes API quota (one request per model) and may
+            take several minutes depending on the number of available models.
+        """
+        model_ids = self.list_ids()
+
+        _PROBE_TOOLS = [
+            {
+                "type": "function",
+                "function": {
+                    "name": "probe",
+                    "description": "Test probe.",
+                    "parameters": {
+                        "type": "object",
+                        "properties": {"x": {"type": "string"}},
+                        "required": ["x"],
+                    },
+                },
+            }
+        ]
+        _PROBE_MESSAGES = [
+            {"role": "user", "content": "Call the probe tool with x='test'."}
+        ]
+
+        capable = []
+        for mid in progress_iter(
+            model_ids, desc="Probing models", unit="model", enabled=not verbose
+        ):
+            try:
+                resp = self._session.post(
+                    f"{self._base_url}/chat/completions",
+                    json={
+                        "model": mid,
+                        "messages": _PROBE_MESSAGES,
+                        "tools": _PROBE_TOOLS,
+                        "max_tokens": 50,
+                    },
+                    timeout=30,
+                )
+                data = resp.json()
+                msg = data.get("choices", [{}])[0].get("message", {})
+                has_tools = bool(msg.get("tool_calls"))
+                if has_tools:
+                    capable.append(mid)
+                if verbose:
+                    status = "tool_calls" if has_tools else "no tools"
+                    print(f"  {mid:<45} {status}")
+            except Exception as e:
+                if verbose:
+                    print(f"  {mid:<45} error: {e}")
+
+        return capable
+
+    def __repr__(self):
+        return f"ModelsService(base_url={self._base_url!r})"