saia-python 0.4.1__py3-none-any.whl

saia_python/__init__.py

@@ -0,0 +1,253 @@
+"""SAIA Python — a wrapper for the GWDG SAIA platform REST API.
+
+Provides both an object-oriented and a functional interface::
+
+    # OOP
+    from saia_python import SAIAClient
+    client = SAIAClient(api_key="...")
+    client.models.list_ids()
+
+    # Functional
+    from saia_python import list_model_ids, chat_completion
+    list_model_ids()
+"""
+
+from __future__ import annotations
+
+import concurrent.futures
+from importlib.metadata import PackageNotFoundError, version
+
+from ._streaming import SSEStream
+from .arcana_references import (
+    ArcanaReference,
+    ParsedReferences,
+    is_arcana_event,
+    parse_arcana_references,
+    parse_reference_entries,
+)
+from .auth import (
+    DEFAULT_BASE_URL,
+    add_arcana_to_config,
+    load_api_key,
+    load_arcana_ids,
+    load_config,
+    load_username,
+    remove_arcana_from_config,
+    resolve_base_url,
+)
+from .client import SAIAClient
+from .documents import ConversionResult
+from .exceptions import APIError, AuthenticationError, RateLimitError, SAIAError
+from .openai_compat import create_openai_client
+from .rate_limits import RateLimitInfo, parse_rate_limits
+from .responses import text_of
+
+try:
+    __version__ = version("saia-python")
+except PackageNotFoundError:
+    __version__ = "0.0.0-dev"
+
+__all__ = [
+    # Meta
+    "__version__",
+    # Client
+    "SAIAClient",
+    "resolve_base_url",
+    "DEFAULT_BASE_URL",
+    "create_openai_client",
+    # Auth
+    "load_api_key",
+    "load_arcana_ids",
+    "load_config",
+    "load_username",
+    "add_arcana_to_config",
+    "remove_arcana_from_config",
+    # Exceptions
+    "SAIAError",
+    "AuthenticationError",
+    "RateLimitError",
+    "APIError",
+    # Rate limits
+    "RateLimitInfo",
+    "parse_rate_limits",
+    # Response helpers
+    "text_of",
+    "SSEStream",
+    # ARCANA reference parsing
+    "ArcanaReference",
+    "ParsedReferences",
+    "parse_arcana_references",
+    "parse_reference_entries",
+    "is_arcana_event",
+    # Functional API
+    "list_models",
+    "list_model_ids",
+    "chat_completion",
+    "transcribe",
+    "translate",
+    "list_arcanas",
+    "get_arcana",
+    "upload_to_arcana",
+    "arcana_chat",
+    "get_rate_limits",
+    "convert_document",
+    "ConversionResult",
+]
+
+
+def _make_client(api_key: str | None = None, base_url: str | None = None) -> SAIAClient:
+    kwargs: dict = {}
+    if api_key is not None:
+        kwargs["api_key"] = api_key
+    if base_url is not None:
+        kwargs["base_url"] = base_url
+    return SAIAClient(**kwargs)
+
+
+# --- Models ---
+
+
+def list_models(
+    *, api_key: str | None = None, base_url: str | None = None
+) -> list[dict]:
+    """List all available models (functional API)."""
+    return _make_client(api_key, base_url).models.list()
+
+
+def list_model_ids(
+    *, api_key: str | None = None, base_url: str | None = None
+) -> list[str]:
+    """List model ID strings (functional API)."""
+    return _make_client(api_key, base_url).models.list_ids()
+
+
+# --- Chat ---
+
+
+def chat_completion(
+    model: str,
+    messages: list[dict],
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    **kwargs,
+) -> dict | SSEStream:
+    """Send a chat completion request (functional API).
+
+    See :meth:`ChatService.completions` for full parameter docs. With
+    ``stream=True`` this returns an iterable ``SSEStream`` instead of a dict.
+    """
+    return _make_client(api_key, base_url).chat.completions(
+        model=model, messages=messages, **kwargs
+    )
+
+
+# --- Voice ---
+
+
+def transcribe(
+    file_path: str,
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    **kwargs,
+) -> str | concurrent.futures.Future[str]:
+    """Transcribe an audio file (functional API).
+
+    See :meth:`VoiceService.transcribe` for full parameter docs. Passing
+    ``wait=False`` returns a :class:`concurrent.futures.Future` instead of
+    the transcription string.
+    """
+    return _make_client(api_key, base_url).voice.transcribe(file_path, **kwargs)
+
+
+def translate(
+    file_path: str,
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    **kwargs,
+) -> str | concurrent.futures.Future[str]:
+    """Translate an audio file to English (functional API).
+
+    See :meth:`VoiceService.translate` for full parameter docs. Passing
+    ``wait=False`` returns a :class:`concurrent.futures.Future` instead of
+    the translation string.
+    """
+    return _make_client(api_key, base_url).voice.translate(file_path, **kwargs)
+
+
+# --- ARCANA ---
+
+
+def list_arcanas(
+    *, api_key: str | None = None, base_url: str | None = None
+) -> list[dict]:
+    """List all arcanas (functional API)."""
+    return _make_client(api_key, base_url).arcana.list()
+
+
+def get_arcana(
+    name: str, *, api_key: str | None = None, base_url: str | None = None
+) -> dict:
+    """Get a specific arcana by name (functional API)."""
+    return _make_client(api_key, base_url).arcana.get(name)
+
+
+def upload_to_arcana(
+    name: str,
+    file_path: str,
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+) -> dict | None:
+    """Upload a file to an arcana (functional API)."""
+    return _make_client(api_key, base_url).arcana.upload(name, file_path)
+
+
+def arcana_chat(
+    model: str,
+    messages: list[dict],
+    arcana_id: str,
+    *,
+    api_key: str | None = None,
+    base_url: str | None = None,
+    **kwargs,
+) -> dict | SSEStream:
+    """Chat with RAG context from an arcana (functional API).
+
+    With ``stream=True`` this returns an iterable ``SSEStream`` instead of a dict.
+    """
+    return _make_client(api_key, base_url).arcana.chat(
+        model=model, messages=messages, arcana_id=arcana_id, **kwargs
+    )
+
+
+# --- Rate Limits ---
+
+
+def get_rate_limits(
+    *, api_key: str | None = None, base_url: str | None = None
+) -> RateLimitInfo:
+    """Fetch current rate-limit status (functional API)."""
+    return _make_client(api_key, base_url).get_rate_limits()
+
+
+# --- Documents ---
+
+
+def convert_document(
+    file_path: str,
+    *,
+    response_type: str = "markdown",
+    api_key: str | None = None,
+    base_url: str | None = None,
+    **kwargs,
+) -> ConversionResult:
+    """Convert a document using the Docling service (functional API).
+
+    See :meth:`DocumentService.convert` for full parameter docs.
+    """
+    return _make_client(api_key, base_url).documents.convert(
+        file_path, response_type=response_type, **kwargs
+    )

saia_python/_http.py

@@ -0,0 +1,71 @@
+"""Shared HTTP plumbing used by more than one service.
+
+Kept in one place so the chat-completion request shape and the
+background-thread ``Session`` helper each have a single implementation,
+rather than being copied across :mod:`saia_python.chat`,
+:mod:`saia_python.arcana`, and :mod:`saia_python.voice`.
+"""
+
+from __future__ import annotations
+
+import requests
+
+from ._streaming import SSEStream
+from .exceptions import raise_for_status
+from .rate_limits import parse_rate_limits
+
+
+def new_session_like(template: requests.Session) -> requests.Session:
+    """Return a fresh :class:`requests.Session` mirroring ``template``'s headers.
+
+    Background-thread work must not reuse the caller's ``Session`` —
+    ``requests.Session`` is not guaranteed thread-safe, and sharing its
+    connection pool across threads can corrupt in-flight requests. Both the
+    non-blocking Voice path and the fire-and-forget ARCANA index trigger spin
+    up their own ``Session`` through this helper so they never race the
+    client's.
+    """
+    session = requests.Session()
+    session.headers.update(template.headers)
+    return session
+
+
+def post_chat_completion(
+    session: requests.Session,
+    url: str,
+    body: dict,
+    *,
+    headers: dict | None = None,
+    stream: bool = False,
+) -> dict | SSEStream:
+    """POST a chat-completion request and normalise the response.
+
+    Shared by :meth:`ChatService.completions` and :meth:`ArcanaService.chat`:
+    both hit the same ``/chat/completions`` endpoint with identical
+    stream/non-stream handling and rate-limit surfacing — only the request
+    ``body`` fields and auth ``headers`` differ, so those stay with the caller.
+
+    Args:
+        session: The authenticated :class:`requests.Session`.
+        url: The fully-qualified ``/chat/completions`` URL.
+        body: The request JSON body (already assembled by the caller).
+        headers: Per-request headers. ``None`` uses the session defaults
+            (the Bearer auth + ``Accept: application/json``).
+        stream: When ``True``, request SSE and return an :class:`SSEStream`.
+
+    Returns:
+        When ``stream=False``: the response dict with an extra
+        ``"_rate_limits"`` key (a JSON-serializable dict). When ``stream=True``:
+        an :class:`SSEStream` whose ``rate_limits`` attribute holds the same dict.
+    """
+    if stream:
+        stream_body = {**body, "stream": True}
+        stream_headers = {**(headers or {}), "Accept": "text/event-stream"}
+        resp = session.post(url, json=stream_body, headers=stream_headers, stream=True)
+        return SSEStream(resp)
+
+    resp = session.post(url, json=body, headers=headers)
+    raise_for_status(resp)
+    result = resp.json()
+    result["_rate_limits"] = parse_rate_limits(resp.headers).to_dict()
+    return result

saia_python/_streaming.py

@@ -0,0 +1,88 @@
+"""Shared SSE (Server-Sent Events) streaming iterator."""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Generator
+from typing import TYPE_CHECKING
+
+from .exceptions import raise_for_status
+from .rate_limits import parse_rate_limits
+
+if TYPE_CHECKING:
+    import requests
+
+
+def iter_sse(response: requests.Response) -> Generator[dict, None, None]:
+    """Yield parsed JSON chunks from an SSE stream.
+
+    Handles the ``data: {...}`` / ``data: [DONE]`` protocol used by
+    the OpenAI-compatible SAIA API.
+
+    Args:
+        response: A streaming :class:`requests.Response` (``stream=True``).
+
+    Yields:
+        Parsed JSON dicts for each ``data:`` line.
+    """
+    raise_for_status(response)
+    try:
+        for line in response.iter_lines(decode_unicode=True):
+            if not line or not line.startswith("data:"):
+                continue
+            payload = line[len("data:") :].strip()
+            if payload == "[DONE]":
+                return
+            try:
+                yield json.loads(payload)
+            except json.JSONDecodeError:
+                continue
+    finally:
+        # Release the underlying connection back to the pool whether the
+        # stream finished, hit [DONE], errored, or the consumer broke early
+        # (GeneratorExit). Without this, an abandoned stream leaks the socket
+        # until garbage collection.
+        response.close()
+
+
+class SSEStream:
+    """Iterable over SSE chunks that also exposes parsed rate-limit info.
+
+    Wraps a streaming :class:`requests.Response` so callers can both iterate
+    the decoded chunks (``for chunk in stream``) and read the rate-limit
+    headers via :attr:`rate_limits` — available immediately, since headers
+    arrive before the streamed body. Mirrors the ``_rate_limits`` key that
+    non-streaming responses carry, giving both paths rate-limit parity.
+
+    Attributes:
+        rate_limits: A JSON-serializable dict of this response's rate-limit
+            headers (the same shape as :class:`~saia_python.RateLimitInfo`
+            via ``to_dict()``).
+    """
+
+    def __init__(self, response: requests.Response):
+        self._response = response
+        self.rate_limits: dict = parse_rate_limits(response.headers).to_dict()
+        self._chunks = iter_sse(response)
+
+    def __iter__(self) -> SSEStream:
+        return self
+
+    def __next__(self) -> dict:
+        return next(self._chunks)
+
+    def close(self) -> None:
+        """Close the stream, releasing the underlying connection.
+
+        Closes both the chunk generator (running its cleanup) and the
+        response directly, so the connection is released even if iteration
+        never started.
+        """
+        self._chunks.close()
+        self._response.close()
+
+    def __enter__(self) -> SSEStream:
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()

saia_python/_util.py

@@ -0,0 +1,29 @@
+"""Small internal utilities shared across services."""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Iterator
+
+
+def progress_iter(
+    items: Iterable,
+    *,
+    desc: str,
+    unit: str = "file",
+    enabled: bool = True,
+) -> Iterator:
+    """Wrap ``items`` in a ``tqdm`` progress bar when available and ``enabled``.
+
+    ``tqdm`` is an optional dependency, so this degrades to a plain iterator
+    when it is not installed (or when ``enabled`` is ``False`` — e.g. a caller
+    that prints its own per-item lines instead of a bar). Centralises the
+    ``try: from tqdm.auto import tqdm`` dance that several services otherwise
+    repeat.
+    """
+    if not enabled:
+        return iter(items)
+    try:
+        from tqdm.auto import tqdm
+    except ImportError:
+        return iter(items)
+    return tqdm(items, desc=desc, unit=unit)