taskiq-llm-queue 0.1.0a1__py3-none-any.whl

taskiq_llm_queue/http_transport.py

@@ -0,0 +1,220 @@
+"""httpx transport that turns OpenAI SDK HTTP requests into LLM queue calls (D88).
+
+Given to `openai.AsyncOpenAI(http_client=httpx.AsyncClient(transport=...))`, it answers
+`/chat/completions` with `llmq.chat` and `/embeddings` with `llmq.embed`, so the SDK and
+LangChain parse replies exactly as from a real server. Queue errors become OpenAI-style error
+bodies with a matching status, so the SDK raises its usual exceptions (429 → `RateLimitError`).
+Chat replies carry how the queue served them under the extra body key `llmq_meta`.
+Needs the `openai` extra (httpx); the core package never imports this module.
+"""
+
+from __future__ import annotations
+
+import array
+import base64
+from typing import TYPE_CHECKING, Any
+
+import httpx
+import orjson
+
+from taskiq_llm_queue.errors import (
+    ConfigError,
+    LLMQueueCancelled,
+    LLMQueueError,
+    LLMQueueTimeout,
+    ProviderError,
+    QueueUnavailable,
+)
+from taskiq_llm_queue.langchain._client import resolve_queue
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.config import Lane
+    from taskiq_llm_queue.envelope import FallbackPolicy
+    from taskiq_llm_queue.langchain._client import ChatResultLike, QueueClient
+
+# Never resolved: every request is answered in-process by the transport.
+QUEUED_BASE_URL = "http://llmq.invalid/v1"
+# Extra key of a queued chat completion body: the call's `CallMeta` (served model, attempts...).
+META_KEY = "llmq_meta"
+
+
+class _RequestRejected(Exception):
+    """A request the transport cannot queue (bad body, streaming, unknown route)."""
+
+    def __init__(self, status: int, message: str, code: str) -> None:
+        super().__init__(message)
+        self.status = status
+        self.code = code
+
+
+class QueuedHTTPTransport(httpx.AsyncBaseTransport, httpx.BaseTransport):
+    """Answer OpenAI-compatible HTTP requests from the LLM queue instead of the network.
+
+    `model_alias` overrides the body's `model` (use it when the SDK model name differs from the
+    queue alias). `llmq=None` uses the default queue, looked up per request. Auth headers are
+    ignored: workers hold the provider credentials. The sync side needs `chat_sync` on the client
+    for chat (embeddings use `embed_sync`); without it the SDK gets a 501 error.
+    """
+
+    def __init__(
+        self,
+        llmq: QueueClient | None = None,
+        model_alias: str | None = None,
+        *,
+        lane: Lane | str | None = None,
+        fallback: FallbackPolicy | None = None,
+        timeout_secs: float | None = None,
+    ) -> None:
+        """Keep the queue, alias and per-call overrides; nothing is resolved until a request."""
+        self._llmq = llmq
+        self._alias = model_alias
+        self._lane = lane
+        self._fallback = fallback
+        self._timeout_secs = timeout_secs
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        """Route one SDK request through the async queue API."""
+        try:
+            route, alias, body = self._parse(request, await request.aread())
+            queue = resolve_queue(self._llmq)
+            if route == "chat":
+                chat = await queue.chat(
+                    alias,
+                    body,
+                    lane=self._lane,
+                    fallback=self._fallback,
+                    timeout_secs=self._timeout_secs,
+                )
+                data = _chat_response(chat)
+            else:
+                texts = _embedding_inputs(body)
+                embedded = await queue.embed(
+                    alias, texts, lane=self._lane, timeout_secs=self._timeout_secs
+                )
+                data = _embedding_response(alias, embedded.vectors, body.get("encoding_format"))
+        except (_RequestRejected, LLMQueueError) as exc:
+            return _error_response(request, exc)
+        return httpx.Response(200, json=data, request=request)
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        """Route one SDK request through the sync queue API (`chat_sync` / `embed_sync`)."""
+        try:
+            route, alias, body = self._parse(request, request.read())
+            queue = resolve_queue(self._llmq)
+            if route == "chat":
+                chat_sync = getattr(queue, "chat_sync", None)
+                if chat_sync is None:
+                    raise _RequestRejected(
+                        501, "sync chat is not supported by this queue; use ainvoke", "sync"
+                    )
+                chat = chat_sync(
+                    alias,
+                    body,
+                    lane=self._lane,
+                    fallback=self._fallback,
+                    timeout_secs=self._timeout_secs,
+                )
+                data = _chat_response(chat)
+            else:
+                texts = _embedding_inputs(body)
+                embedded = queue.embed_sync(
+                    alias, texts, lane=self._lane, timeout_secs=self._timeout_secs
+                )
+                data = _embedding_response(alias, embedded.vectors, body.get("encoding_format"))
+        except (_RequestRejected, LLMQueueError) as exc:
+            return _error_response(request, exc)
+        return httpx.Response(200, json=data, request=request)
+
+    def _parse(self, request: httpx.Request, raw: bytes) -> tuple[str, str, dict[str, Any]]:
+        """Return (route, alias, JSON body), rejecting what the queue cannot serve."""
+        path = request.url.path.rstrip("/")
+        if request.method != "POST":
+            raise _RequestRejected(405, f"{request.method} {path} is not queued", "method")
+        if path.endswith("/chat/completions"):
+            route = "chat"
+        elif path.endswith("/embeddings"):
+            route = "embeddings"
+        else:
+            raise _RequestRejected(404, f"route {path} is not supported by the queue", "route")
+        try:
+            body = orjson.loads(raw or b"{}")
+        except ValueError as exc:
+            raise _RequestRejected(400, f"request body is not JSON: {exc}", "body") from exc
+        if not isinstance(body, dict):
+            raise _RequestRejected(400, "request body must be a JSON object", "body")
+        if body.get("stream"):
+            raise _RequestRejected(400, "streaming is not supported through the queue", "stream")
+        alias = self._alias or body.get("model")
+        if not isinstance(alias, str) or not alias:
+            raise _RequestRejected(400, "no model alias in the transport or the body", "model")
+        return route, alias, body
+
+
+def _chat_response(result: ChatResultLike) -> dict[str, Any]:
+    """The completion body plus `llmq_meta` (the SDK keeps unknown keys as extras)."""
+    return {**result.response, META_KEY: result.meta.model_dump(mode="json")}
+
+
+def _embedding_inputs(body: dict[str, Any]) -> list[str]:
+    """Texts from an embeddings body; token-id inputs are rejected (queue embeds text only)."""
+    value = body.get("input")
+    if isinstance(value, str):
+        return [value]
+    if isinstance(value, list) and all(isinstance(v, str) for v in value):
+        return value
+    raise _RequestRejected(400, "embeddings input must be a string or a list of strings", "input")
+
+
+def _embedding_response(
+    alias: str, vectors: list[list[float]], encoding_format: object
+) -> dict[str, Any]:
+    """OpenAI embeddings body; base64 is float32 in native order, as the SDK decodes it."""
+    as_base64 = encoding_format == "base64"
+    return {
+        "object": "list",
+        "model": alias,
+        "data": [
+            {
+                "object": "embedding",
+                "index": i,
+                "embedding": (
+                    base64.b64encode(array.array("f", vec).tobytes()).decode("ascii")
+                    if as_base64
+                    else vec
+                ),
+            }
+            for i, vec in enumerate(vectors)
+        ],
+        "usage": {"prompt_tokens": 0, "total_tokens": 0},
+    }
+
+
+def _error_response(request: httpx.Request, exc: Exception) -> httpx.Response:
+    """Map a queue error to the status and body an OpenAI server would send.
+
+    `ProviderError` uses its `status_code` when the client sets one, else 429 when retryable
+    (the SDK raises `RateLimitError`) and 400 when not; timeouts are 504, queue outages 503.
+    """
+    headers = {"x-llmq-error": type(exc).__name__}
+    if isinstance(exc, _RequestRejected):
+        status, err_type, code = exc.status, "invalid_request_error", exc.code
+    elif isinstance(exc, ProviderError):
+        explicit = getattr(exc, "status_code", None)
+        status = explicit if isinstance(explicit, int) else (429 if exc.retryable else 400)
+        err_type = "rate_limit_error" if status == 429 else "provider_error"
+        code = "provider_error"
+        if exc.retry_after_secs is not None:
+            headers["retry-after"] = f"{exc.retry_after_secs:g}"
+            headers["retry-after-ms"] = str(int(exc.retry_after_secs * 1000))
+    elif isinstance(exc, LLMQueueCancelled):
+        status, err_type, code = 499, "cancelled", "cancelled"
+    elif isinstance(exc, LLMQueueTimeout):
+        status, err_type, code = 504, "timeout", "queue_timeout"
+    elif isinstance(exc, QueueUnavailable):
+        status, err_type, code = 503, "service_unavailable", "queue_unavailable"
+    elif isinstance(exc, ConfigError):
+        status, err_type, code = 404, "invalid_request_error", "model_not_found"
+    else:
+        status, err_type, code = 500, "server_error", "llmq_error"
+    body = {"error": {"message": str(exc), "type": err_type, "param": None, "code": code}}
+    return httpx.Response(status, json=body, headers=headers, request=request)

taskiq_llm_queue/keys.py

@@ -0,0 +1,87 @@
+"""Every Redis/Valkey key and channel name, built from one configurable prefix.
+
+Keeping names in one place lets several environments share a Redis without collisions and
+keeps the state schema (versioned under `schema`) easy to migrate. Keys that one Lua script
+touches together share a hash tag (`{...}`), so they map to one slot on cluster-mode Valkey.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+# 2: hash tags for cluster mode (`{m:<model>:<target>}`, `{r:<request_id>}`, `{cfg}`).
+SCHEMA_VERSION = 2
+
+
+@dataclass(frozen=True, slots=True)
+class KeySpace:
+    """Key builder for one prefix, e.g. `llmq:prod` -> `llmq:prod:{r:<id>}:job`.
+
+    Tags group the keys of one script: `{r:<request_id>}` (claim, cancel, params, result),
+    `{m:<model>:<target>}` (limiter state, leases, breaker) and `{cfg}` (live config, audit).
+    """
+
+    prefix: str
+
+    def _request(self, request_id: str, name: str) -> str:
+        return f"{self.prefix}:{{r:{request_id}}}:{name}"
+
+    def _target(self, model: str, target: str, name: str) -> str:
+        return f"{self.prefix}:{{m:{model}:{target}}}:{name}"
+
+    def job(self, request_id: str) -> str:
+        """Claim state of one task (hash): state, worker, heartbeat, failures."""
+        return self._request(request_id, "job")
+
+    def params(self, request_id: str) -> str:
+        """Large params stored by reference (string)."""
+        return self._request(request_id, "params")
+
+    def cancel(self, request_id: str) -> str:
+        """Cancel flag for one task (string)."""
+        return self._request(request_id, "cancel")
+
+    def result(self, task_id: str) -> str:
+        """JSON result of one task, deleted on read (string); same slot as its claim."""
+        return self._request(task_id, "res")
+
+    def limiter(self, model: str, target: str) -> str:
+        """Adaptive limiter state for one model target (hash, no TTL)."""
+        return self._target(model, target, "ac")
+
+    def leases(self, model: str, target: str) -> str:
+        """In-flight slot leases scored by expiry (sorted set)."""
+        return self._target(model, target, "lease")
+
+    def breaker(self, model: str, target: str) -> str:
+        """Circuit breaker state for one model target (hash, no TTL)."""
+        return self._target(model, target, "cb")
+
+    def done_channel(self, pod: str) -> str:
+        """Pub/sub channel that wakes one pod's result dispatcher."""
+        return f"{self.prefix}:done:{pod}"
+
+    @property
+    def routing(self) -> str:
+        """Live canary routing config (JSON string)."""
+        return f"{self.prefix}:{{cfg}}:routing"
+
+    @property
+    def fallback(self) -> str:
+        """Live fallback chains (JSON string)."""
+        return f"{self.prefix}:{{cfg}}:fallback"
+
+    @property
+    def audit(self) -> str:
+        """Audit stream of config changes (who, before, after)."""
+        return f"{self.prefix}:{{cfg}}:audit"
+
+    @property
+    def pods(self) -> str:
+        """Live pods of this prefix (sorted set scored by expiry)."""
+        return f"{self.prefix}:pods"
+
+    @property
+    def schema(self) -> str:
+        """State schema version marker."""
+        return f"{self.prefix}:schema"

taskiq_llm_queue/langchain/__init__.py

@@ -0,0 +1,64 @@
+"""LangChain adapters that send model calls through the LLM queue (extras `langchain-*`).
+
+The adapter classes are loaded on first attribute access, so importing this package (for
+example to call `set_default_queue` from `LLMQueue.start()`) never imports langchain, boto3 or
+httpx. `QueuedChatBedrockConverse` / `QueuedBedrockEmbeddings` need `[langchain-aws]`;
+`QueuedChatOpenAI` / `QueuedOpenAIEmbeddings` need `[openai,langchain-openai]`.
+"""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import TYPE_CHECKING, Any
+
+from taskiq_llm_queue.langchain._client import (
+    ChatResultLike,
+    ConverseResultLike,
+    EmbedResultLike,
+    QueueClient,
+    get_default_queue,
+    set_default_queue,
+)
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.langchain.bedrock import (
+        QueuedBedrockEmbeddings,
+        QueuedChatBedrockConverse,
+    )
+    from taskiq_llm_queue.langchain.openai import QueuedChatOpenAI, QueuedOpenAIEmbeddings
+
+# Lazy exports: attribute name -> (module, extra that provides its dependencies).
+_LAZY = {
+    "QueuedChatBedrockConverse": ("bedrock", "langchain-aws"),
+    "QueuedBedrockEmbeddings": ("bedrock", "langchain-aws"),
+    "QueuedChatOpenAI": ("openai", "openai,langchain-openai"),
+    "QueuedOpenAIEmbeddings": ("openai", "openai,langchain-openai"),
+}
+
+__all__ = [
+    "ChatResultLike",
+    "ConverseResultLike",
+    "EmbedResultLike",
+    "QueueClient",
+    "QueuedBedrockEmbeddings",
+    "QueuedChatBedrockConverse",
+    "QueuedChatOpenAI",
+    "QueuedOpenAIEmbeddings",
+    "get_default_queue",
+    "set_default_queue",
+]
+
+
+def __getattr__(name: str) -> Any:
+    """Import an adapter's module on first use, naming the extra to install if it is missing."""
+    if name not in _LAZY:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+    module, extra = _LAZY[name]
+    try:
+        value = getattr(import_module(f"{__name__}.{module}"), name)
+    except ImportError as exc:
+        raise ImportError(
+            f"{name} needs the optional dependencies: uv add 'taskiq-llm-queue[{extra}]'"
+        ) from exc
+    globals()[name] = value
+    return value

taskiq_llm_queue/langchain/_client.py

@@ -0,0 +1,142 @@
+"""The client interface the adapters call, and the process-wide default queue (D85).
+
+Adapters type against `QueueClient` instead of importing `LLMQueue`, so they work with any object
+that has the same methods (tests use a fake). Classes created at import time find the running
+queue on first call through `get_default_queue`. This module never imports langchain or httpx.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any, Protocol, cast, runtime_checkable
+
+from taskiq_llm_queue.errors import LLMQueueError
+
+if TYPE_CHECKING:
+    from taskiq_llm_queue.config import Lane
+    from taskiq_llm_queue.envelope import CallMeta, FallbackPolicy
+
+
+class ConverseResultLike(Protocol):
+    """What `converse` returns: the exact Bedrock Converse dict and how it was served."""
+
+    @property
+    def response(self) -> dict[str, Any]: ...
+
+    @property
+    def meta(self) -> CallMeta: ...
+
+
+class ChatResultLike(Protocol):
+    """What `chat` returns: the OpenAI-compatible completion body and how it was served."""
+
+    @property
+    def response(self) -> dict[str, Any]: ...
+
+    @property
+    def meta(self) -> CallMeta: ...
+
+
+class EmbedResultLike(Protocol):
+    """What `embed` returns: one vector per text, in input order, and how they were served."""
+
+    @property
+    def vectors(self) -> list[list[float]]: ...
+
+    @property
+    def meta(self) -> CallMeta: ...
+
+
+@runtime_checkable
+class QueueClient(Protocol):
+    """The subset of `LLMQueue` the LangChain adapters and the HTTP transport use.
+
+    `lane=None` means "take the lane from `current_context()`"; adapters pass only explicit
+    overrides. Sync methods run on the caller's thread (direct path under the shared limiter).
+    """
+
+    async def converse(
+        self,
+        model: str,
+        request: dict[str, Any],
+        *,
+        lane: Lane | str | None = None,
+        fallback: FallbackPolicy | None = None,
+        timeout_secs: float | None = None,
+        deadline_at: float | None = None,
+        request_id: str | None = None,
+    ) -> ConverseResultLike: ...
+
+    def converse_sync(
+        self,
+        model: str,
+        request: dict[str, Any],
+        *,
+        lane: Lane | str | None = None,
+        fallback: FallbackPolicy | None = None,
+        timeout_secs: float | None = None,
+        deadline_at: float | None = None,
+        request_id: str | None = None,
+    ) -> ConverseResultLike: ...
+
+    async def embed(
+        self,
+        model: str,
+        texts: list[str],
+        *,
+        lane: Lane | str | None = None,
+        input_type: str | None = None,
+        timeout_secs: float | None = None,
+        deadline_at: float | None = None,
+        request_id: str | None = None,
+    ) -> EmbedResultLike: ...
+
+    def embed_sync(
+        self,
+        model: str,
+        texts: list[str],
+        *,
+        lane: Lane | str | None = None,
+        input_type: str | None = None,
+        timeout_secs: float | None = None,
+        deadline_at: float | None = None,
+        request_id: str | None = None,
+    ) -> EmbedResultLike: ...
+
+    async def chat(
+        self,
+        model: str,
+        request: dict[str, Any],
+        *,
+        lane: Lane | str | None = None,
+        fallback: FallbackPolicy | None = None,
+        timeout_secs: float | None = None,
+        deadline_at: float | None = None,
+        request_id: str | None = None,
+    ) -> ChatResultLike: ...
+
+
+_default_queue: QueueClient | None = None
+
+
+def set_default_queue(llmq: QueueClient | None) -> None:
+    """Register the queue that adapters built without `llmq=` use; `None` clears it.
+
+    `LLMQueue.start()` is expected to call this, and `stop()` to clear it.
+    """
+    global _default_queue
+    _default_queue = llmq
+
+
+def get_default_queue() -> QueueClient:
+    """The registered default queue; raises `LLMQueueError` when none has been set yet."""
+    if _default_queue is None:
+        raise LLMQueueError(
+            "no LLM queue: pass llmq=... to the adapter or call "
+            "taskiq_llm_queue.langchain.set_default_queue(llmq) (LLMQueue.start() does this)"
+        )
+    return _default_queue
+
+
+def resolve_queue(explicit: Any) -> QueueClient:
+    """The adapter's own queue if it has one, else the default, looked up on every call."""
+    return cast("QueueClient", explicit) if explicit is not None else get_default_queue()