agenta 0.27.0a9__py3-none-any.whl → 0.27.0a12__py3-none-any.whl

agenta/client/backend/core/pydantic_utilities.py

@@ -10,6 +10,7 @@ import typing_extensions
 import pydantic
 
 from .datetime_utils import serialize_datetime
+from .serialization import convert_and_respect_annotation_metadata
 
 IS_PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
 
@@ -56,11 +57,14 @@ Model = typing.TypeVar("Model", bound=pydantic.BaseModel)
 
 
 def parse_obj_as(type_: typing.Type[T], object_: typing.Any) -> T:
+    dealiased_object = convert_and_respect_annotation_metadata(
+        object_=object_, annotation=type_, direction="read"
+    )
     if IS_PYDANTIC_V2:
         adapter = pydantic.TypeAdapter(type_)  # type: ignore # Pydantic v2
-        return adapter.validate_python(object_)
+        return adapter.validate_python(dealiased_object)
     else:
-        return pydantic.parse_obj_as(type_, object_)
+        return pydantic.parse_obj_as(type_, dealiased_object)
 
 
 def to_jsonable_with_fallback(
@@ -75,9 +79,53 @@ def to_jsonable_with_fallback(
 
 
 class UniversalBaseModel(pydantic.BaseModel):
-    class Config:
-        populate_by_name = True
-        json_encoders = {dt.datetime: serialize_datetime}
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(
+            # Allow fields begining with `model_` to be used in the model
+            protected_namespaces=(),
+        )  # type: ignore # Pydantic v2
+
+        @pydantic.model_serializer(mode="wrap", when_used="json")  # type: ignore # Pydantic v2
+        def serialize_model(
+            self, handler: pydantic.SerializerFunctionWrapHandler
+        ) -> typing.Any:  # type: ignore # Pydantic v2
+            serialized = handler(self)
+            data = {
+                k: serialize_datetime(v) if isinstance(v, dt.datetime) else v
+                for k, v in serialized.items()
+            }
+            return data
+
+    else:
+
+        class Config:
+            smart_union = True
+            json_encoders = {dt.datetime: serialize_datetime}
+
+    @classmethod
+    def model_construct(
+        cls: typing.Type["Model"],
+        _fields_set: typing.Optional[typing.Set[str]] = None,
+        **values: typing.Any,
+    ) -> "Model":
+        dealiased_object = convert_and_respect_annotation_metadata(
+            object_=values, annotation=cls, direction="read"
+        )
+        return cls.construct(_fields_set, **dealiased_object)
+
+    @classmethod
+    def construct(
+        cls: typing.Type["Model"],
+        _fields_set: typing.Optional[typing.Set[str]] = None,
+        **values: typing.Any,
+    ) -> "Model":
+        dealiased_object = convert_and_respect_annotation_metadata(
+            object_=values, annotation=cls, direction="read"
+        )
+        if IS_PYDANTIC_V2:
+            return super().model_construct(_fields_set, **dealiased_object)  # type: ignore # Pydantic v2
+        else:
+            return super().construct(_fields_set, **dealiased_object)
 
     def json(self, **kwargs: typing.Any) -> str:
         kwargs_with_defaults: typing.Any = {
@@ -95,30 +143,97 @@ class UniversalBaseModel(pydantic.BaseModel):
         Override the default dict method to `exclude_unset` by default. This function patches
         `exclude_unset` to work include fields within non-None default values.
         """
-        _fields_set = self.__fields_set__
-
-        fields = _get_model_fields(self.__class__)
-        for name, field in fields.items():
-            if name not in _fields_set:
-                default = _get_field_default(field)
-
-                # If the default values are non-null act like they've been set
-                # This effectively allows exclude_unset to work like exclude_none where
-                # the latter passes through intentionally set none values.
-                if default != None:
-                    _fields_set.add(name)
-
-        kwargs_with_defaults_exclude_unset: typing.Any = {
-            "by_alias": True,
-            "exclude_unset": True,
-            "include": _fields_set,
-            **kwargs,
-        }
-
+        # Note: the logic here is multi-plexed given the levers exposed in Pydantic V1 vs V2
+        # Pydantic V1's .dict can be extremely slow, so we do not want to call it twice.
+        #
+        # We'd ideally do the same for Pydantic V2, but it shells out to a library to serialize models
+        # that we have less control over, and this is less intrusive than custom serializers for now.
         if IS_PYDANTIC_V2:
-            return super().model_dump(**kwargs_with_defaults_exclude_unset)  # type: ignore # Pydantic v2
+            kwargs_with_defaults_exclude_unset: typing.Any = {
+                **kwargs,
+                "by_alias": True,
+                "exclude_unset": True,
+                "exclude_none": False,
+            }
+            kwargs_with_defaults_exclude_none: typing.Any = {
+                **kwargs,
+                "by_alias": True,
+                "exclude_none": True,
+                "exclude_unset": False,
+            }
+            dict_dump = deep_union_pydantic_dicts(
+                super().model_dump(**kwargs_with_defaults_exclude_unset),  # type: ignore # Pydantic v2
+                super().model_dump(**kwargs_with_defaults_exclude_none),  # type: ignore # Pydantic v2
+            )
+
+        else:
+            _fields_set = self.__fields_set__.copy()
+
+            fields = _get_model_fields(self.__class__)
+            for name, field in fields.items():
+                if name not in _fields_set:
+                    default = _get_field_default(field)
+
+                    # If the default values are non-null act like they've been set
+                    # This effectively allows exclude_unset to work like exclude_none where
+                    # the latter passes through intentionally set none values.
+                    if default is not None or (
+                        "exclude_unset" in kwargs and not kwargs["exclude_unset"]
+                    ):
+                        _fields_set.add(name)
+
+                        if default is not None:
+                            self.__fields_set__.add(name)
+
+            kwargs_with_defaults_exclude_unset_include_fields: typing.Any = {
+                "by_alias": True,
+                "exclude_unset": True,
+                "include": _fields_set,
+                **kwargs,
+            }
+
+            dict_dump = super().dict(
+                **kwargs_with_defaults_exclude_unset_include_fields
+            )
+
+        return convert_and_respect_annotation_metadata(
+            object_=dict_dump, annotation=self.__class__, direction="write"
+        )
+
+
+def _union_list_of_pydantic_dicts(
+    source: typing.List[typing.Any], destination: typing.List[typing.Any]
+) -> typing.List[typing.Any]:
+    converted_list: typing.List[typing.Any] = []
+    for i, item in enumerate(source):
+        destination_value = destination[i]  # type: ignore
+        if isinstance(item, dict):
+            converted_list.append(deep_union_pydantic_dicts(item, destination_value))
+        elif isinstance(item, list):
+            converted_list.append(
+                _union_list_of_pydantic_dicts(item, destination_value)
+            )
         else:
-            return super().dict(**kwargs_with_defaults_exclude_unset)
+            converted_list.append(item)
+    return converted_list
+
+
+def deep_union_pydantic_dicts(
+    source: typing.Dict[str, typing.Any], destination: typing.Dict[str, typing.Any]
+) -> typing.Dict[str, typing.Any]:
+    for key, value in source.items():
+        node = destination.setdefault(key, {})
+        if isinstance(value, dict):
+            deep_union_pydantic_dicts(value, node)
+        # Note: we do not do this same processing for sets given we do not have sets of models
+        # and given the sets are unordered, the processing of the set and matching objects would
+        # be non-trivial.
+        elif isinstance(value, list):
+            destination[key] = _union_list_of_pydantic_dicts(value, node)
+        else:
+            destination[key] = value
+
+    return destination
 
 
 if IS_PYDANTIC_V2:
@@ -145,11 +260,11 @@ def encode_by_type(o: typing.Any) -> typing.Any:
             return encoder(o)
 
 
-def update_forward_refs(model: typing.Type["Model"]) -> None:
+def update_forward_refs(model: typing.Type["Model"], **localns: typing.Any) -> None:
     if IS_PYDANTIC_V2:
         model.model_rebuild(raise_errors=False)  # type: ignore # Pydantic v2
     else:
-        model.update_forward_refs()
+        model.update_forward_refs(**localns)
 
 
 # Mirrors Pydantic's internal typing

agenta/client/backend/core/request_options.py

@@ -23,6 +23,8 @@ class RequestOptions(typing.TypedDict, total=False):
         - additional_query_parameters: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's query parameters dict
 
         - additional_body_parameters: typing.Dict[str, typing.Any]. A dictionary containing additional parameters to spread into the request's body parameters dict
+
+        - chunk_size: int. The size, in bytes, to process each chunk of data being streamed back within the response. This equates to leveraging `chunk_size` within `requests` or `httpx`, and is only leveraged for file downloads.
     """
 
     timeout_in_seconds: NotRequired[int]
@@ -30,3 +32,4 @@ class RequestOptions(typing.TypedDict, total=False):
     additional_headers: NotRequired[typing.Dict[str, typing.Any]]
     additional_query_parameters: NotRequired[typing.Dict[str, typing.Any]]
     additional_body_parameters: NotRequired[typing.Dict[str, typing.Any]]
+    chunk_size: NotRequired[int]

agenta/client/backend/core/serialization.py

@@ -1,10 +1,13 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import collections
+import inspect
 import typing
 
 import typing_extensions
 
+import pydantic
+
 
 class FieldMetadata:
     """
@@ -29,6 +32,7 @@ def convert_and_respect_annotation_metadata(
     object_: typing.Any,
     annotation: typing.Any,
     inner_type: typing.Optional[typing.Any] = None,
+    direction: typing.Literal["read", "write"],
 ) -> typing.Any:
     """
     Respect the metadata annotations on a field, such as aliasing. This function effectively
@@ -56,49 +60,79 @@ def convert_and_respect_annotation_metadata(
         inner_type = annotation
 
     clean_type = _remove_annotations(inner_type)
+    # Pydantic models
+    if (
+        inspect.isclass(clean_type)
+        and issubclass(clean_type, pydantic.BaseModel)
+        and isinstance(object_, typing.Mapping)
+    ):
+        return _convert_mapping(object_, clean_type, direction)
+    # TypedDicts
     if typing_extensions.is_typeddict(clean_type) and isinstance(
         object_, typing.Mapping
     ):
-        return _convert_typeddict(object_, clean_type)
+        return _convert_mapping(object_, clean_type, direction)
 
     if (
-        # If you're iterating on a string, do not bother to coerce it to a sequence.
-        (not isinstance(object_, str))
-        and (
-            (
-                (
-                    typing_extensions.get_origin(clean_type) == typing.List
-                    or typing_extensions.get_origin(clean_type) == list
-                    or clean_type == typing.List
-                )
-                and isinstance(object_, typing.List)
-            )
-            or (
-                (
-                    typing_extensions.get_origin(clean_type) == typing.Set
-                    or typing_extensions.get_origin(clean_type) == set
-                    or clean_type == typing.Set
-                )
-                and isinstance(object_, typing.Set)
+        typing_extensions.get_origin(clean_type) == typing.Dict
+        or typing_extensions.get_origin(clean_type) == dict
+        or clean_type == typing.Dict
+    ) and isinstance(object_, typing.Dict):
+        key_type = typing_extensions.get_args(clean_type)[0]
+        value_type = typing_extensions.get_args(clean_type)[1]
+
+        return {
+            key: convert_and_respect_annotation_metadata(
+                object_=value,
+                annotation=annotation,
+                inner_type=value_type,
+                direction=direction,
             )
-            or (
-                (
-                    typing_extensions.get_origin(clean_type) == typing.Sequence
-                    or typing_extensions.get_origin(clean_type)
-                    == collections.abc.Sequence
-                    or clean_type == typing.Sequence
+            for key, value in object_.items()
+        }
+
+    # If you're iterating on a string, do not bother to coerce it to a sequence.
+    if not isinstance(object_, str):
+        if (
+            typing_extensions.get_origin(clean_type) == typing.Set
+            or typing_extensions.get_origin(clean_type) == set
+            or clean_type == typing.Set
+        ) and isinstance(object_, typing.Set):
+            inner_type = typing_extensions.get_args(clean_type)[0]
+            return {
+                convert_and_respect_annotation_metadata(
+                    object_=item,
+                    annotation=annotation,
+                    inner_type=inner_type,
+                    direction=direction,
                 )
-                and isinstance(object_, typing.Sequence)
+                for item in object_
+            }
+        elif (
+            (
+                typing_extensions.get_origin(clean_type) == typing.List
+                or typing_extensions.get_origin(clean_type) == list
+                or clean_type == typing.List
             )
-        )
-    ):
-        inner_type = typing_extensions.get_args(clean_type)[0]
-        return [
-            convert_and_respect_annotation_metadata(
-                object_=item, annotation=annotation, inner_type=inner_type
+            and isinstance(object_, typing.List)
+        ) or (
+            (
+                typing_extensions.get_origin(clean_type) == typing.Sequence
+                or typing_extensions.get_origin(clean_type) == collections.abc.Sequence
+                or clean_type == typing.Sequence
             )
-            for item in object_
-        ]
+            and isinstance(object_, typing.Sequence)
+        ):
+            inner_type = typing_extensions.get_args(clean_type)[0]
+            return [
+                convert_and_respect_annotation_metadata(
+                    object_=item,
+                    annotation=annotation,
+                    inner_type=inner_type,
+                    direction=direction,
+                )
+                for item in object_
+            ]
 
     if typing_extensions.get_origin(clean_type) == typing.Union:
         # We should be able to ~relatively~ safely try to convert keys against all
@@ -107,7 +141,10 @@ def convert_and_respect_annotation_metadata(
         # Or if another member aliases a field of the same name that another member does not.
         for member in typing_extensions.get_args(clean_type):
             object_ = convert_and_respect_annotation_metadata(
-                object_=object_, annotation=annotation, inner_type=member
+                object_=object_,
+                annotation=annotation,
+                inner_type=member,
+                direction=direction,
             )
         return object_
 
@@ -120,19 +157,37 @@ def convert_and_respect_annotation_metadata(
     return object_
 
 
-def _convert_typeddict(
-    object_: typing.Mapping[str, object], expected_type: typing.Any
+def _convert_mapping(
+    object_: typing.Mapping[str, object],
+    expected_type: typing.Any,
+    direction: typing.Literal["read", "write"],
 ) -> typing.Mapping[str, object]:
     converted_object: typing.Dict[str, object] = {}
     annotations = typing_extensions.get_type_hints(expected_type, include_extras=True)
+    aliases_to_field_names = _get_alias_to_field_name(annotations)
     for key, value in object_.items():
-        type_ = annotations.get(key)
+        if direction == "read" and key in aliases_to_field_names:
+            dealiased_key = aliases_to_field_names.get(key)
+            if dealiased_key is not None:
+                type_ = annotations.get(dealiased_key)
+        else:
+            type_ = annotations.get(key)
+        # Note you can't get the annotation by the field name if you're in read mode, so you must check the aliases map
+        #
+        # So this is effectively saying if we're in write mode, and we don't have a type, or if we're in read mode and we don't have an alias
+        # then we can just pass the value through as is
         if type_ is None:
             converted_object[key] = value
+        elif direction == "read" and key not in aliases_to_field_names:
+            converted_object[key] = convert_and_respect_annotation_metadata(
+                object_=value, annotation=type_, direction=direction
+            )
         else:
             converted_object[
-                _alias_key(key, type_)
-            ] = convert_and_respect_annotation_metadata(object_=value, annotation=type_)
+                _alias_key(key, type_, direction, aliases_to_field_names)
+            ] = convert_and_respect_annotation_metadata(
+                object_=value, annotation=type_, direction=direction
+            )
     return converted_object
 
 
@@ -165,7 +220,39 @@ def _remove_annotations(type_: typing.Any) -> typing.Any:
     return type_
 
 
-def _alias_key(key: str, type_: typing.Any) -> str:
+def get_alias_to_field_mapping(type_: typing.Any) -> typing.Dict[str, str]:
+    annotations = typing_extensions.get_type_hints(type_, include_extras=True)
+    return _get_alias_to_field_name(annotations)
+
+
+def get_field_to_alias_mapping(type_: typing.Any) -> typing.Dict[str, str]:
+    annotations = typing_extensions.get_type_hints(type_, include_extras=True)
+    return _get_field_to_alias_name(annotations)
+
+
+def _get_alias_to_field_name(
+    field_to_hint: typing.Dict[str, typing.Any],
+) -> typing.Dict[str, str]:
+    aliases = {}
+    for field, hint in field_to_hint.items():
+        maybe_alias = _get_alias_from_type(hint)
+        if maybe_alias is not None:
+            aliases[maybe_alias] = field
+    return aliases
+
+
+def _get_field_to_alias_name(
+    field_to_hint: typing.Dict[str, typing.Any],
+) -> typing.Dict[str, str]:
+    aliases = {}
+    for field, hint in field_to_hint.items():
+        maybe_alias = _get_alias_from_type(hint)
+        if maybe_alias is not None:
+            aliases[field] = maybe_alias
+    return aliases
+
+
+def _get_alias_from_type(type_: typing.Any) -> typing.Optional[str]:
     maybe_annotated_type = _get_annotation(type_)
 
     if maybe_annotated_type is not None:
@@ -175,5 +262,15 @@ def _alias_key(key: str, type_: typing.Any) -> str:
         for annotation in annotations:
             if isinstance(annotation, FieldMetadata) and annotation.alias is not None:
                 return annotation.alias
+    return None
+
 
-    return key
+def _alias_key(
+    key: str,
+    type_: typing.Any,
+    direction: typing.Literal["read", "write"],
+    aliases_to_field_names: typing.Dict[str, str],
+) -> str:
+    if direction == "read":
+        return aliases_to_field_names.get(key, key)
+    return _get_alias_from_type(type_=type_) or key

agenta/client/backend/evaluations/client.py

@@ -10,6 +10,7 @@ from json.decoder import JSONDecodeError
 from ..core.api_error import ApiError
 from ..types.evaluation import Evaluation
 from ..types.llm_run_rate_limit import LlmRunRateLimit
+from ..core.serialization import convert_and_respect_annotation_metadata
 from ..core.jsonable_encoder import jsonable_encoder
 from ..types.evaluation_scenario import EvaluationScenario
 from ..core.client_wrapper import AsyncClientWrapper
@@ -249,7 +250,9 @@ class EvaluationsClient:
                 "variant_ids": variant_ids,
                 "evaluators_configs": evaluators_configs,
                 "testset_id": testset_id,
-                "rate_limit": rate_limit,
+                "rate_limit": convert_and_respect_annotation_metadata(
+                    object_=rate_limit, annotation=LlmRunRateLimit, direction="write"
+                ),
                 "lm_providers_keys": lm_providers_keys,
                 "correct_answer_column": correct_answer_column,
             },
@@ -959,7 +962,9 @@ class AsyncEvaluationsClient:
                 "variant_ids": variant_ids,
                 "evaluators_configs": evaluators_configs,
                 "testset_id": testset_id,
-                "rate_limit": rate_limit,
+                "rate_limit": convert_and_respect_annotation_metadata(
+                    object_=rate_limit, annotation=LlmRunRateLimit, direction="write"
+                ),
                 "lm_providers_keys": lm_providers_keys,
                 "correct_answer_column": correct_answer_column,
             },