polza-extras 0.2.0__py3-none-any.whl

polza_extras/__init__.py

@@ -0,0 +1,57 @@
+from ._models import (
+    Balance,
+    ChatResult,
+    KeyInfo,
+    ModelInfo,
+    ModelPage,
+    Pricing,
+    ReasoningInfo,
+    StreamEvent,
+    ToolCall,
+    ToolCallDelta,
+    Usage,
+)
+from .client import DEFAULT_BASE_URL, AsyncPolza, Polza
+from .errors import (
+    AuthError,
+    EmptyResponseError,
+    ForbiddenError,
+    InsufficientFundsError,
+    NotFoundError,
+    PolzaConnectionError,
+    PolzaError,
+    RateLimitError,
+    RequestError,
+    ServerError,
+    StreamError,
+)
+from .retry import RetryPolicy
+
+__all__ = [
+    "DEFAULT_BASE_URL",
+    "AsyncPolza",
+    "AuthError",
+    "Balance",
+    "ChatResult",
+    "EmptyResponseError",
+    "ForbiddenError",
+    "InsufficientFundsError",
+    "KeyInfo",
+    "ModelInfo",
+    "ModelPage",
+    "NotFoundError",
+    "Polza",
+    "PolzaConnectionError",
+    "PolzaError",
+    "Pricing",
+    "RateLimitError",
+    "ReasoningInfo",
+    "RequestError",
+    "RetryPolicy",
+    "ServerError",
+    "StreamError",
+    "StreamEvent",
+    "ToolCall",
+    "ToolCallDelta",
+    "Usage",
+]

polza_extras/_json.py

@@ -0,0 +1,13 @@
+"""Typed access to parsed JSON: containers get a known shape; a missing or wrong one is empty."""
+
+from __future__ import annotations
+
+from typing import Any, cast
+
+
+def obj(value: Any) -> dict[str, Any]:
+    return cast("dict[str, Any]", value) if isinstance(value, dict) else {}
+
+
+def arr(value: Any) -> list[Any]:
+    return cast("list[Any]", value) if isinstance(value, list) else []

polza_extras/_models.py

@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+import json
+from datetime import datetime
+from decimal import Decimal
+from typing import Any, cast
+
+from pydantic import BaseModel, Field
+
+from ._json import obj
+
+
+class Usage(BaseModel):
+    """Token counts. `cached_tokens` and `cache_write_tokens` are part of `prompt_tokens`;
+    `reasoning_tokens` are part of `completion_tokens`."""
+
+    prompt_tokens: int = 0
+    completion_tokens: int = 0
+    total_tokens: int = 0
+    cached_tokens: int = 0
+    cache_write_tokens: int = 0
+    reasoning_tokens: int = 0
+
+
+class ToolCall(BaseModel):
+    """A function call the model asked for; `arguments` is the JSON string as the model wrote it."""
+
+    id: str
+    name: str
+    arguments: str = ""
+
+    def parsed_arguments(self) -> dict[str, Any]:
+        """`arguments` decoded; raises ValueError when they are not a JSON object."""
+        value: Any = json.loads(self.arguments or "{}")
+        if not isinstance(value, dict):
+            raise ValueError(f"tool call arguments are not a JSON object: {self.arguments!r}")
+        return cast("dict[str, Any]", value)
+
+
+class ChatResult(BaseModel):
+    text: str
+    tool_calls: list[ToolCall] = Field(default_factory=list[ToolCall])
+    reasoning: str | None = None
+    finish_reason: str | None = None
+    usage: Usage
+    cost_rub: Decimal | None = None
+    model: str | None = None
+    provider: str | None = None
+    id: str | None = None
+    raw: dict[str, Any]
+
+
+class ToolCallDelta(BaseModel):
+    """A streamed piece of a tool call. The first piece of call `index` has `id` and `name`;
+    the `arguments` fragments of one index concatenate to the full JSON string."""
+
+    index: int
+    id: str | None = None
+    name: str | None = None
+    arguments: str = ""
+
+
+class StreamEvent(BaseModel):
+    """One streamed chunk: a text delta, or the final chunk with usage and ruble cost."""
+
+    delta: str = ""
+    tool_calls: list[ToolCallDelta] = Field(default_factory=list[ToolCallDelta])
+    reasoning: str | None = None
+    finish_reason: str | None = None
+    usage: Usage | None = None
+    cost_rub: Decimal | None = None
+
+
+_MILLION = Decimal(1_000_000)
+
+
+class Pricing(BaseModel):
+    """Model prices in rubles per million tokens (None when the catalog does not list a price)."""
+
+    prompt_per_million: Decimal | None = None
+    completion_per_million: Decimal | None = None
+    internal_reasoning_per_million: Decimal | None = None
+    input_cache_read_per_million: Decimal | None = None
+    input_cache_write_per_million: Decimal | None = None
+    currency: str = "RUB"
+
+    def estimate(
+        self, prompt_tokens: int, completion_tokens: int, cached_tokens: int = 0
+    ) -> Decimal:
+        """Approximate cost of one call; `cached_tokens` are billed at the cache-read price."""
+        prompt = self.prompt_per_million or Decimal(0)
+        cache_read = (
+            self.input_cache_read_per_million
+            if self.input_cache_read_per_million is not None
+            else prompt
+        )
+        completion = self.completion_per_million or Decimal(0)
+        total = (
+            prompt * (prompt_tokens - cached_tokens)
+            + cache_read * cached_tokens
+            + completion * completion_tokens
+        )
+        return total / _MILLION
+
+
+class ReasoningInfo(BaseModel):
+    mandatory: bool = False
+    default_enabled: bool | None = None
+    supported_efforts: list[str] = Field(default_factory=list)
+    default_effort: str | None = None
+
+
+class ModelInfo(BaseModel):
+    """A catalog entry; pricing and parameters come from the model's top provider."""
+
+    id: str
+    name: str = ""
+    type: str | None = None
+    context_length: int | None = None
+    max_completion_tokens: int | None = None
+    reasoning: ReasoningInfo | None = None
+    pricing: Pricing | None = None
+    supported_parameters: list[str] = Field(default_factory=list)
+    raw: dict[str, Any] = Field(default_factory=dict, repr=False)
+
+    @classmethod
+    def from_api(cls, data: dict[str, Any]) -> ModelInfo:
+        top = obj(data.get("top_provider"))
+        return cls(
+            id=data["id"],
+            name=data.get("name") or "",
+            type=data.get("type"),
+            context_length=data.get("context_length"),
+            max_completion_tokens=data.get("max_completion_tokens"),
+            reasoning=data.get("reasoning") or top.get("reasoning"),
+            pricing=top.get("pricing") or data.get("pricing"),
+            supported_parameters=top.get("supported_parameters") or [],
+            raw=data,
+        )
+
+
+class ModelPage(BaseModel):
+    items: list[ModelInfo]
+    page: int = 1
+    limit: int | None = None
+    total: int | None = None
+    total_pages: int | None = None
+
+
+class Balance(BaseModel):
+    amount: Decimal
+    available: Decimal
+    reserved: Decimal
+    spent: Decimal
+    updated_at: datetime | None = None
+
+
+class KeyInfo(BaseModel):
+    label: str | None = None
+    limit: Decimal | None = None
+    limit_remaining: Decimal | None = None
+    limit_reset: str | None = None
+    usage: Decimal | None = None
+    usage_daily: Decimal | None = None
+    usage_weekly: Decimal | None = None
+    usage_monthly: Decimal | None = None

polza_extras/chat.py

@@ -0,0 +1,120 @@
+from __future__ import annotations
+
+from decimal import Decimal
+from typing import Any
+
+from ._json import arr, obj
+from ._models import ChatResult, StreamEvent, ToolCall, ToolCallDelta, Usage
+from .errors import EmptyResponseError
+
+
+def _cost_rub(usage: dict[str, Any]) -> Decimal | None:
+    for key in ("cost_rub", "cost"):
+        value = usage.get(key)
+        if isinstance(value, (int, float, str)):
+            return Decimal(str(value))
+    return None
+
+
+def _int(data: dict[str, Any], key: str) -> int:
+    return int(data.get(key) or 0)
+
+
+def _usage(usage: dict[str, Any]) -> Usage:
+    prompt_details = obj(usage.get("prompt_tokens_details"))
+    completion_details = obj(usage.get("completion_tokens_details"))
+    return Usage(
+        prompt_tokens=_int(usage, "prompt_tokens"),
+        completion_tokens=_int(usage, "completion_tokens"),
+        total_tokens=_int(usage, "total_tokens"),
+        cached_tokens=_int(prompt_details, "cached_tokens"),
+        cache_write_tokens=_int(prompt_details, "cache_write_tokens"),
+        reasoning_tokens=_int(completion_details, "reasoning_tokens"),
+    )
+
+
+def _tool_calls(message: dict[str, Any]) -> list[ToolCall]:
+    calls: list[ToolCall] = []
+    for item in arr(message.get("tool_calls")):
+        call = obj(item)
+        function = obj(call.get("function"))
+        calls.append(
+            ToolCall(
+                id=call.get("id") or "",
+                name=function.get("name") or "",
+                arguments=function.get("arguments") or "",
+            )
+        )
+    return calls
+
+
+def parse_chat(data: dict[str, Any], reasoning_fallback: bool = False) -> ChatResult:
+    choices = arr(data.get("choices"))
+    if not choices:
+        raise EmptyResponseError("Polza returned a response without choices")
+    choice = obj(choices[0])
+    message = obj(choice.get("message"))
+    finish_reason: str | None = choice.get("finish_reason")
+    reasoning: str | None = message.get("reasoning") or message.get("reasoning_content") or None
+    text: str = message.get("content") or ""
+    tool_calls = _tool_calls(message)
+    usage = obj(data.get("usage"))
+    parsed_usage = _usage(usage)
+    if not text and not tool_calls:
+        if reasoning and reasoning_fallback:
+            text = reasoning
+        else:
+            hint = ""
+            if finish_reason == "length" and (reasoning or message.get("reasoning_details")):
+                hint = "; max_tokens was spent on thinking"
+            raise EmptyResponseError(
+                f"Polza returned empty content (finish_reason={finish_reason}{hint})",
+                finish_reason=finish_reason,
+                reasoning=reasoning,
+                usage=parsed_usage,
+                raw=data,
+                cost_rub=_cost_rub(usage),
+            )
+    return ChatResult(
+        text=text,
+        tool_calls=tool_calls,
+        reasoning=reasoning,
+        finish_reason=finish_reason,
+        usage=parsed_usage,
+        cost_rub=_cost_rub(usage),
+        model=data.get("model"),
+        provider=data.get("provider"),
+        id=data.get("id"),
+        raw=data,
+    )
+
+
+def _tool_call_deltas(delta: dict[str, Any]) -> list[ToolCallDelta]:
+    deltas: list[ToolCallDelta] = []
+    for position, item in enumerate(arr(delta.get("tool_calls"))):
+        call = obj(item)
+        function = obj(call.get("function"))
+        deltas.append(
+            ToolCallDelta(
+                index=call.get("index", position),
+                id=call.get("id"),
+                name=function.get("name"),
+                arguments=function.get("arguments") or "",
+            )
+        )
+    return deltas
+
+
+def parse_stream_chunk(data: dict[str, Any]) -> StreamEvent:
+    choices = arr(data.get("choices"))
+    choice = obj(choices[0]) if choices else {}
+    delta = obj(choice.get("delta"))
+    usage = obj(data.get("usage"))
+    return StreamEvent(
+        delta=delta.get("content") or "",
+        tool_calls=_tool_call_deltas(delta),
+        reasoning=delta.get("reasoning") or delta.get("reasoning_content") or None,
+        finish_reason=choice.get("finish_reason"),
+        usage=_usage(usage) if usage else None,
+        cost_rub=_cost_rub(usage) if usage else None,
+    )

polza_extras/client.py

@@ -0,0 +1,452 @@
+from __future__ import annotations
+
+import asyncio
+import os
+import random
+import time
+from collections.abc import AsyncIterator, Awaitable, Callable, Iterator
+from typing import Any, Self, TypeVar
+
+import openai
+from pydantic import BaseModel
+
+from ._json import arr, obj
+from ._models import Balance, ChatResult, KeyInfo, ModelInfo, ModelPage, StreamEvent
+from .chat import parse_chat, parse_stream_chunk
+from .errors import PolzaConnectionError, PolzaError, StreamError, error_from_response
+from .retry import RetryPolicy
+
+DEFAULT_BASE_URL = "https://polza.ai/api/v1"
+API_KEY_ENV = "POLZA_AI_API_KEY"
+
+T = TypeVar("T")
+
+
+# --- shared, I/O-free pieces of Polza and AsyncPolza ---------------------------------------------
+
+
+def _api_key(api_key: str | None) -> str:
+    key = api_key or os.environ.get(API_KEY_ENV)
+    if not key:
+        raise ValueError(f"Polza API key missing: pass api_key or set {API_KEY_ENV}")
+    return key
+
+
+def _translate(exc: openai.APIStatusError | openai.APIConnectionError) -> PolzaError:
+    """The PolzaError for an openai SDK request error."""
+    if isinstance(exc, openai.APIConnectionError):
+        return PolzaConnectionError(str(exc))
+    try:
+        body = exc.response.json()
+    except ValueError:
+        body = None
+    return error_from_response(exc.status_code, body, dict(exc.response.headers))
+
+
+def _stream_error(exc: openai.APIError) -> StreamError:
+    body = obj(exc.body)
+    return StreamError(
+        200,
+        str(body.get("message") or exc.message),
+        code=body.get("code"),
+        trace_id=body.get("trace_id"),
+        metadata=body.get("metadata"),
+    )
+
+
+def _schema_format(json_schema: dict[str, Any] | type[BaseModel], name: str) -> dict[str, Any]:
+    schema = json_schema if isinstance(json_schema, dict) else json_schema.model_json_schema()
+    return {"type": "json_schema", "json_schema": {"name": name, "schema": schema, "strict": True}}
+
+
+def _chat_params(
+    model: str,
+    messages: list[dict[str, Any]],
+    max_tokens: int | None,
+    temperature: float | None,
+    json: bool,
+    extra: dict[str, Any],
+    stream: bool = False,
+    json_schema: dict[str, Any] | type[BaseModel] | None = None,
+    schema_name: str = "response",
+) -> dict[str, Any]:
+    params: dict[str, Any] = {"model": model, "messages": messages}
+    if max_tokens is not None:
+        params["max_tokens"] = max_tokens
+    if temperature is not None:
+        params["temperature"] = temperature
+    if json_schema is not None:
+        if json or "response_format" in extra:
+            raise ValueError("pass json_schema alone, without json=True or response_format")
+        extra["response_format"] = _schema_format(json_schema, schema_name)
+    if json:
+        extra.setdefault("response_format", {"type": "json_object"})
+    if extra:
+        params["extra_body"] = extra
+    if stream:
+        params["stream"] = True
+        params["stream_options"] = {"include_usage": True}
+    return params
+
+
+def _get_options(
+    params: dict[str, Any] | None, language: str | None
+) -> Any:  # openai RequestOptions
+    headers = {"accept-language": language} if language else {}
+    return {
+        "params": {k: v for k, v in (params or {}).items() if v is not None},
+        "headers": headers,
+    }
+
+
+def _catalog_query(
+    search: str | None,
+    type: str | None,
+    providers: list[str] | None,
+    page: int,
+    limit: int | None,
+    filters: dict[str, Any],
+) -> dict[str, Any]:
+    return {
+        "search": search,
+        "type": type,
+        "providers": ",".join(providers) if providers else None,
+        "page": page,
+        "limit": limit,
+        **filters,
+    }
+
+
+def _parse_page(data: Any, page: int, limit: int | None) -> ModelPage:
+    data = obj(data)
+    meta = obj(data.get("meta"))
+    return ModelPage(
+        items=[ModelInfo.from_api(obj(m)) for m in arr(data.get("data"))],
+        page=meta.get("page", page),
+        limit=meta.get("limit", limit),
+        total=meta.get("total"),
+        total_pages=meta.get("totalPages"),
+    )
+
+
+def _balance_url(base_url: str) -> str:
+    return base_url.rstrip("/").removesuffix("/v1") + "/v2/balance"
+
+
+def _parse_balance(data: Any) -> Balance:
+    return Balance(
+        amount=data["amount"],
+        available=data["available"],
+        reserved=data["reservedAmount"],
+        spent=data["spentAmount"],
+        updated_at=data.get("updatedAt"),
+    )
+
+
+class Polza:
+    """Polza.ai facade: an OpenAI SDK client plus Polza-specific cost, errors and retries."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str = DEFAULT_BASE_URL,
+        retry: RetryPolicy | None = None,
+        timeout: float = 60.0,
+        sleep: Callable[[float], None] = time.sleep,
+        http_client: Any | None = None,
+    ) -> None:
+        self.base_url = base_url
+        self.retry = retry or RetryPolicy()
+        self._sleep = sleep
+        self._rng = random.Random()
+        # Use the SDK's own client flavor (httpx or its vendored httpx2, depending on the openai
+        # version) so errors are translated correctly; the Polza-only endpoints reuse it.
+        self.http: Any = http_client or openai.DefaultHttpxClient(timeout=timeout)
+        # SDK retries are off: Polza has its own retry rules (see RetryPolicy).
+        self.openai = openai.OpenAI(
+            base_url=base_url, api_key=_api_key(api_key), max_retries=0, http_client=self.http
+        )
+        # Typed as Any: create(**params) cannot be matched against the SDK's overloads.
+        self._completions: Any = self.openai.chat.completions
+
+    def close(self) -> None:
+        """Close the HTTP connections (also the http_client passed in)."""
+        self.openai.close()
+
+    def __enter__(self) -> Self:
+        return self
+
+    def __exit__(self, *exc: object) -> None:
+        self.close()
+
+    def _with_retries(self, call: Callable[[], T]) -> T:
+        attempt = 0
+        while True:
+            try:
+                return call()
+            except (openai.APIStatusError, openai.APIConnectionError) as exc:
+                error = _translate(exc)
+            delay = self.retry.delay_for(error, attempt, self._rng)
+            if delay is None:
+                raise error
+            self._sleep(delay)
+            attempt += 1
+
+    def chat(
+        self,
+        model: str,
+        messages: list[dict[str, Any]],
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+        reasoning_fallback: bool = False,
+        json: bool = False,
+        json_schema: dict[str, Any] | type[BaseModel] | None = None,
+        schema_name: str = "response",
+        **extra: Any,
+    ) -> ChatResult:
+        """One chat completion. `json=True` asks for any JSON object; `json_schema` (a JSON Schema
+        dict or a pydantic model class) asks for strict structured output under `schema_name`.
+        Other keyword arguments (tools, reasoning, provider, ...) go to Polza as is."""
+        params = _chat_params(
+            model,
+            messages,
+            max_tokens,
+            temperature,
+            json,
+            extra,
+            stream=False,
+            json_schema=json_schema,
+            schema_name=schema_name,
+        )
+
+        def call() -> dict[str, Any]:
+            raw = self._completions.with_raw_response.create(**params)
+            data: dict[str, Any] = raw.http_response.json()
+            return data
+
+        return parse_chat(self._with_retries(call), reasoning_fallback)
+
+    def chat_stream(
+        self,
+        model: str,
+        messages: list[dict[str, Any]],
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+        json: bool = False,
+        json_schema: dict[str, Any] | type[BaseModel] | None = None,
+        schema_name: str = "response",
+        **extra: Any,
+    ) -> Iterator[StreamEvent]:
+        """Stream a chat completion. The last event carries `usage` and `cost_rub`.
+
+        Errors before the first chunk are retried like `chat()`; an error event inside the stream
+        raises StreamError (text already yielded stays yielded).
+        """
+        params = _chat_params(
+            model,
+            messages,
+            max_tokens,
+            temperature,
+            json,
+            extra,
+            stream=True,
+            json_schema=json_schema,
+            schema_name=schema_name,
+        )
+        stream = self._with_retries(lambda: self._completions.create(**params))
+        try:
+            for chunk in stream:
+                yield parse_stream_chunk(chunk.to_dict())
+        except openai.APIError as exc:
+            raise _stream_error(exc) from exc
+        finally:
+            stream.close()  # also when the caller stops reading early
+
+    # --- Polza-only endpoints ------------------------------------------------
+
+    def _get(
+        self, path: str, params: dict[str, Any] | None = None, language: str | None = None
+    ) -> Any:
+        """GET a Polza endpoint through the SDK client, so errors and retries match chat()."""
+        options = _get_options(params, language)
+        return self._with_retries(lambda: self.openai.get(path, cast_to=object, options=options))
+
+    def models(
+        self,
+        search: str | None = None,
+        type: str | None = None,
+        providers: list[str] | None = None,
+        page: int = 1,
+        limit: int | None = None,
+        language: str = "ru",
+        **filters: Any,
+    ) -> ModelPage:
+        """One page of /models/catalog. Extra filters (e.g. sortBy, contextLengthMin) pass as is."""
+        data = self._get(
+            "/models/catalog",
+            _catalog_query(search, type, providers, page, limit, filters),
+            language,
+        )
+        return _parse_page(data, page, limit)
+
+    def model(self, model_id: str, language: str = "ru") -> ModelInfo:
+        """Details of one model; raises NotFoundError for an unknown id."""
+        data = self._get(f"/models/{model_id}", language=language)
+        return ModelInfo.from_api(data.get("data", data))
+
+    def balance(self) -> Balance:
+        """Organization balance in rubles (GET /api/v2/balance)."""
+        return _parse_balance(self._get(_balance_url(self.base_url)))
+
+    def key_info(self) -> KeyInfo:
+        """Limits and usage of the current API key (GET /v1/key)."""
+        data = self._get("/key")
+        return KeyInfo.model_validate(data.get("data", data))
+
+
+class AsyncPolza:
+    """Async mirror of Polza on openai.AsyncOpenAI; same methods, awaited."""
+
+    def __init__(
+        self,
+        api_key: str | None = None,
+        base_url: str = DEFAULT_BASE_URL,
+        retry: RetryPolicy | None = None,
+        timeout: float = 60.0,
+        sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
+        http_client: Any | None = None,
+    ) -> None:
+        self.base_url = base_url
+        self.retry = retry or RetryPolicy()
+        self._sleep = sleep
+        self._rng = random.Random()
+        self.http: Any = http_client or openai.DefaultAsyncHttpxClient(timeout=timeout)
+        self.openai = openai.AsyncOpenAI(
+            base_url=base_url, api_key=_api_key(api_key), max_retries=0, http_client=self.http
+        )
+        # Typed as Any: create(**params) cannot be matched against the SDK's overloads.
+        self._completions: Any = self.openai.chat.completions
+
+    async def close(self) -> None:
+        await self.openai.close()
+
+    async def __aenter__(self) -> Self:
+        return self
+
+    async def __aexit__(self, *exc: object) -> None:
+        await self.close()
+
+    async def _with_retries(self, call: Callable[[], Awaitable[T]]) -> T:
+        attempt = 0
+        while True:
+            try:
+                return await call()
+            except (openai.APIStatusError, openai.APIConnectionError) as exc:
+                error = _translate(exc)
+            delay = self.retry.delay_for(error, attempt, self._rng)
+            if delay is None:
+                raise error
+            await self._sleep(delay)
+            attempt += 1
+
+    async def chat(
+        self,
+        model: str,
+        messages: list[dict[str, Any]],
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+        reasoning_fallback: bool = False,
+        json: bool = False,
+        json_schema: dict[str, Any] | type[BaseModel] | None = None,
+        schema_name: str = "response",
+        **extra: Any,
+    ) -> ChatResult:
+        """Async version of Polza.chat."""
+        params = _chat_params(
+            model,
+            messages,
+            max_tokens,
+            temperature,
+            json,
+            extra,
+            stream=False,
+            json_schema=json_schema,
+            schema_name=schema_name,
+        )
+
+        async def call() -> dict[str, Any]:
+            raw = await self._completions.with_raw_response.create(**params)
+            data: dict[str, Any] = raw.http_response.json()
+            return data
+
+        return parse_chat(await self._with_retries(call), reasoning_fallback)
+
+    async def chat_stream(
+        self,
+        model: str,
+        messages: list[dict[str, Any]],
+        max_tokens: int | None = None,
+        temperature: float | None = None,
+        json: bool = False,
+        json_schema: dict[str, Any] | type[BaseModel] | None = None,
+        schema_name: str = "response",
+        **extra: Any,
+    ) -> AsyncIterator[StreamEvent]:
+        """Async version of Polza.chat_stream: `async for event in polza.chat_stream(...)`."""
+        params = _chat_params(
+            model,
+            messages,
+            max_tokens,
+            temperature,
+            json,
+            extra,
+            stream=True,
+            json_schema=json_schema,
+            schema_name=schema_name,
+        )
+
+        async def call() -> Any:
+            return await self._completions.create(**params)
+
+        stream = await self._with_retries(call)
+        try:
+            async for chunk in stream:
+                yield parse_stream_chunk(chunk.to_dict())
+        except openai.APIError as exc:
+            raise _stream_error(exc) from exc
+        finally:
+            await stream.close()
+
+    async def _get(
+        self, path: str, params: dict[str, Any] | None = None, language: str | None = None
+    ) -> Any:
+        options = _get_options(params, language)
+
+        async def call() -> Any:
+            return await self.openai.get(path, cast_to=object, options=options)
+
+        return await self._with_retries(call)
+
+    async def models(
+        self,
+        search: str | None = None,
+        type: str | None = None,
+        providers: list[str] | None = None,
+        page: int = 1,
+        limit: int | None = None,
+        language: str = "ru",
+        **filters: Any,
+    ) -> ModelPage:
+        query = _catalog_query(search, type, providers, page, limit, filters)
+        return _parse_page(await self._get("/models/catalog", query, language), page, limit)
+
+    async def model(self, model_id: str, language: str = "ru") -> ModelInfo:
+        data = await self._get(f"/models/{model_id}", language=language)
+        return ModelInfo.from_api(data.get("data", data))
+
+    async def balance(self) -> Balance:
+        return _parse_balance(await self._get(_balance_url(self.base_url)))
+
+    async def key_info(self) -> KeyInfo:
+        data = await self._get("/key")
+        return KeyInfo.model_validate(data.get("data", data))

polza_extras/errors.py

@@ -0,0 +1,161 @@
+from __future__ import annotations
+
+from collections.abc import Mapping
+from typing import TYPE_CHECKING, Any, Literal
+
+from ._json import obj
+
+if TYPE_CHECKING:
+    from decimal import Decimal
+
+    from ._models import Usage
+
+
+class PolzaError(Exception):
+    """Any error returned by the Polza API, parsed from its error body."""
+
+    def __init__(
+        self,
+        status: int,
+        message: str,
+        code: str | None = None,
+        trace_id: str | None = None,
+        metadata: dict[str, Any] | None = None,
+        details: dict[str, Any] | None = None,
+        retry_after: float | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.retry_after = retry_after
+        self.status = status
+        self.message = message
+        self.code = code
+        self.trace_id = trace_id
+        self.metadata = metadata or {}
+        self.details = details or {}
+
+    def __str__(self) -> str:
+        parts = [str(self.status)]
+        if self.code:
+            parts.append(self.code)
+        text = f"[{' '.join(parts)}] {self.message}"
+        raw = self.metadata.get("raw")
+        if raw and str(raw) not in self.message:
+            text += f": {raw}"
+        if self.trace_id:
+            text += f" (trace_id={self.trace_id})"
+        return text
+
+
+class AuthError(PolzaError):
+    """401: missing, invalid or revoked API key."""
+
+
+class InsufficientFundsError(PolzaError):
+    """402: balance or spending limit exhausted (or a parallel-request reserve conflict)."""
+
+
+class ForbiddenError(PolzaError):
+    """403: access denied, including requests rejected by moderation."""
+
+
+class RateLimitError(PolzaError):
+    """429. `scope` is "own" for Polza's per-organization limit (has Retry-After), or "provider"."""
+
+    @property
+    def scope(self) -> Literal["own", "provider"]:
+        return "own" if self.retry_after is not None else "provider"
+
+
+class ServerError(PolzaError):
+    """408, 5xx: provider timeout, bad gateway, unavailable."""
+
+
+class RequestError(PolzaError):
+    """Other 4xx: bad request, not found, payload too large."""
+
+
+class NotFoundError(RequestError):
+    """404: unknown model or endpoint (a model id without the provider/ prefix is a 404 too)."""
+
+
+class PolzaConnectionError(PolzaError):
+    """The request never got an HTTP response (DNS, connect, timeout)."""
+
+    def __init__(self, message: str) -> None:
+        super().__init__(0, message)
+
+
+class StreamError(PolzaError):
+    """An error event inside a 200 stream (Polza sends `error` + `finish_reason: "error"`).
+
+    Text received before it was already yielded; the request is not retried mid-stream.
+    """
+
+
+class EmptyResponseError(PolzaError):
+    """A 2xx response without usable text: no choices, or empty content.
+
+    `reasoning` carries the model's reasoning text, if any; with `finish_reason == "length"` it
+    usually means `max_tokens` was spent on thinking. `usage` and `raw` (the response body) are
+    kept because the request is still billed.
+    """
+
+    def __init__(
+        self,
+        message: str,
+        finish_reason: str | None = None,
+        reasoning: str | None = None,
+        usage: Usage | None = None,
+        raw: dict[str, Any] | None = None,
+        cost_rub: Decimal | None = None,
+    ) -> None:
+        super().__init__(200, message)
+        self.finish_reason = finish_reason
+        self.reasoning = reasoning
+        self.usage = usage
+        self.raw = raw or {}
+        self.cost_rub = cost_rub  # what the failed call was billed
+
+
+_BY_STATUS: dict[int, type[PolzaError]] = {
+    401: AuthError,
+    402: InsufficientFundsError,
+    403: ForbiddenError,
+    404: NotFoundError,
+    408: ServerError,
+}
+
+
+def _header(headers: Mapping[str, str], name: str) -> str | None:
+    for key, value in headers.items():
+        if key.lower() == name.lower():
+            return value
+    return None
+
+
+def _retry_after(headers: Mapping[str, str]) -> float | None:
+    raw = _header(headers, "retry-after")
+    try:
+        return float(raw) if raw is not None else None
+    except ValueError:
+        return None
+
+
+def error_from_response(status: int, body: Any, headers: Mapping[str, str]) -> PolzaError:
+    error = obj(obj(body).get("error"))
+    top_trace = obj(body).get("trace_id")
+    kwargs: dict[str, Any] = {
+        "message": str(error.get("message") or f"Polza API returned HTTP {status}"),
+        "code": error.get("code"),
+        "trace_id": error.get("trace_id") or top_trace,
+        "metadata": error.get("metadata"),
+        "details": error.get("details"),
+    }
+    kwargs["retry_after"] = _retry_after(headers)
+    if status == 429:
+        return RateLimitError(status, **kwargs)
+    if status in _BY_STATUS:
+        return _BY_STATUS[status](status, **kwargs)
+    if status >= 500:
+        return ServerError(status, **kwargs)
+    return RequestError(status, **kwargs)

polza_extras/retry.py

@@ -0,0 +1,51 @@
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+
+from .errors import PolzaError, RateLimitError
+
+_BACKOFF_STATUSES = frozenset({0, 408, 500, 502})
+
+
+@dataclass(frozen=True)
+class RetryPolicy:
+    """Retry rules from the Polza docs.
+
+    408/500/502: retry after a pause. 503: honor Retry-After, else back off.
+    429: wait Retry-After (own limit) or back off (provider limit). 402 is retried only when
+    `retry_402` is set, because Polza gives no machine-readable way to tell a parallel-load
+    reserve conflict from an empty balance. 400/401/403/404/409/413 are never retried.
+    """
+
+    max_retries: int = 3
+    base_delay: float = 1.0
+    max_delay: float = 30.0
+    max_wait: float = 60.0
+    jitter: float = 0.25
+    retry_402: bool = False
+    retry_402_attempts: int = 2
+
+    def delay_for(
+        self, error: PolzaError, attempt: int, rng: random.Random | None = None
+    ) -> float | None:
+        """Seconds to sleep before retry number `attempt` (0-based), or None to give up."""
+        if attempt >= self.max_retries:
+            return None
+        status = error.status
+        if status == 402:
+            if not self.retry_402 or attempt >= self.retry_402_attempts:
+                return None
+            return self._jittered(self.base_delay, rng)
+        if isinstance(error, RateLimitError) and error.retry_after is not None:
+            return error.retry_after if error.retry_after <= self.max_wait else None
+        if status == 503 and error.retry_after is not None:
+            return error.retry_after
+        if status in _BACKOFF_STATUSES or status in (429, 503):
+            return self._jittered(min(self.base_delay * 2**attempt, self.max_delay), rng)
+        return None
+
+    def _jittered(self, delay: float, rng: random.Random | None) -> float:
+        if not self.jitter:
+            return delay
+        return delay * (1 + (rng or random).uniform(-self.jitter, self.jitter))

polza_extras-0.2.0.dist-info/METADATA

@@ -0,0 +1,189 @@
+Metadata-Version: 2.5
+Name: polza-extras
+Version: 0.2.0
+Summary: Polza.ai helpers on top of the OpenAI SDK: ruble cost, typed errors, retry rules, model catalog, balance
+Project-URL: Homepage, https://github.com/still-coding/polza-extras
+Project-URL: Issues, https://github.com/still-coding/polza-extras/issues
+Project-URL: Changelog, https://github.com/still-coding/polza-extras/blob/main/CHANGELOG.md
+Author: still-coding
+License-Expression: MIT
+License-File: LICENSE
+Keywords: llm,openai,polza,polza.ai,rubles,russia
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Typing :: Typed
+Requires-Python: >=3.11
+Requires-Dist: httpx>=0.27
+Requires-Dist: openai>=1.50
+Requires-Dist: pydantic>=2.7
+Description-Content-Type: text/markdown
+
+# polza-extras
+
+Small helpers for [Polza.ai](https://polza.ai) on top of the official `openai` SDK.
+
+Polza is OpenAI-compatible, so for chat you can use the `openai` SDK directly. This library covers what the SDK doesn't know about:
+
+- **Ruble cost** of each call (`usage.cost_rub`) as `Decimal`.
+- **Typed errors** parsed from Polza's error body, with `trace_id` and the provider's real reason (`metadata.raw`), which is often hidden behind a generic «Некорректный запрос».
+- **Polza's retry rules** (which statuses to retry, `Retry-After`, own vs provider rate limits).
+- **Reasoning models**: a clear `EmptyResponseError` when thinking used up `max_tokens`, instead of an empty string.
+- **Tool calls, usage details, structured output**: parsed `tool_calls`, cached/reasoning token counts, `json_schema=` with a dict or a pydantic model.
+- **Catalog, balance and key info**: `models()`, `model()`, `balance()`, `key_info()`, with ruble prices as `Decimal` and a cost estimate.
+
+## Install
+
+```bash
+uv add polza-extras     # or: pip install polza-extras
+```
+
+Python 3.11+. Dependencies: `openai`, `pydantic`.
+
+## Quickstart
+
+```python
+from polza_extras import Polza
+
+polza = Polza()  # reads POLZA_AI_API_KEY; base URL https://polza.ai/api/v1
+
+result = polza.chat(
+    "anthropic/claude-haiku-5.5",          # always provider/model; a bare name is a 404
+    [{"role": "user", "content": "Привет!"}],
+    max_tokens=200,
+    reasoning={"type": "disabled"},        # any extra Polza field goes into the request body
+)
+print(result.text, result.cost_rub)       # 'Привет! …' Decimal('0.0123')
+
+data = polza.chat(model, messages, json=True)   # response_format={"type": "json_object"}
+
+polza.openai                               # the configured openai.OpenAI client, for anything else
+polza.close()                              # or: with Polza() as polza: ...
+```
+
+### Tool calls
+
+```python
+result = polza.chat(model, messages, tools=[weather_tool])   # tools pass through as is
+for call in result.tool_calls:                # [] for a plain text reply
+    call.id, call.name, call.arguments        # arguments: the JSON string from the model
+    args = call.parsed_arguments()            # dict; ValueError if not a JSON object
+```
+
+A reply with tool calls usually has no text: `result.text == ""` and `finish_reason == "tool_calls"`. `EmptyResponseError` is raised only when there is neither text nor tool calls.
+
+In a stream, `event.tool_calls` holds the raw `ToolCallDelta` pieces: the first piece of a call has `index`, `id` and `name`; join the `arguments` fragments of one `index` to get the full JSON.
+
+### Usage details
+
+```python
+u = result.usage
+u.prompt_tokens, u.completion_tokens, u.total_tokens
+u.cached_tokens        # prompt tokens read from the prompt cache
+u.cache_write_tokens   # prompt tokens written to the prompt cache
+u.reasoning_tokens     # completion tokens spent on thinking
+```
+
+They are 0 when Polza does not report them. The last stream event carries the same `usage`.
+
+### Structured output
+
+```python
+from pydantic import BaseModel
+
+class City(BaseModel):
+    name: str
+    population_millions: float
+
+result = polza.chat(model, messages, json_schema=City, schema_name="city")
+city = City.model_validate_json(result.text)
+```
+
+`json_schema` takes a pydantic model class or a JSON Schema dict and sends `response_format={"type": "json_schema", "json_schema": {"name": schema_name, "schema": ..., "strict": True}}`. Not every provider supports it: check `json_schema_strict` in the model's provider capabilities. Strict mode on some providers also wants `"additionalProperties": false` and every property in `required`; pass a dict when you need full control. Use either `json=True` (any JSON object) or `json_schema`, not both.
+
+### Streaming
+
+```python
+for event in polza.chat_stream("anthropic/claude-haiku-5.5", messages, max_tokens=200):
+    print(event.delta, end="", flush=True)   # event.reasoning holds thinking deltas, if any
+print(event.cost_rub)                        # the last event carries usage and cost
+```
+
+Errors before the first chunk are retried like `chat()`. An error event inside the stream raises `StreamError`; text already yielded stays yielded, and the request is not retried. Stopping early (`break`) closes the connection.
+
+### Async
+
+`AsyncPolza` has the same methods, awaited (`async for` for `chat_stream`):
+
+```python
+from polza_extras import AsyncPolza
+
+polza = AsyncPolza()
+result = await polza.chat("anthropic/claude-haiku-5.5", messages)
+async for event in polza.chat_stream("anthropic/claude-haiku-5.5", messages):
+    ...
+```
+
+### Catalog, balance, key
+
+```python
+page = polza.models(search="claude", type="chat", limit=10)   # one catalog page
+for m in page.items:
+    print(m.id, m.pricing.prompt_per_million, m.reasoning.mandatory)
+
+sonnet = polza.model("anthropic/claude-sonnet-5.5")            # NotFoundError for unknown ids
+sonnet.pricing.estimate(prompt_tokens=1500, completion_tokens=80)  # Decimal rubles
+
+polza.balance()    # Balance(amount, available, reserved, spent, updated_at)
+polza.key_info()   # KeyInfo(label, limit, limit_remaining, limit_reset, usage, usage_daily, ...)
+```
+
+## Errors
+
+All errors subclass `PolzaError` (`status`, `code`, `message`, `trace_id`, `metadata`, `details`, `retry_after`).
+
+| Status | Class | Retried by default |
+|---|---|---|
+| 401 | `AuthError` | no |
+| 402 | `InsufficientFundsError` | no (opt-in `retry_402` for parallel-load reserve conflicts) |
+| 403 | `ForbiddenError` | no |
+| 404 | `NotFoundError` | no |
+| 400, 413, other 4xx | `RequestError` | no |
+| 429 with `Retry-After` | `RateLimitError(scope="own")` | yes, waits `Retry-After` (fails fast if above `max_wait`) |
+| 429 without `Retry-After` | `RateLimitError(scope="provider")` | yes, backoff 1, 2, 4, 8 s |
+| 408, 502 | `ServerError` | yes, backoff |
+| 503 | `ServerError` | yes, honors `Retry-After` |
+| connection failure | `PolzaConnectionError` | yes, backoff |
+| 200 with no text and no tool calls | `EmptyResponseError` | no; carries `finish_reason`, `reasoning`, `usage`, `cost_rub` (still billed), `raw` |
+| error event inside a stream | `StreamError` | no |
+
+Polza doesn't charge for failed requests, so retrying is safe. Tune with `Polza(retry=RetryPolicy(max_retries=..., base_delay=..., max_wait=...))`.
+
+### Reasoning models
+
+Thinking tokens count against `max_tokens`. If the model spends them all on thinking, `content` comes back empty and `chat()` raises `EmptyResponseError` with `finish_reason="length"`. Fixes:
+
+- turn thinking off: `reasoning={"type": "disabled"}` (where allowed),
+- lower it: `reasoning={"effort": "low"}`,
+- raise `max_tokens`.
+
+Some models (e.g. `anthropic/claude-sonnet-5.5`) make thinking mandatory; `model(id).reasoning.mandatory` tells you.
+
+## Development
+
+```bash
+uv sync
+uv run pytest                 # offline, against an in-process fake server; fails below 95% coverage
+POLZA_AI_API_KEY=... uv run pytest -m live --no-cov   # real API smoke tests, costs ~0.1 ₽
+uv run ruff check src tests && uv run ruff format --check src tests && uv run pyright
+ln -sf ../../scripts/pre-commit .git/hooks/pre-commit   # all of the above before each commit
+```
+
+## License
+
+MIT

polza_extras-0.2.0.dist-info/RECORD

@@ -0,0 +1,12 @@
+polza_extras/__init__.py,sha256=8aBkrSncshcSeRtukXNK1XZ3nY8CSdImMPgIrUQ14gk,1039
+polza_extras/_json.py,sha256=3eQ4IKBxqWPvo1bqUSpaRLa6tVJ_BzgGCtPf9jivx4k,388
+polza_extras/_models.py,sha256=YRYyZqsoM6blXHZcSN7BP8C03JI8IGictorRhPatY_U,5180
+polza_extras/chat.py,sha256=zZU0l_14veij84D9m9Q95JAiy_3anJIafGv35Lf5doc,4213
+polza_extras/client.py,sha256=TBZNAs24xwB4c6FiLfWd_SXMHb05cfi0oIW832DKqhs,15412
+polza_extras/errors.py,sha256=v34B1DTDOFzHhzpMxOTYlpoxQCEbZqvB5oOsf38__No,4868
+polza_extras/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+polza_extras/retry.py,sha256=wUQ20oYWsNUtjWOUF5-eZJoiYYqFyvuXyTDVxIJaH4I,1957
+polza_extras-0.2.0.dist-info/METADATA,sha256=dlXHA2cN_hTOm4D23cgh_vHk8f70VD71mbP7qHcMBQg,8325
+polza_extras-0.2.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+polza_extras-0.2.0.dist-info/licenses/LICENSE,sha256=DsRhA3-G0ezNENtBc0WESVItyrPiPfi4RDMfldHRsGw,1082
+polza_extras-0.2.0.dist-info/RECORD,,

polza_extras-0.2.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

polza_extras-0.2.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 polza-extras contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.