athena-intelligence 0.1.125__py3-none-any.whl → 0.1.126__py3-none-any.whl

athena/core/__init__.py

@@ -3,28 +3,45 @@
 from .api_error import ApiError
 from .client_wrapper import AsyncClientWrapper, BaseClientWrapper, SyncClientWrapper
 from .datetime_utils import serialize_datetime
-from .file import File, convert_file_dict_to_httpx_tuples
+from .file import File, convert_file_dict_to_httpx_tuples, with_content_type
 from .http_client import AsyncHttpClient, HttpClient
 from .jsonable_encoder import jsonable_encoder
-from .pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
+from .pydantic_utilities import (
+    IS_PYDANTIC_V2,
+    UniversalBaseModel,
+    UniversalRootModel,
+    parse_obj_as,
+    universal_field_validator,
+    universal_root_validator,
+    update_forward_refs,
+)
 from .query_encoder import encode_query
 from .remove_none_from_dict import remove_none_from_dict
 from .request_options import RequestOptions
+from .serialization import FieldMetadata, convert_and_respect_annotation_metadata
 
 __all__ = [
     "ApiError",
     "AsyncClientWrapper",
     "AsyncHttpClient",
     "BaseClientWrapper",
+    "FieldMetadata",
     "File",
     "HttpClient",
+    "IS_PYDANTIC_V2",
     "RequestOptions",
     "SyncClientWrapper",
+    "UniversalBaseModel",
+    "UniversalRootModel",
+    "convert_and_respect_annotation_metadata",
     "convert_file_dict_to_httpx_tuples",
-    "deep_union_pydantic_dicts",
     "encode_query",
     "jsonable_encoder",
-    "pydantic_v1",
+    "parse_obj_as",
     "remove_none_from_dict",
     "serialize_datetime",
+    "universal_field_validator",
+    "universal_root_validator",
+    "update_forward_refs",
+    "with_content_type",
 ]

athena/core/client_wrapper.py

@@ -1,10 +1,9 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-
 import httpx
-
-from .http_client import AsyncHttpClient, HttpClient
+from .http_client import HttpClient
+from .http_client import AsyncHttpClient
 
 
 class BaseClientWrapper:
@@ -17,7 +16,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "athena-intelligence",
-            "X-Fern-SDK-Version": "0.1.125",
+            "X-Fern-SDK-Version": "0.1.126",
         }
         headers["X-API-KEY"] = self.api_key
         return headers
@@ -36,9 +35,9 @@ class SyncClientWrapper(BaseClientWrapper):
         super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
         self.httpx_client = HttpClient(
             httpx_client=httpx_client,
-            base_headers=self.get_headers(),
-            base_timeout=self.get_timeout(),
-            base_url=self.get_base_url(),
+            base_headers=self.get_headers,
+            base_timeout=self.get_timeout,
+            base_url=self.get_base_url,
         )
 
 
@@ -49,7 +48,7 @@ class AsyncClientWrapper(BaseClientWrapper):
         super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
         self.httpx_client = AsyncHttpClient(
             httpx_client=httpx_client,
-            base_headers=self.get_headers(),
-            base_timeout=self.get_timeout(),
-            base_url=self.get_base_url(),
+            base_headers=self.get_headers,
+            base_timeout=self.get_timeout,
+            base_url=self.get_base_url,
         )

athena/core/file.py

@@ -1,25 +1,30 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import typing
+from typing import IO, Dict, List, Mapping, Optional, Tuple, Union, cast
 
 # File typing inspired by the flexibility of types within the httpx library
 # https://github.com/encode/httpx/blob/master/httpx/_types.py
-FileContent = typing.Union[typing.IO[bytes], bytes, str]
-File = typing.Union[
+FileContent = Union[IO[bytes], bytes, str]
+File = Union[
     # file (or bytes)
     FileContent,
     # (filename, file (or bytes))
-    typing.Tuple[typing.Optional[str], FileContent],
+    Tuple[Optional[str], FileContent],
     # (filename, file (or bytes), content_type)
-    typing.Tuple[typing.Optional[str], FileContent, typing.Optional[str]],
+    Tuple[Optional[str], FileContent, Optional[str]],
     # (filename, file (or bytes), content_type, headers)
-    typing.Tuple[typing.Optional[str], FileContent, typing.Optional[str], typing.Mapping[str, str]],
+    Tuple[
+        Optional[str],
+        FileContent,
+        Optional[str],
+        Mapping[str, str],
+    ],
 ]
 
 
 def convert_file_dict_to_httpx_tuples(
-    d: typing.Dict[str, typing.Union[File, typing.List[File]]]
-) -> typing.List[typing.Tuple[str, File]]:
+    d: Dict[str, Union[File, List[File]]],
+) -> List[Tuple[str, File]]:
     """
     The format we use is a list of tuples, where the first element is the
     name of the file and the second is the file object. Typically HTTPX wants
@@ -36,3 +41,27 @@ def convert_file_dict_to_httpx_tuples(
         else:
             httpx_tuples.append((key, file_like))
     return httpx_tuples
+
+
+def with_content_type(*, file: File, default_content_type: str) -> File:
+    """
+    This function resolves to the file's content type, if provided, and defaults
+    to the default_content_type value if not.
+    """
+    if isinstance(file, tuple):
+        if len(file) == 2:
+            filename, content = cast(Tuple[Optional[str], FileContent], file)  # type: ignore
+            return (filename, content, default_content_type)
+        elif len(file) == 3:
+            filename, content, file_content_type = cast(Tuple[Optional[str], FileContent, Optional[str]], file)  # type: ignore
+            out_content_type = file_content_type or default_content_type
+            return (filename, content, out_content_type)
+        elif len(file) == 4:
+            filename, content, file_content_type, headers = cast(  # type: ignore
+                Tuple[Optional[str], FileContent, Optional[str], Mapping[str, str]], file
+            )
+            out_content_type = file_content_type or default_content_type
+            return (filename, content, out_content_type, headers)
+        else:
+            raise ValueError(f"Unexpected tuple length: {len(file)}")
+    return (None, file, default_content_type)

athena/core/http_client.py

@@ -2,6 +2,7 @@
 
 import asyncio
 import email.utils
+import json
 import re
 import time
 import typing
@@ -84,12 +85,13 @@ def _retry_timeout(response: httpx.Response, retries: int) -> float:
 
 
 def _should_retry(response: httpx.Response) -> bool:
-    retriable_400s = [429, 408, 409]
-    return response.status_code >= 500 or response.status_code in retriable_400s
+    retryable_400s = [429, 408, 409]
+    return response.status_code >= 500 or response.status_code in retryable_400s
 
 
 def remove_omit_from_dict(
-    original: typing.Dict[str, typing.Optional[typing.Any]], omit: typing.Optional[typing.Any]
+    original: typing.Dict[str, typing.Optional[typing.Any]],
+    omit: typing.Optional[typing.Any],
 ) -> typing.Dict[str, typing.Any]:
     if omit is None:
         return original
@@ -106,14 +108,18 @@ def maybe_filter_request_body(
     omit: typing.Optional[typing.Any],
 ) -> typing.Optional[typing.Any]:
     if data is None:
-        return None
+        return (
+            jsonable_encoder(request_options.get("additional_body_parameters", {})) or {}
+            if request_options is not None
+            else None
+        )
     elif not isinstance(data, typing.Mapping):
         data_content = jsonable_encoder(data)
     else:
         data_content = {
             **(jsonable_encoder(remove_omit_from_dict(data, omit))),  # type: ignore
             **(
-                jsonable_encoder(request_options.get("additional_body_parameters", {}))
+                jsonable_encoder(request_options.get("additional_body_parameters", {})) or {}
                 if request_options is not None
                 else {}
             ),
@@ -121,14 +127,34 @@ def maybe_filter_request_body(
     return data_content
 
 
+# Abstracted out for testing purposes
+def get_request_body(
+    *,
+    json: typing.Optional[typing.Any],
+    data: typing.Optional[typing.Any],
+    request_options: typing.Optional[RequestOptions],
+    omit: typing.Optional[typing.Any],
+) -> typing.Tuple[typing.Optional[typing.Any], typing.Optional[typing.Any]]:
+    json_body = None
+    data_body = None
+    if data is not None:
+        data_body = maybe_filter_request_body(data, request_options, omit)
+    else:
+        # If both data and json are None, we send json data in the event extra properties are specified
+        json_body = maybe_filter_request_body(json, request_options, omit)
+
+    # If you have an empty JSON body, you should just send None
+    return (json_body if json_body != {} else None), data_body if data_body != {} else None
+
+
 class HttpClient:
     def __init__(
         self,
         *,
         httpx_client: httpx.Client,
-        base_timeout: typing.Optional[float],
-        base_headers: typing.Dict[str, str],
-        base_url: typing.Optional[str] = None,
+        base_timeout: typing.Callable[[], typing.Optional[float]],
+        base_headers: typing.Callable[[], typing.Dict[str, str]],
+        base_url: typing.Optional[typing.Callable[[], str]] = None,
     ):
         self.base_url = base_url
         self.base_timeout = base_timeout
@@ -136,7 +162,10 @@ class HttpClient:
         self.httpx_client = httpx_client
 
     def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
-        base_url = self.base_url if maybe_base_url is None else maybe_base_url
+        base_url = maybe_base_url
+        if self.base_url is not None and base_url is None:
+            base_url = self.base_url()
+
         if base_url is None:
             raise ValueError("A base_url is required to make this request, please provide one and try again.")
         return base_url
@@ -154,25 +183,27 @@ class HttpClient:
         files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
         headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
         request_options: typing.Optional[RequestOptions] = None,
-        retries: int = 0,
+        retries: int = 2,
         omit: typing.Optional[typing.Any] = None,
     ) -> httpx.Response:
         base_url = self.get_base_url(base_url)
         timeout = (
             request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self.base_timeout
+            else self.base_timeout()
         )
 
+        json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
         response = self.httpx_client.request(
             method=method,
             url=urllib.parse.urljoin(f"{base_url}/", path),
             headers=jsonable_encoder(
                 remove_none_from_dict(
                     {
-                        **self.base_headers,
+                        **self.base_headers(),
                         **(headers if headers is not None else {}),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+                        **(request_options.get("additional_headers", {}) or {} if request_options is not None else {}),
                     }
                 )
             ),
@@ -183,7 +214,7 @@ class HttpClient:
                             {
                                 **(params if params is not None else {}),
                                 **(
-                                    request_options.get("additional_query_parameters", {})
+                                    request_options.get("additional_query_parameters", {}) or {}
                                     if request_options is not None
                                     else {}
                                 ),
@@ -193,10 +224,14 @@ class HttpClient:
                     )
                 )
             ),
-            json=maybe_filter_request_body(json, request_options, omit),
-            data=maybe_filter_request_body(data, request_options, omit),
+            json=json_body,
+            data=data_body,
             content=content,
-            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            files=(
+                convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+                if (files is not None and files is not omit)
+                else None
+            ),
             timeout=timeout,
         )
 
@@ -234,23 +269,25 @@ class HttpClient:
         files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
         headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
         request_options: typing.Optional[RequestOptions] = None,
-        retries: int = 0,
+        retries: int = 2,
         omit: typing.Optional[typing.Any] = None,
     ) -> typing.Iterator[httpx.Response]:
         base_url = self.get_base_url(base_url)
         timeout = (
             request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self.base_timeout
+            else self.base_timeout()
         )
 
+        json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
         with self.httpx_client.stream(
             method=method,
             url=urllib.parse.urljoin(f"{base_url}/", path),
             headers=jsonable_encoder(
                 remove_none_from_dict(
                     {
-                        **self.base_headers,
+                        **self.base_headers(),
                         **(headers if headers is not None else {}),
                         **(request_options.get("additional_headers", {}) if request_options is not None else {}),
                     }
@@ -273,10 +310,14 @@ class HttpClient:
                     )
                 )
             ),
-            json=maybe_filter_request_body(json, request_options, omit),
-            data=maybe_filter_request_body(data, request_options, omit),
+            json=json_body,
+            data=data_body,
             content=content,
-            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            files=(
+                convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+                if (files is not None and files is not omit)
+                else None
+            ),
             timeout=timeout,
         ) as stream:
             yield stream
@@ -287,9 +328,9 @@ class AsyncHttpClient:
         self,
         *,
         httpx_client: httpx.AsyncClient,
-        base_timeout: typing.Optional[float],
-        base_headers: typing.Dict[str, str],
-        base_url: typing.Optional[str] = None,
+        base_timeout: typing.Callable[[], typing.Optional[float]],
+        base_headers: typing.Callable[[], typing.Dict[str, str]],
+        base_url: typing.Optional[typing.Callable[[], str]] = None,
     ):
         self.base_url = base_url
         self.base_timeout = base_timeout
@@ -297,7 +338,10 @@ class AsyncHttpClient:
         self.httpx_client = httpx_client
 
     def get_base_url(self, maybe_base_url: typing.Optional[str]) -> str:
-        base_url = self.base_url if maybe_base_url is None else maybe_base_url
+        base_url = maybe_base_url
+        if self.base_url is not None and base_url is None:
+            base_url = self.base_url()
+
         if base_url is None:
             raise ValueError("A base_url is required to make this request, please provide one and try again.")
         return base_url
@@ -315,16 +359,18 @@ class AsyncHttpClient:
         files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
         headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
         request_options: typing.Optional[RequestOptions] = None,
-        retries: int = 0,
+        retries: int = 2,
         omit: typing.Optional[typing.Any] = None,
     ) -> httpx.Response:
         base_url = self.get_base_url(base_url)
         timeout = (
             request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self.base_timeout
+            else self.base_timeout()
         )
 
+        json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
         # Add the input to each of these and do None-safety checks
         response = await self.httpx_client.request(
             method=method,
@@ -332,9 +378,9 @@ class AsyncHttpClient:
             headers=jsonable_encoder(
                 remove_none_from_dict(
                     {
-                        **self.base_headers,
+                        **self.base_headers(),
                         **(headers if headers is not None else {}),
-                        **(request_options.get("additional_headers", {}) if request_options is not None else {}),
+                        **(request_options.get("additional_headers", {}) or {} if request_options is not None else {}),
                     }
                 )
             ),
@@ -345,7 +391,7 @@ class AsyncHttpClient:
                             {
                                 **(params if params is not None else {}),
                                 **(
-                                    request_options.get("additional_query_parameters", {})
+                                    request_options.get("additional_query_parameters", {}) or {}
                                     if request_options is not None
                                     else {}
                                 ),
@@ -355,10 +401,14 @@ class AsyncHttpClient:
                     )
                 )
             ),
-            json=maybe_filter_request_body(json, request_options, omit),
-            data=maybe_filter_request_body(data, request_options, omit),
+            json=json_body,
+            data=data_body,
             content=content,
-            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            files=(
+                convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+                if files is not None
+                else None
+            ),
             timeout=timeout,
         )
 
@@ -395,23 +445,25 @@ class AsyncHttpClient:
         files: typing.Optional[typing.Dict[str, typing.Optional[typing.Union[File, typing.List[File]]]]] = None,
         headers: typing.Optional[typing.Dict[str, typing.Any]] = None,
         request_options: typing.Optional[RequestOptions] = None,
-        retries: int = 0,
+        retries: int = 2,
         omit: typing.Optional[typing.Any] = None,
     ) -> typing.AsyncIterator[httpx.Response]:
         base_url = self.get_base_url(base_url)
         timeout = (
             request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else self.base_timeout
+            else self.base_timeout()
         )
 
+        json_body, data_body = get_request_body(json=json, data=data, request_options=request_options, omit=omit)
+
         async with self.httpx_client.stream(
             method=method,
             url=urllib.parse.urljoin(f"{base_url}/", path),
             headers=jsonable_encoder(
                 remove_none_from_dict(
                     {
-                        **self.base_headers,
+                        **self.base_headers(),
                         **(headers if headers is not None else {}),
                         **(request_options.get("additional_headers", {}) if request_options is not None else {}),
                     }
@@ -434,10 +486,14 @@ class AsyncHttpClient:
                     )
                 )
             ),
-            json=maybe_filter_request_body(json, request_options, omit),
-            data=maybe_filter_request_body(data, request_options, omit),
+            json=json_body,
+            data=data_body,
             content=content,
-            files=convert_file_dict_to_httpx_tuples(remove_none_from_dict(files)) if files is not None else None,
+            files=(
+                convert_file_dict_to_httpx_tuples(remove_omit_from_dict(remove_none_from_dict(files), omit))
+                if files is not None
+                else None
+            ),
             timeout=timeout,
         ) as stream:
             yield stream

athena/core/jsonable_encoder.py

@@ -8,33 +8,27 @@ Taken from FastAPI, and made a bit simpler
 https://github.com/tiangolo/fastapi/blob/master/fastapi/encoders.py
 """
 
+import base64
 import dataclasses
 import datetime as dt
-from collections import defaultdict
 from enum import Enum
 from pathlib import PurePath
 from types import GeneratorType
-from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
+from typing import Any, Callable, Dict, List, Optional, Set, Union
+
+import pydantic
 
 from .datetime_utils import serialize_datetime
-from .pydantic_utilities import pydantic_v1
+from .pydantic_utilities import (
+    IS_PYDANTIC_V2,
+    encode_by_type,
+    to_jsonable_with_fallback,
+)
 
 SetIntStr = Set[Union[int, str]]
 DictIntStrAny = Dict[Union[int, str], Any]
 
 
-def generate_encoders_by_class_tuples(
-    type_encoder_map: Dict[Any, Callable[[Any], Any]]
-) -> Dict[Callable[[Any], Any], Tuple[Any, ...]]:
-    encoders_by_class_tuples: Dict[Callable[[Any], Any], Tuple[Any, ...]] = defaultdict(tuple)
-    for type_, encoder in type_encoder_map.items():
-        encoders_by_class_tuples[encoder] += (type_,)
-    return encoders_by_class_tuples
-
-
-encoders_by_class_tuples = generate_encoders_by_class_tuples(pydantic_v1.json.ENCODERS_BY_TYPE)
-
-
 def jsonable_encoder(obj: Any, custom_encoder: Optional[Dict[Any, Callable[[Any], Any]]] = None) -> Any:
     custom_encoder = custom_encoder or {}
     if custom_encoder:
@@ -44,17 +38,24 @@ def jsonable_encoder(obj: Any, custom_encoder: Optional[Dict[Any, Callable[[Any]
             for encoder_type, encoder_instance in custom_encoder.items():
                 if isinstance(obj, encoder_type):
                     return encoder_instance(obj)
-    if isinstance(obj, pydantic_v1.BaseModel):
-        encoder = getattr(obj.__config__, "json_encoders", {})
+    if isinstance(obj, pydantic.BaseModel):
+        if IS_PYDANTIC_V2:
+            encoder = getattr(obj.model_config, "json_encoders", {})  # type: ignore # Pydantic v2
+        else:
+            encoder = getattr(obj.__config__, "json_encoders", {})  # type: ignore # Pydantic v1
         if custom_encoder:
             encoder.update(custom_encoder)
         obj_dict = obj.dict(by_alias=True)
         if "__root__" in obj_dict:
             obj_dict = obj_dict["__root__"]
+        if "root" in obj_dict:
+            obj_dict = obj_dict["root"]
         return jsonable_encoder(obj_dict, custom_encoder=encoder)
     if dataclasses.is_dataclass(obj):
-        obj_dict = dataclasses.asdict(obj)
+        obj_dict = dataclasses.asdict(obj)  # type: ignore
         return jsonable_encoder(obj_dict, custom_encoder=custom_encoder)
+    if isinstance(obj, bytes):
+        return base64.b64encode(obj).decode("utf-8")
     if isinstance(obj, Enum):
         return obj.value
     if isinstance(obj, PurePath):
@@ -80,20 +81,21 @@ def jsonable_encoder(obj: Any, custom_encoder: Optional[Dict[Any, Callable[[Any]
             encoded_list.append(jsonable_encoder(item, custom_encoder=custom_encoder))
         return encoded_list
 
-    if type(obj) in pydantic_v1.json.ENCODERS_BY_TYPE:
-        return pydantic_v1.json.ENCODERS_BY_TYPE[type(obj)](obj)
-    for encoder, classes_tuple in encoders_by_class_tuples.items():
-        if isinstance(obj, classes_tuple):
-            return encoder(obj)
+    def fallback_serializer(o: Any) -> Any:
+        attempt_encode = encode_by_type(o)
+        if attempt_encode is not None:
+            return attempt_encode
 
-    try:
-        data = dict(obj)
-    except Exception as e:
-        errors: List[Exception] = []
-        errors.append(e)
         try:
-            data = vars(obj)
+            data = dict(o)
         except Exception as e:
+            errors: List[Exception] = []
             errors.append(e)
-            raise ValueError(errors) from e
-    return jsonable_encoder(data, custom_encoder=custom_encoder)
+            try:
+                data = vars(o)
+            except Exception as e:
+                errors.append(e)
+                raise ValueError(errors) from e
+        return jsonable_encoder(data, custom_encoder=custom_encoder)
+
+    return to_jsonable_with_fallback(obj, fallback_serializer)