orq-ai-sdk 4.2.0rc28__py3-none-any.whl → 4.2.6__py3-none-any.whl

orq_ai_sdk/models/createimageop.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 from .publiccontact import PublicContact, PublicContactTypedDict
+from .publicidentity import PublicIdentity, PublicIdentityTypedDict
 from orq_ai_sdk.types import (
     BaseModel,
     Nullable,
@@ -9,9 +10,10 @@ from orq_ai_sdk.types import (
     UNSET,
     UNSET_SENTINEL,
 )
+import pydantic
 from pydantic import model_serializer
-from typing import Any, Dict, List, Literal, Optional
-from typing_extensions import NotRequired, TypedDict, deprecated
+from typing import List, Literal, Optional
+from typing_extensions import Annotated, NotRequired, TypedDict
 
 
 Background = Literal[
@@ -80,6 +82,22 @@ class CreateImageRetry(BaseModel):
     on_codes: Optional[List[float]] = None
     r"""HTTP status codes that trigger retry logic"""
 
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(["count", "on_codes"])
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m
+
 
 class CreateImageFallbacksTypedDict(TypedDict):
     model: str
@@ -91,51 +109,6 @@ class CreateImageFallbacks(BaseModel):
     r"""Fallback model identifier"""
 
 
-@deprecated(
-    "warning: ** DEPRECATED ** - This will be removed in a future release, please migrate away from it as soon as possible."
-)
-class CreateImageContactTypedDict(TypedDict):
-    r"""@deprecated Use identity instead. Information about the contact making the request."""
-
-    id: str
-    r"""Unique identifier for the contact"""
-    display_name: NotRequired[str]
-    r"""Display name of the contact"""
-    email: NotRequired[str]
-    r"""Email address of the contact"""
-    metadata: NotRequired[List[Dict[str, Any]]]
-    r"""A hash of key/value pairs containing any other data about the contact"""
-    logo_url: NotRequired[str]
-    r"""URL to the contact's avatar or logo"""
-    tags: NotRequired[List[str]]
-    r"""A list of tags associated with the contact"""
-
-
-@deprecated(
-    "warning: ** DEPRECATED ** - This will be removed in a future release, please migrate away from it as soon as possible."
-)
-class CreateImageContact(BaseModel):
-    r"""@deprecated Use identity instead. Information about the contact making the request."""
-
-    id: str
-    r"""Unique identifier for the contact"""
-
-    display_name: Optional[str] = None
-    r"""Display name of the contact"""
-
-    email: Optional[str] = None
-    r"""Email address of the contact"""
-
-    metadata: Optional[List[Dict[str, Any]]] = None
-    r"""A hash of key/value pairs containing any other data about the contact"""
-
-    logo_url: Optional[str] = None
-    r"""URL to the contact's avatar or logo"""
-
-    tags: Optional[List[str]] = None
-    r"""A list of tags associated with the contact"""
-
-
 CreateImageType = Literal["exact_match",]
 
 
@@ -155,32 +128,74 @@ class CreateImageCache(BaseModel):
     ttl: Optional[float] = 1800
     r"""Time to live for cached responses in seconds. Maximum 259200 seconds (3 days)."""
 
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(["ttl"])
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m
+
 
 CreateImageLoadBalancerType = Literal["weight_based",]
 
 
-class CreateImageLoadBalancer1TypedDict(TypedDict):
-    type: CreateImageLoadBalancerType
+class CreateImageLoadBalancerModelsTypedDict(TypedDict):
     model: str
     r"""Model identifier for load balancing"""
     weight: NotRequired[float]
     r"""Weight assigned to this model for load balancing"""
 
 
-class CreateImageLoadBalancer1(BaseModel):
-    type: CreateImageLoadBalancerType
-
+class CreateImageLoadBalancerModels(BaseModel):
     model: str
     r"""Model identifier for load balancing"""
 
     weight: Optional[float] = 0.5
     r"""Weight assigned to this model for load balancing"""
 
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(["weight"])
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m
+
+
+class CreateImageLoadBalancer1TypedDict(TypedDict):
+    type: CreateImageLoadBalancerType
+    models: List[CreateImageLoadBalancerModelsTypedDict]
+
+
+class CreateImageLoadBalancer1(BaseModel):
+    type: CreateImageLoadBalancerType
+
+    models: List[CreateImageLoadBalancerModels]
+
 
 CreateImageLoadBalancerTypedDict = CreateImageLoadBalancer1TypedDict
+r"""Array of models with weights for load balancing requests"""
 
 
 CreateImageLoadBalancer = CreateImageLoadBalancer1
+r"""Array of models with weights for load balancing requests"""
 
 
 class CreateImageTimeoutTypedDict(TypedDict):
@@ -204,12 +219,13 @@ class CreateImageOrqTypedDict(TypedDict):
     r"""Retry configuration for the request"""
     fallbacks: NotRequired[List[CreateImageFallbacksTypedDict]]
     r"""Array of fallback models to use if primary model fails"""
-    identity: NotRequired[PublicContactTypedDict]
+    identity: NotRequired[PublicIdentityTypedDict]
     r"""Information about the identity making the request. If the identity does not exist, it will be created automatically."""
-    contact: NotRequired[CreateImageContactTypedDict]
+    contact: NotRequired[PublicContactTypedDict]
+    r"""@deprecated Use identity instead. Information about the contact making the request."""
     cache: NotRequired[CreateImageCacheTypedDict]
     r"""Cache configuration for the request."""
-    load_balancer: NotRequired[List[CreateImageLoadBalancerTypedDict]]
+    load_balancer: NotRequired[CreateImageLoadBalancerTypedDict]
     r"""Array of models with weights for load balancing requests"""
     timeout: NotRequired[CreateImageTimeoutTypedDict]
     r"""Timeout configuration to apply to the request. If the request exceeds the timeout, it will be retried or fallback to the next model if configured."""
@@ -225,20 +241,53 @@ class CreateImageOrq(BaseModel):
     fallbacks: Optional[List[CreateImageFallbacks]] = None
     r"""Array of fallback models to use if primary model fails"""
 
-    identity: Optional[PublicContact] = None
+    identity: Optional[PublicIdentity] = None
     r"""Information about the identity making the request. If the identity does not exist, it will be created automatically."""
 
-    contact: Optional[CreateImageContact] = None
+    contact: Annotated[
+        Optional[PublicContact],
+        pydantic.Field(
+            deprecated="warning: ** DEPRECATED ** - This will be removed in a future release, please migrate away from it as soon as possible."
+        ),
+    ] = None
+    r"""@deprecated Use identity instead. Information about the contact making the request."""
 
     cache: Optional[CreateImageCache] = None
     r"""Cache configuration for the request."""
 
-    load_balancer: Optional[List[CreateImageLoadBalancer]] = None
+    load_balancer: Optional[CreateImageLoadBalancer] = None
     r"""Array of models with weights for load balancing requests"""
 
     timeout: Optional[CreateImageTimeout] = None
     r"""Timeout configuration to apply to the request. If the request exceeds the timeout, it will be retried or fallback to the next model if configured."""
 
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(
+            [
+                "name",
+                "retry",
+                "fallbacks",
+                "identity",
+                "contact",
+                "cache",
+                "load_balancer",
+                "timeout",
+            ]
+        )
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m
+
 
 class CreateImageRequestBodyTypedDict(TypedDict):
     r"""input"""
@@ -308,52 +357,51 @@ class CreateImageRequestBody(BaseModel):
 
     @model_serializer(mode="wrap")
     def serialize_model(self, handler):
-        optional_fields = [
-            "background",
-            "moderation",
-            "n",
-            "output_compression",
-            "output_format",
-            "quality",
-            "response_format",
-            "size",
-            "style",
-            "orq",
-        ]
-        nullable_fields = [
-            "background",
-            "moderation",
-            "n",
-            "output_compression",
-            "output_format",
-            "quality",
-            "response_format",
-            "size",
-            "style",
-        ]
-        null_default_fields = []
-
+        optional_fields = set(
+            [
+                "background",
+                "moderation",
+                "n",
+                "output_compression",
+                "output_format",
+                "quality",
+                "response_format",
+                "size",
+                "style",
+                "orq",
+            ]
+        )
+        nullable_fields = set(
+            [
+                "background",
+                "moderation",
+                "n",
+                "output_compression",
+                "output_format",
+                "quality",
+                "response_format",
+                "size",
+                "style",
+            ]
+        )
         serialized = handler(self)
-
         m = {}
 
         for n, f in type(self).model_fields.items():
             k = f.alias or n
             val = serialized.get(k)
-            serialized.pop(k, None)
-
-            optional_nullable = k in optional_fields and k in nullable_fields
-            is_set = (
-                self.__pydantic_fields_set__.intersection({n})
-                or k in null_default_fields
-            )  # pylint: disable=no-member
-
-            if val is not None and val != UNSET_SENTINEL:
-                m[k] = val
-            elif val != UNSET_SENTINEL and (
-                not k in optional_fields or (optional_nullable and is_set)
-            ):
-                m[k] = val
+            is_nullable_and_explicitly_set = (
+                k in nullable_fields
+                and (self.__pydantic_fields_set__.intersection({n}))  # pylint: disable=no-member
+            )
+
+            if val != UNSET_SENTINEL:
+                if (
+                    val is not None
+                    or k not in optional_fields
+                    or is_nullable_and_explicitly_set
+                ):
+                    m[k] = val
 
         return m
 
@@ -379,31 +427,26 @@ class CreateImageData(BaseModel):
 
     @model_serializer(mode="wrap")
     def serialize_model(self, handler):
-        optional_fields = ["revised_prompt", "b64_json", "url"]
-        nullable_fields = ["revised_prompt"]
-        null_default_fields = []
-
+        optional_fields = set(["revised_prompt", "b64_json", "url"])
+        nullable_fields = set(["revised_prompt"])
         serialized = handler(self)
-
         m = {}
 
         for n, f in type(self).model_fields.items():
             k = f.alias or n
             val = serialized.get(k)
-            serialized.pop(k, None)
-
-            optional_nullable = k in optional_fields and k in nullable_fields
-            is_set = (
-                self.__pydantic_fields_set__.intersection({n})
-                or k in null_default_fields
-            )  # pylint: disable=no-member
-
-            if val is not None and val != UNSET_SENTINEL:
-                m[k] = val
-            elif val != UNSET_SENTINEL and (
-                not k in optional_fields or (optional_nullable and is_set)
-            ):
-                m[k] = val
+            is_nullable_and_explicitly_set = (
+                k in nullable_fields
+                and (self.__pydantic_fields_set__.intersection({n}))  # pylint: disable=no-member
+            )
+
+            if val != UNSET_SENTINEL:
+                if (
+                    val is not None
+                    or k not in optional_fields
+                    or is_nullable_and_explicitly_set
+                ):
+                    m[k] = val
 
         return m
 
@@ -418,6 +461,22 @@ class InputTokensDetails(BaseModel):
 
     text_tokens: Optional[float] = None
 
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(["image_tokens", "text_tokens"])
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m
+
 
 class CreateImageUsageTypedDict(TypedDict):
     input_tokens_details: InputTokensDetailsTypedDict
@@ -435,6 +494,22 @@ class CreateImageUsage(BaseModel):
 
     total_tokens: Optional[float] = None
 
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(["input_tokens", "output_tokens", "total_tokens"])
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m
+
 
 class CreateImageResponseBodyTypedDict(TypedDict):
     r"""Represents an image generation response from the API."""
@@ -456,3 +531,19 @@ class CreateImageResponseBody(BaseModel):
     r"""Represents the url or the content of an image generated."""
 
     usage: Optional[CreateImageUsage] = None
+
+    @model_serializer(mode="wrap")
+    def serialize_model(self, handler):
+        optional_fields = set(["usage"])
+        serialized = handler(self)
+        m = {}
+
+        for n, f in type(self).model_fields.items():
+            k = f.alias or n
+            val = serialized.get(k)
+
+            if val != UNSET_SENTINEL:
+                if val is not None or k not in optional_fields:
+                    m[k] = val
+
+        return m