Move some types around

Author: anbuzin

src/ai/models/__init__.py

@@ -25,15 +25,14 @@
 """
 
 from ..types.stream import StreamResultLike
-from .ai_gateway.types import GenerateParams, ImageParams, VideoParams
 from .core.adapters import register_check, register_generate, register_stream
 from .core.api import check_connection, generate, stream
 from .core.catalog import get_models, get_providers, register_catalog
 from .core.catalog import model as model
 from .core.client import _PROVIDER_DEFAULTS, Client
 from .core.model import Model, ModelCost
 from .core.proto import CheckConnFn, GenerateFn, StreamFn
-from .core.stream_result import StreamResult
+from .core.types import GenerateParams, ImageParams, StreamResult, VideoParams
 
 __all__ = [
     # Core types

src/ai/models/ai_gateway/__init__.py

@@ -6,12 +6,8 @@
 """
 
 from . import errors
-from .types import GenerateParams, ImageParams, VideoParams
 
 __all__ = [
-    "GenerateParams",
-    "ImageParams",
-    "VideoParams",
     "errors",
 ]
 

src/ai/models/ai_gateway/generate.py

@@ -1,7 +1,6 @@
 """AI Gateway v3 generation adapter — image-model and video-model endpoints.
 
-Provides typed parameter objects (:class:`ImageParams`, :class:`VideoParams`)
-and a unified :func:`generate` entry point that dispatches based on param type
+Unified :func:`generate` entry point that dispatches based on param type
 and validates against model capabilities.
 """
 
@@ -16,10 +15,10 @@
 from ..core import client as client_
 from ..core import model as model_
 from ..core.helpers import files
+from ..core.types import GenerateParams as GenerateParams
+from ..core.types import ImageParams as ImageParams
+from ..core.types import VideoParams as VideoParams
 from . import _common, errors
-from .types import GenerateParams as GenerateParams
-from .types import ImageParams as ImageParams
-from .types import VideoParams as VideoParams
 
 # ---------------------------------------------------------------------------
 # Image generation — /image-model

src/ai/models/ai_gateway/types.py

@@ -7,16 +7,19 @@
 from .client import Client
 from .model import Model, ModelCost
 from .proto import CheckConnFn, GenerateFn, StreamFn
-from .stream_result import StreamResult
+from .types import GenerateParams, ImageParams, StreamResult, VideoParams
 
 __all__ = [
     "CheckConnFn",
     "Client",
     "GenerateFn",
+    "GenerateParams",
+    "ImageParams",
     "Model",
     "ModelCost",
     "StreamFn",
     "StreamResult",
+    "VideoParams",
     "check_connection",
     "generate",
     "get_models",

src/ai/models/core/__init__.py

@@ -14,10 +14,10 @@
 from ...types import messages as messages_
 from ...types import stream as stream_
 from ...types import tools as tools_
-from ..ai_gateway import types as ai_gw_types
-from . import adapters, stream_result
+from . import adapters
 from . import client as client_
 from . import model as model_
+from . import types as types_
 
 
 async def stream(
@@ -50,7 +50,7 @@ async def stream(
     async def _real(call: middleware_.ModelContext) -> stream_.StreamResultLike:
         c = call.client or client_.auto_client(call.model)
         adapter_fn = adapters.get_stream_adapter(call.model.adapter)
-        return stream_result.StreamResult(
+        return types_.StreamResult(
             adapter_fn(
                 c,
                 call.model,
@@ -68,7 +68,7 @@ async def _real(call: middleware_.ModelContext) -> stream_.StreamResultLike:
 async def generate(
     model: model_.Model,
     messages: list[messages_.Message],
-    params: ai_gw_types.GenerateParams | None = None,
+    params: types_.GenerateParams | None = None,
     *,
     client: client_.Client | None = None,
 ) -> messages_.Message:

src/ai/models/core/api.py

@@ -1,12 +1,67 @@
-"""StreamResult — concrete wrapper around a message stream."""
+"""Core model-layer types — parameter objects and StreamResult.
+
+Parameter types (:class:`ImageParams`, :class:`VideoParams`) live here
+because they parameterise the public :func:`generate` API.
+
+:class:`StreamResult` is the concrete wrapper returned by :func:`stream`.
+"""
 
 from __future__ import annotations
 
 from collections.abc import AsyncGenerator
 from typing import Any
 
+import pydantic
+
 from ...types import messages as messages_
 
+# ---------------------------------------------------------------------------
+# Generation parameter types
+# ---------------------------------------------------------------------------
+
+_PARAMS_CONFIG = pydantic.ConfigDict(frozen=True, populate_by_name=True)
+
+
+class ImageParams(pydantic.BaseModel):
+    """Parameters for image generation (``/image-model`` endpoint)."""
+
+    model_config = _PARAMS_CONFIG
+
+    n: int = 1
+    size: str | None = None
+    aspect_ratio: str | None = pydantic.Field(
+        default=None, serialization_alias="aspectRatio"
+    )
+    seed: int | None = None
+    provider_options: dict[str, Any] = pydantic.Field(
+        default_factory=dict, serialization_alias="providerOptions"
+    )
+
+
+class VideoParams(pydantic.BaseModel):
+    """Parameters for video generation (``/video-model`` endpoint)."""
+
+    model_config = _PARAMS_CONFIG
+
+    n: int = 1
+    aspect_ratio: str | None = pydantic.Field(
+        default=None, serialization_alias="aspectRatio"
+    )
+    resolution: str | None = None
+    duration: int | None = None
+    fps: int | None = None
+    seed: int | None = None
+    provider_options: dict[str, Any] = pydantic.Field(
+        default_factory=dict, serialization_alias="providerOptions"
+    )
+
+
+GenerateParams = ImageParams | VideoParams
+
+# ---------------------------------------------------------------------------
+# StreamResult
+# ---------------------------------------------------------------------------
+
 
 class StreamResult:
     """Wrapper around a message stream. Async-iterable; collects the final result.