machgen-client 0.1.0__py3-none-any.whl

machgen/client/__init__.py

@@ -0,0 +1,36 @@
+from machgen.client._models import (
+    GenerateResponse,
+    ModerationResult,
+    ModerationStage,
+    TaskMetadata,
+    TaskOutputType,
+    TaskStatusResponse,
+    UploadResponse,
+)
+from machgen.client.api import (
+    ImageConfig,
+    TaskInput,
+    TaskStatus,
+    TaskUpdate,
+    VideoConfig,
+)
+from machgen.client.client import MachGenClient, SseRetryConfig
+from machgen.client.task_handle import TaskHandle
+
+__all__ = [
+    "GenerateResponse",
+    "ImageConfig",
+    "MachGenClient",
+    "ModerationResult",
+    "ModerationStage",
+    "SseRetryConfig",
+    "TaskHandle",
+    "TaskInput",
+    "TaskMetadata",
+    "TaskOutputType",
+    "TaskStatus",
+    "TaskStatusResponse",
+    "TaskUpdate",
+    "UploadResponse",
+    "VideoConfig",
+]

machgen/client/_models.py

@@ -0,0 +1,138 @@
+"""Wire-level response models shared between client and server.
+
+The server ([frontend/api_server.py](frontend/api_server.py)) and the external
+client both import these so the JSON schema stays in lock-step.
+
+**Schema evolution policy.** Every model here sets ``extra="ignore"`` so an
+older peer (client *or* server) silently drops fields it doesn't know about.
+That means new fields can be added without coordinating a release across both
+sides — but it requires:
+
+- New fields are always **optional** (carry a default).
+- Existing required fields are **never removed or renamed**.
+- Enum members are **never removed** (additions are tolerated by the receiver
+  only if the receiver doesn't dispatch on the enum value).
+
+These rules are enforced by reviewer discipline, not by code; the test
+[tests/api/test_wire_compat.py](tests/api/test_wire_compat.py) pins the
+``extra="ignore"`` behavior so a regression flips a red light.
+"""
+
+from __future__ import annotations
+
+from enum import StrEnum
+
+from pydantic import BaseModel, ConfigDict, Field
+
+from machgen.client.api import TaskStatus
+
+_WIRE_MODEL_CONFIG = ConfigDict(extra="ignore")
+
+
+class TaskMetadata(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    prompt: str = Field(description="Prompt the task was submitted with.")
+    model: str | None = Field(default=None, description="Resolved model id.")
+    task_type: str | None = Field(default=None, description="Resolved task type.")
+    seed: int | None = Field(default=None, description="Seed used for generation.")
+    fps: int | None = Field(default=None, description="Frames per second (video).")
+    height: int | None = Field(default=None, description="Output height in pixels.")
+    width: int | None = Field(default=None, description="Output width in pixels.")
+    aspect_ratio: str | None = Field(default=None, description="Output aspect ratio.")
+    duration_secs: int | None = Field(
+        default=None, description="Clip duration in seconds (video)."
+    )
+    src_image_urls: list[str] | None = Field(
+        default=None, description="Source / reference image refs."
+    )
+    src_video_url: str | None = Field(
+        default=None, description="Source video ref (video-to-video)."
+    )
+    acceleration: str | None = Field(
+        default=None, description="Acceleration profile applied by the backend."
+    )
+
+
+class TaskOutputType(StrEnum):
+    """A task may have 1 or more types of output."""
+
+    TEXT = "text"
+    VIDEO = "video"
+    IMAGE = "image"
+    AUDIO = "audio"
+
+
+class ModerationStage(StrEnum):
+    INPUT = "input"
+    OUTPUT = "output"
+
+
+class ModerationResult(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    flagged: bool = Field(
+        default=False, description="Whether moderation matched a block category."
+    )
+    stage: ModerationStage | None = Field(
+        default=None, description="Which stage flagged (input or output)."
+    )
+    surface: str | None = Field(
+        default=None, description='Surface that tripped: "text" or "image".'
+    )
+    categories: list[str] = Field(
+        default_factory=list, description="Matched block categories."
+    )
+    source: str | None = Field(
+        default=None,
+        description='Who produced the verdict: "machgen" (OpenAI) or "vendor".',
+    )
+    moderation_time_secs: float | None = Field(
+        default=None, description="Wall-clock time spent in the moderation call(s)."
+    )
+
+
+class GenerateResponse(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    task_id: str = Field(
+        description="Server-assigned id; poll status and download the asset with it."
+    )
+    metadata: TaskMetadata = Field(description="Echo of the resolved task parameters.")
+    task_output: dict[TaskOutputType, str] | None = Field(
+        default=None,
+        description="Per output-type asset locators, as they become available.",
+    )
+    moderation: ModerationResult | None = Field(
+        default=None,
+        description="Set when moderation rejected the request at admission. The task was admitted then immediately failed (refunded); no generation work was done.",
+    )
+
+
+class TaskStatusResponse(GenerateResponse):
+    status: TaskStatus = Field(description="Current task status.")
+    error_msg: str | None = Field(
+        default=None, description="Failure reason, set when status is FAILED."
+    )
+    generation_time_secs: float | None = Field(
+        default=None, description="Time spent on actual generation."
+    )
+    upload_time_secs: float | None = Field(
+        default=None, description="Time spent uploading the asset after generation."
+    )
+    queue_time_secs: float | None = Field(
+        default=None,
+        description="Time spent waiting in the queue for generation to begin.",
+    )
+    moderation: ModerationResult | None = Field(
+        default=None,
+        description="Structured moderation verdict when the task failed or was blocked by moderation.",
+    )
+
+
+class UploadResponse(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    artifact_path: str = Field(
+        description="Storage path of the uploaded input; reference it as @input/<artifact_path>."
+    )

machgen/client/_sse_executor.py

@@ -0,0 +1,82 @@
+"""Process-global thread pool for ``TaskHandle`` SSE consumers.
+
+All handles across all ``MachGenClient`` instances submit their stream
+consumption work to one ``ThreadPoolExecutor`` so the per-process worker
+count stays bounded — the alternative (one ad-hoc daemon thread per
+handle) blows up when callers spin up many clients with many handles.
+
+Cap: ``MACHGEN_CLIENT_MAX_SSE_WORKERS`` env var (default 256). The pool
+grows lazily up to that cap.
+
+Atexit handling: at program shutdown we signal stop on every live stream
+state before the pool's own atexit hook fires, so workers blocked on
+``iter_sse()`` are unstuck and the pool's join doesn't hang. Atexit is
+LIFO and the pool's hook is registered the first time a
+``ThreadPoolExecutor`` is created — we register ours after that, so ours
+runs first.
+"""
+
+from __future__ import annotations
+
+import atexit
+import logging
+import os
+import threading
+import weakref
+from concurrent.futures import Future, ThreadPoolExecutor
+from typing import TYPE_CHECKING, Callable
+
+if TYPE_CHECKING:
+    from machgen.client.task_handle import _StreamState
+
+logger = logging.getLogger(__name__)
+
+_DEFAULT_MAX_WORKERS = int(os.environ.get("MACHGEN_CLIENT_MAX_SSE_WORKERS", "256"))
+
+_executor: ThreadPoolExecutor | None = None
+_executor_lock = threading.Lock()
+_active_states: "weakref.WeakSet[_StreamState]" = weakref.WeakSet()
+_active_lock = threading.Lock()
+
+
+def _ensure_executor() -> ThreadPoolExecutor:
+    global _executor
+    if _executor is not None:
+        return _executor
+    with _executor_lock:
+        if _executor is None:
+            _executor = ThreadPoolExecutor(
+                max_workers=_DEFAULT_MAX_WORKERS,
+                thread_name_prefix="machgen-sse",
+            )
+            atexit.register(_atexit_shutdown)
+    return _executor
+
+
+def submit(fn: Callable[..., None], *args, **kwargs) -> Future:
+    """Submit an SSE consumer function to the shared pool."""
+    return _ensure_executor().submit(fn, *args, **kwargs)
+
+
+def register_state(state: "_StreamState") -> None:
+    """Track *state* so atexit can stop its worker before pool shutdown."""
+    with _active_lock:
+        _active_states.add(state)
+
+
+def _atexit_shutdown() -> None:
+    """Signal stop on every live state, then shut down the pool.
+
+    Runs before ``concurrent.futures.thread._python_exit`` (LIFO atexit
+    order). ``signal_stop`` wakes any worker blocked on ``iter_sse()``;
+    the pool's subsequent join then proceeds promptly.
+    """
+    with _active_lock:
+        states = list(_active_states)
+    for state in states:
+        try:
+            state.signal_stop(terminal_message="process exiting")
+        except Exception:
+            logger.warning("error stopping SSE worker at exit", exc_info=True)
+    if _executor is not None:
+        _executor.shutdown(wait=False)

machgen/client/api.py

@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+from enum import StrEnum
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+_WIRE_MODEL_CONFIG = ConfigDict(extra="ignore", frozen=True)
+
+
+class TaskStatus(StrEnum):
+    PENDING = "PENDING"
+    RUNNING = "RUNNING"
+    # Generation + upload both done; task_output is populated. Terminal.
+    COMPLETED = "COMPLETED"
+    FAILED = "FAILED"
+
+
+class VideoConfig(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    fps: int | None = Field(
+        default=None,
+        description="Video frames per second. If omitted the default FPS would be used based on the model.",
+    )
+    duration_secs: int = Field(
+        description="Video duration in seconds.",
+    )
+    height: int | None = Field(
+        default=None,
+        description="Output height in pixels",
+    )
+    width: int | None = Field(
+        default=None,
+        description="Output width in pixels.",
+    )
+    aspect_ratio: str | None = Field(
+        default=None,
+        description=(
+            "Output aspect ratio. "
+            "The default is 16:9 if omitted. "
+            "The width of the output will be updated to match the height based on the aspect ratio, "
+            "rounded up to the nearest integer. "
+        ),
+    )
+    infer_steps: int | None = Field(
+        default=None,
+        description=(
+            "Number of denoising / inference steps. Higher values trade more "
+            "compute for potentially finer detail. This is a best-effort match: "
+            "if the model does not support it the model default "
+            "is used. "
+        ),
+    )
+    audio: bool | None = Field(
+        default=None,
+        description=(
+            "Whether the video should include audio. "
+            "**Note:** some models do no support audio, "
+            "or the audio is always on (e.g. Veo 3.1). "
+            "In those cases this field has no effect."
+        ),
+    )
+    guidance_scale: list[float] | None = Field(
+        default=None,
+        description=(
+            "Classifier-free guidance scale(s). Controls how strongly the output "
+            "adheres to the prompt: higher values follow the prompt more closely "
+            "at the cost of diversity. If the model supports multiple guidance "
+            "scales, these will be applied in a sequence (e.g. per stage or per "
+            "denoising phase). This is a best-effort match: if the model does "
+            "not support it, or does not support the number of scales provided, "
+            "the model default is used. "
+        ),
+    )
+
+    # Accept the legacy scalar shape from tasks stored before guidance_scale
+    # became a list, so old rows keep loading after the schema change.
+    @field_validator("guidance_scale", mode="before")
+    @classmethod
+    def _coerce_scalar_guidance_scale(cls, v: object) -> object:
+        if isinstance(v, (int, float)) and not isinstance(v, bool):
+            return [float(v)]
+        return v
+
+
+class ImageConfig(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    height: int | None = Field(
+        default=None,
+        description="Output height in pixels",
+    )
+    width: int | None = Field(
+        default=None,
+        description="Output width in pixels.",
+    )
+    aspect_ratio: str | None = Field(
+        default=None,
+        description=(
+            "Output aspect ratio. "
+            "The default is 1:1 if omitted. "
+            "The width of the output will be updated to match the height based on the aspect ratio, "
+            "rounded up to the nearest integer. "
+        ),
+    )
+    infer_steps: int | None = Field(
+        default=None,
+        description=(
+            "Number of denoising / inference steps. Higher values trade more "
+            "compute for potentially finer detail. This is a best-effort match: "
+            "if the model does not support it the model default "
+            "is used. "
+        ),
+    )
+    guidance_scale: list[float] | None = Field(
+        default=None,
+        description=(
+            "Classifier-free guidance scale(s). Controls how strongly the output "
+            "adheres to the prompt: higher values follow the prompt more closely "
+            "at the cost of diversity. If the model supports multiple guidance "
+            "scales, these will be applied in a sequence (e.g. per stage or per "
+            "denoising phase). This is a best-effort match: if the model does "
+            "not support it, or does not support the number of scales provided, "
+            "the model default is used. "
+        ),
+    )
+
+    # Accept the legacy scalar shape from tasks stored before guidance_scale
+    # became a list, so old rows keep loading after the schema change.
+    @field_validator("guidance_scale", mode="before")
+    @classmethod
+    def _coerce_scalar_guidance_scale(cls, v: object) -> object:
+        if isinstance(v, (int, float)) and not isinstance(v, bool):
+            return [float(v)]
+        return v
+
+
+class TaskUpdate(BaseModel):
+    model_config = _WIRE_MODEL_CONFIG
+
+    status: TaskStatus
+    progress: float | None = None
+
+
+class TaskInput(BaseModel):
+    """
+    Public input model for :meth:`machgen.client.MachGenClient.submit_task`.
+
+    Model/task_type need to match the supported model list.
+    """
+
+    model_config = _WIRE_MODEL_CONFIG
+
+    prompt: str = Field(description="Text prompt driving generation.")
+    multi_prompt: list[str] | None = Field(
+        default=None,
+        description=(
+            "Per-shot prompts for multi-shot passthrough vendors; length must "
+            "match src_image_urls when used as reference frames. "
+            "**This is only supported/needed for Kling-v3 R2V.**"
+        ),
+    )
+    model: str = Field(description="Model id, e.g. 'Wan2.2-T2V-A14B', 'Kling-v3'.")
+    task_type: str = Field(description="one of T2I, I2I, T2V, I2V, R2V")
+    moderate: bool = Field(
+        default=True,
+        description=(
+            "Whether this request is screened by content moderation (the prompt "
+            "and any reference images, plus the generated image). Defaults to "
+            "true and most callers should leave it so. Setting it false is a "
+            "privileged option: it is honored only for accounts explicitly "
+            "allowlisted to bypass moderation; from any other account the "
+            "request is rejected with HTTP 403."
+        ),
+    )
+    video_config: VideoConfig | None = Field(
+        default=None, description="Required for video task types."
+    )
+    image_config: ImageConfig | None = Field(
+        default=None, description="Required for image task types."
+    )
+    seed: int | None = Field(
+        default=None,
+        description="Seed for reproducible generation. If not specified, a random seed will be used.",
+    )
+    src_image_urls: list[str] | None = Field(
+        default=None,
+        description=(
+            "Source / reference image URLs. "
+            "Only needed for tasks that require input images like I2I, I2V, R2V "
+            "Refer to the API docs for concrete examples how to use this and what inputs are allowed. "
+            "For image to video tasks, this can optionally specify 1 (first frame) or 2 (first & last frames) input images. "
+        ),
+    )
+    subject_to_image_ids: dict[str, list[int]] | None = Field(
+        default=None,
+        description=(
+            "R2V only: maps a subject name to the indices of its reference "
+            "images in src_image_urls, e.g. {'alice': [0, 1], 'bob': [2]}. The "
+            "prompt may address a subject via '@name'. Honored by vendors with "
+            "named subjects (Vidu reference2video)."
+        ),
+    )