athena-intelligence 0.1.44__tar.gz → 0.1.49__tar.gz

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: athena-intelligence
-Version: 0.1.44
+Version: 0.1.49
 Summary: 
 Requires-Python: >=3.8,<4.0
 Classifier: Programming Language :: Python :: 3

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "athena-intelligence"
-version = "0.1.44"
+version = "0.1.49"
 description = ""
 readme = "README.md"
 authors = []
@@ -15,8 +15,18 @@ pydantic = ">= 1.9.2"
 typing_extensions = ">= 4.0.0"
 
 [tool.poetry.dev-dependencies]
-mypy = "^1.8.0"
+mypy = "1.9.0"
 pytest = "^7.4.0"
+pytest-asyncio = "^0.23.5"
+python-dateutil = "^2.9.0"
+
+[tool.pytest.ini_options]
+testpaths = [ "tests" ]
+asyncio_mode = "auto"
+
+[tool.mypy]
+plugins = ["pydantic.mypy"]
+
 
 [build-system]
 requires = ["poetry-core"]

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/__init__.py

@@ -10,9 +10,11 @@ from .types import (
     GetSnippetsResponse,
     HttpValidationError,
     LangchainDocumentsRequestOut,
+    LlmModel,
     MessageOut,
     MessageOutDto,
     Model,
+    PlanExecuteOut,
     Report,
     Snippet,
     SqlResults,
@@ -25,9 +27,9 @@ from .types import (
     ValidationErrorLocItem,
 )
 from .errors import UnprocessableEntityError
-from . import chain, dataset, message, query, report, search, snippet, tools
-from .chain import StructuredParseInParsingModel
+from . import chain, dataset, message, query, report, search, snippet, tasks, tools
 from .environment import AthenaEnvironment
+from .version import __version__
 
 __all__ = [
     "AthenaEnvironment",
@@ -40,14 +42,15 @@ __all__ = [
     "GetSnippetsResponse",
     "HttpValidationError",
     "LangchainDocumentsRequestOut",
+    "LlmModel",
     "MessageOut",
     "MessageOutDto",
     "Model",
+    "PlanExecuteOut",
     "Report",
     "Snippet",
     "SqlResults",
     "StatusEnum",
-    "StructuredParseInParsingModel",
     "StructuredParseResult",
     "ToolModels",
     "Tools",
@@ -55,6 +58,7 @@ __all__ = [
     "UrlResult",
     "ValidationError",
     "ValidationErrorLocItem",
+    "__version__",
     "chain",
     "dataset",
     "message",
@@ -62,5 +66,6 @@ __all__ = [
     "report",
     "search",
     "snippet",
+    "tasks",
     "tools",
 ]

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/base_client.py

@@ -13,6 +13,7 @@ from .query.client import AsyncQueryClient, QueryClient
 from .report.client import AsyncReportClient, ReportClient
 from .search.client import AsyncSearchClient, SearchClient
 from .snippet.client import AsyncSnippetClient, SnippetClient
+from .tasks.client import AsyncTasksClient, TasksClient
 from .tools.client import AsyncToolsClient, ToolsClient
 
 
@@ -29,7 +30,9 @@ class BaseAthena:
 
         - api_key: str.
 
-        - timeout: typing.Optional[float]. The timeout to be used, in seconds, for requests by default the timeout is 60 seconds.
+        - timeout: typing.Optional[float]. The timeout to be used, in seconds, for requests by default the timeout is 60 seconds, unless a custom httpx client is used, in which case a default is not set.
+
+        - follow_redirects: typing.Optional[bool]. Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
 
         - httpx_client: typing.Optional[httpx.Client]. The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
     ---
@@ -46,13 +49,20 @@ class BaseAthena:
         base_url: typing.Optional[str] = None,
         environment: AthenaEnvironment = AthenaEnvironment.DEFAULT,
         api_key: str,
-        timeout: typing.Optional[float] = 60,
+        timeout: typing.Optional[float] = None,
+        follow_redirects: typing.Optional[bool] = True,
         httpx_client: typing.Optional[httpx.Client] = None
     ):
+        _defaulted_timeout = timeout if timeout is not None else 60 if httpx_client is None else None
         self._client_wrapper = SyncClientWrapper(
             base_url=_get_base_url(base_url=base_url, environment=environment),
             api_key=api_key,
-            httpx_client=httpx.Client(timeout=timeout) if httpx_client is None else httpx_client,
+            httpx_client=httpx_client
+            if httpx_client is not None
+            else httpx.Client(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
+            if follow_redirects is not None
+            else httpx.Client(timeout=_defaulted_timeout),
+            timeout=_defaulted_timeout,
         )
         self.message = MessageClient(client_wrapper=self._client_wrapper)
         self.dataset = DatasetClient(client_wrapper=self._client_wrapper)
@@ -62,6 +72,7 @@ class BaseAthena:
         self.search = SearchClient(client_wrapper=self._client_wrapper)
         self.chain = ChainClient(client_wrapper=self._client_wrapper)
         self.tools = ToolsClient(client_wrapper=self._client_wrapper)
+        self.tasks = TasksClient(client_wrapper=self._client_wrapper)
 
 
 class AsyncBaseAthena:
@@ -77,7 +88,9 @@ class AsyncBaseAthena:
 
         - api_key: str.
 
-        - timeout: typing.Optional[float]. The timeout to be used, in seconds, for requests by default the timeout is 60 seconds.
+        - timeout: typing.Optional[float]. The timeout to be used, in seconds, for requests by default the timeout is 60 seconds, unless a custom httpx client is used, in which case a default is not set.
+
+        - follow_redirects: typing.Optional[bool]. Whether the default httpx client follows redirects or not, this is irrelevant if a custom httpx client is passed in.
 
         - httpx_client: typing.Optional[httpx.AsyncClient]. The httpx client to use for making requests, a preconfigured client is used by default, however this is useful should you want to pass in any custom httpx configuration.
     ---
@@ -94,13 +107,20 @@ class AsyncBaseAthena:
         base_url: typing.Optional[str] = None,
         environment: AthenaEnvironment = AthenaEnvironment.DEFAULT,
         api_key: str,
-        timeout: typing.Optional[float] = 60,
+        timeout: typing.Optional[float] = None,
+        follow_redirects: typing.Optional[bool] = True,
         httpx_client: typing.Optional[httpx.AsyncClient] = None
     ):
+        _defaulted_timeout = timeout if timeout is not None else 60 if httpx_client is None else None
         self._client_wrapper = AsyncClientWrapper(
             base_url=_get_base_url(base_url=base_url, environment=environment),
             api_key=api_key,
-            httpx_client=httpx.AsyncClient(timeout=timeout) if httpx_client is None else httpx_client,
+            httpx_client=httpx_client
+            if httpx_client is not None
+            else httpx.AsyncClient(timeout=_defaulted_timeout, follow_redirects=follow_redirects)
+            if follow_redirects is not None
+            else httpx.AsyncClient(timeout=_defaulted_timeout),
+            timeout=_defaulted_timeout,
         )
         self.message = AsyncMessageClient(client_wrapper=self._client_wrapper)
         self.dataset = AsyncDatasetClient(client_wrapper=self._client_wrapper)
@@ -110,6 +130,7 @@ class AsyncBaseAthena:
         self.search = AsyncSearchClient(client_wrapper=self._client_wrapper)
         self.chain = AsyncChainClient(client_wrapper=self._client_wrapper)
         self.tools = AsyncToolsClient(client_wrapper=self._client_wrapper)
+        self.tasks = AsyncTasksClient(client_wrapper=self._client_wrapper)
 
 
 def _get_base_url(*, base_url: typing.Optional[str] = None, environment: AthenaEnvironment) -> str:

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/chain/client.py

@@ -7,17 +7,13 @@ from json.decoder import JSONDecodeError
 from ..core.api_error import ApiError
 from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from ..core.jsonable_encoder import jsonable_encoder
+from ..core.pydantic_utilities import pydantic_v1
 from ..core.remove_none_from_dict import remove_none_from_dict
 from ..core.request_options import RequestOptions
 from ..errors.unprocessable_entity_error import UnprocessableEntityError
 from ..types.http_validation_error import HttpValidationError
+from ..types.llm_model import LlmModel
 from ..types.structured_parse_result import StructuredParseResult
-from .types.structured_parse_in_parsing_model import StructuredParseInParsingModel
-
-try:
-    import pydantic.v1 as pydantic  # type: ignore
-except ImportError:
-    import pydantic  # type: ignore
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -32,7 +28,7 @@ class ChainClient:
         *,
         text_input: str,
         custom_type_dict: typing.Dict[str, typing.Any],
-        parsing_model: typing.Optional[StructuredParseInParsingModel] = OMIT,
+        model: LlmModel,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> StructuredParseResult:
         """
@@ -41,33 +37,32 @@ class ChainClient:
 
             - custom_type_dict: typing.Dict[str, typing.Any]. A dictionary of field names and their default values.
 
-            - parsing_model: typing.Optional[StructuredParseInParsingModel]. The model to be used for parsing.
+            - model: LlmModel.
 
             - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
         ---
-        from athena import StructuredParseInParsingModel
+        from athena import LlmModel
         from athena.client import Athena
 
-        client = Athena(api_key="YOUR_API_KEY", )
-        client.chain.structured_parse(text_input='Athena is an AI-native analytics platform and artificial employee built to accelerate analytics workflows
-                        by offering enterprise teams co-pilot and auto-pilot modes. Athena learns your workflow as a co-pilot,
-                        allowing you to hand over controls to her for autonomous execution with confidence."
-
-                        Give me all of the modes Athena provides.', custom_type_dict={"modes": {}}, parsing_model=StructuredParseInParsingModel.GPT_4_TURBO, )
+        client = Athena(
+            api_key="YOUR_API_KEY",
+        )
+        client.chain.structured_parse(
+            text_input='Athena is an AI-native analytics platform and artificial employee built to accelerate analytics workflows \n                by offering enterprise teams co-pilot and auto-pilot modes. Athena learns your workflow as a co-pilot, \n                allowing you to hand over controls to her for autonomous execution with confidence." \n                \n                Give me all of the modes Athena provides.',
+            custom_type_dict={"modes": {}},
+            model=LlmModel.GPT_35_TURBO,
+        )
         """
-        _request: typing.Dict[str, typing.Any] = {"text_input": text_input, "custom_type_dict": custom_type_dict}
-        if parsing_model is not OMIT:
-            _request["parsing_model"] = parsing_model
         _response = self._client_wrapper.httpx_client.request(
-            "POST",
-            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
+            method="POST",
+            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
             params=jsonable_encoder(
                 request_options.get("additional_query_parameters") if request_options is not None else None
             ),
-            json=jsonable_encoder(_request)
+            json=jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model})
             if request_options is None or request_options.get("additional_body_parameters") is None
             else {
-                **jsonable_encoder(_request),
+                **jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model}),
                 **(jsonable_encoder(remove_none_from_dict(request_options.get("additional_body_parameters", {})))),
             },
             headers=jsonable_encoder(
@@ -80,14 +75,16 @@ class ChainClient:
             ),
             timeout=request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else 60,
+            else self._client_wrapper.get_timeout(),
             retries=0,
             max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
         )
         if 200 <= _response.status_code < 300:
-            return pydantic.parse_obj_as(StructuredParseResult, _response.json())  # type: ignore
+            return pydantic_v1.parse_obj_as(StructuredParseResult, _response.json())  # type: ignore
         if _response.status_code == 422:
-            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
+            raise UnprocessableEntityError(
+                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
+            )
         try:
             _response_json = _response.json()
         except JSONDecodeError:
@@ -104,7 +101,7 @@ class AsyncChainClient:
         *,
         text_input: str,
         custom_type_dict: typing.Dict[str, typing.Any],
-        parsing_model: typing.Optional[StructuredParseInParsingModel] = OMIT,
+        model: LlmModel,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> StructuredParseResult:
         """
@@ -113,33 +110,32 @@ class AsyncChainClient:
 
             - custom_type_dict: typing.Dict[str, typing.Any]. A dictionary of field names and their default values.
 
-            - parsing_model: typing.Optional[StructuredParseInParsingModel]. The model to be used for parsing.
+            - model: LlmModel.
 
             - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
         ---
-        from athena import StructuredParseInParsingModel
+        from athena import LlmModel
         from athena.client import AsyncAthena
 
-        client = AsyncAthena(api_key="YOUR_API_KEY", )
-        await client.chain.structured_parse(text_input='Athena is an AI-native analytics platform and artificial employee built to accelerate analytics workflows
-                        by offering enterprise teams co-pilot and auto-pilot modes. Athena learns your workflow as a co-pilot,
-                        allowing you to hand over controls to her for autonomous execution with confidence."
-
-                        Give me all of the modes Athena provides.', custom_type_dict={"modes": {}}, parsing_model=StructuredParseInParsingModel.GPT_4_TURBO, )
+        client = AsyncAthena(
+            api_key="YOUR_API_KEY",
+        )
+        await client.chain.structured_parse(
+            text_input='Athena is an AI-native analytics platform and artificial employee built to accelerate analytics workflows \n                by offering enterprise teams co-pilot and auto-pilot modes. Athena learns your workflow as a co-pilot, \n                allowing you to hand over controls to her for autonomous execution with confidence." \n                \n                Give me all of the modes Athena provides.',
+            custom_type_dict={"modes": {}},
+            model=LlmModel.GPT_35_TURBO,
+        )
         """
-        _request: typing.Dict[str, typing.Any] = {"text_input": text_input, "custom_type_dict": custom_type_dict}
-        if parsing_model is not OMIT:
-            _request["parsing_model"] = parsing_model
         _response = await self._client_wrapper.httpx_client.request(
-            "POST",
-            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
+            method="POST",
+            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/structured-parse"),
             params=jsonable_encoder(
                 request_options.get("additional_query_parameters") if request_options is not None else None
             ),
-            json=jsonable_encoder(_request)
+            json=jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model})
             if request_options is None or request_options.get("additional_body_parameters") is None
             else {
-                **jsonable_encoder(_request),
+                **jsonable_encoder({"text_input": text_input, "custom_type_dict": custom_type_dict, "model": model}),
                 **(jsonable_encoder(remove_none_from_dict(request_options.get("additional_body_parameters", {})))),
             },
             headers=jsonable_encoder(
@@ -152,14 +148,16 @@ class AsyncChainClient:
             ),
             timeout=request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else 60,
+            else self._client_wrapper.get_timeout(),
             retries=0,
             max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
         )
         if 200 <= _response.status_code < 300:
-            return pydantic.parse_obj_as(StructuredParseResult, _response.json())  # type: ignore
+            return pydantic_v1.parse_obj_as(StructuredParseResult, _response.json())  # type: ignore
         if _response.status_code == 422:
-            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
+            raise UnprocessableEntityError(
+                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
+            )
         try:
             _response_json = _response.json()
         except JSONDecodeError:

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/core/__init__.py

@@ -6,6 +6,7 @@ from .datetime_utils import serialize_datetime
 from .file import File, convert_file_dict_to_httpx_tuples
 from .http_client import AsyncHttpClient, HttpClient
 from .jsonable_encoder import jsonable_encoder
+from .pydantic_utilities import pydantic_v1
 from .remove_none_from_dict import remove_none_from_dict
 from .request_options import RequestOptions
 
@@ -20,6 +21,7 @@ __all__ = [
     "SyncClientWrapper",
     "convert_file_dict_to_httpx_tuples",
     "jsonable_encoder",
+    "pydantic_v1",
     "remove_none_from_dict",
     "serialize_datetime",
 ]

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/core/client_wrapper.py

@@ -8,15 +8,16 @@ from .http_client import AsyncHttpClient, HttpClient
 
 
 class BaseClientWrapper:
-    def __init__(self, *, api_key: str, base_url: str):
+    def __init__(self, *, api_key: str, base_url: str, timeout: typing.Optional[float] = None):
         self.api_key = api_key
         self._base_url = base_url
+        self._timeout = timeout
 
     def get_headers(self) -> typing.Dict[str, str]:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "athena-intelligence",
-            "X-Fern-SDK-Version": "0.1.44",
+            "X-Fern-SDK-Version": "0.1.49",
         }
         headers["X-API-KEY"] = self.api_key
         return headers
@@ -24,14 +25,21 @@ class BaseClientWrapper:
     def get_base_url(self) -> str:
         return self._base_url
 
+    def get_timeout(self) -> typing.Optional[float]:
+        return self._timeout
+
 
 class SyncClientWrapper(BaseClientWrapper):
-    def __init__(self, *, api_key: str, base_url: str, httpx_client: httpx.Client):
-        super().__init__(api_key=api_key, base_url=base_url)
+    def __init__(
+        self, *, api_key: str, base_url: str, timeout: typing.Optional[float] = None, httpx_client: httpx.Client
+    ):
+        super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
         self.httpx_client = HttpClient(httpx_client=httpx_client)
 
 
 class AsyncClientWrapper(BaseClientWrapper):
-    def __init__(self, *, api_key: str, base_url: str, httpx_client: httpx.AsyncClient):
-        super().__init__(api_key=api_key, base_url=base_url)
+    def __init__(
+        self, *, api_key: str, base_url: str, timeout: typing.Optional[float] = None, httpx_client: httpx.AsyncClient
+    ):
+        super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
         self.httpx_client = AsyncHttpClient(httpx_client=httpx_client)

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/core/http_client.py

@@ -5,6 +5,7 @@ import email.utils
 import re
 import time
 import typing
+from contextlib import asynccontextmanager, contextmanager
 from functools import wraps
 from random import random
 
@@ -98,8 +99,10 @@ class HttpClient:
         return response
 
     @wraps(httpx.Client.stream)
+    @contextmanager
     def stream(self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any) -> typing.Any:
-        return self.httpx_client.stream(*args, **kwargs)
+        with self.httpx_client.stream(*args, **kwargs) as stream:
+            yield stream
 
 
 class AsyncHttpClient:
@@ -118,8 +121,10 @@ class AsyncHttpClient:
                 return await self.request(max_retries=max_retries, retries=retries + 1, *args, **kwargs)
         return response
 
-    @wraps(httpx.AsyncClient.request)
+    @wraps(httpx.AsyncClient.stream)
+    @asynccontextmanager
     async def stream(
         self, *args: typing.Any, max_retries: int = 0, retries: int = 0, **kwargs: typing.Any
     ) -> typing.Any:
-        return self.httpx_client.stream(*args, **kwargs)
+        async with self.httpx_client.stream(*args, **kwargs) as stream:
+            yield stream

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/core/jsonable_encoder.py

@@ -16,12 +16,8 @@ from pathlib import PurePath
 from types import GeneratorType
 from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
 
-try:
-    import pydantic.v1 as pydantic  # type: ignore
-except ImportError:
-    import pydantic  # type: ignore
-
 from .datetime_utils import serialize_datetime
+from .pydantic_utilities import pydantic_v1
 
 SetIntStr = Set[Union[int, str]]
 DictIntStrAny = Dict[Union[int, str], Any]
@@ -36,7 +32,7 @@ def generate_encoders_by_class_tuples(
     return encoders_by_class_tuples
 
 
-encoders_by_class_tuples = generate_encoders_by_class_tuples(pydantic.json.ENCODERS_BY_TYPE)
+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:
@@ -48,7 +44,7 @@ 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.BaseModel):
+    if isinstance(obj, pydantic_v1.BaseModel):
         encoder = getattr(obj.__config__, "json_encoders", {})
         if custom_encoder:
             encoder.update(custom_encoder)
@@ -65,10 +61,10 @@ def jsonable_encoder(obj: Any, custom_encoder: Optional[Dict[Any, Callable[[Any]
         return str(obj)
     if isinstance(obj, (str, int, float, type(None))):
         return obj
-    if isinstance(obj, dt.date):
-        return str(obj)
     if isinstance(obj, dt.datetime):
         return serialize_datetime(obj)
+    if isinstance(obj, dt.date):
+        return str(obj)
     if isinstance(obj, dict):
         encoded_dict = {}
         allowed_keys = set(obj.keys())
@@ -84,8 +80,8 @@ 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.json.ENCODERS_BY_TYPE:
-        return pydantic.json.ENCODERS_BY_TYPE[type(obj)](obj)
+    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)

athena_intelligence-0.1.49/src/athena/core/pydantic_utilities.py

@@ -0,0 +1,12 @@
+# This file was auto-generated by Fern from our API Definition.
+
+import pydantic
+
+IS_PYDANTIC_V2 = pydantic.VERSION.startswith("2.")
+
+if IS_PYDANTIC_V2:
+    import pydantic.v1 as pydantic_v1  # type: ignore  # nopycln: import
+else:
+    import pydantic as pydantic_v1  # type: ignore  # nopycln: import
+
+__all__ = ["pydantic_v1"]

{athena_intelligence-0.1.44 → athena_intelligence-0.1.49}/src/athena/dataset/client.py

@@ -7,17 +7,13 @@ from json.decoder import JSONDecodeError
 from ..core.api_error import ApiError
 from ..core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from ..core.jsonable_encoder import jsonable_encoder
+from ..core.pydantic_utilities import pydantic_v1
 from ..core.remove_none_from_dict import remove_none_from_dict
 from ..core.request_options import RequestOptions
 from ..errors.unprocessable_entity_error import UnprocessableEntityError
 from ..types.get_datasets_response import GetDatasetsResponse
 from ..types.http_validation_error import HttpValidationError
 
-try:
-    import pydantic.v1 as pydantic  # type: ignore
-except ImportError:
-    import pydantic  # type: ignore
-
 
 class DatasetClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
@@ -46,8 +42,8 @@ class DatasetClient:
         client.dataset.get()
         """
         _response = self._client_wrapper.httpx_client.request(
-            "GET",
-            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
+            method="GET",
+            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
             params=jsonable_encoder(
                 remove_none_from_dict(
                     {
@@ -71,14 +67,16 @@ class DatasetClient:
             ),
             timeout=request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else 60,
+            else self._client_wrapper.get_timeout(),
             retries=0,
             max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
         )
         if 200 <= _response.status_code < 300:
-            return pydantic.parse_obj_as(GetDatasetsResponse, _response.json())  # type: ignore
+            return pydantic_v1.parse_obj_as(GetDatasetsResponse, _response.json())  # type: ignore
         if _response.status_code == 422:
-            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
+            raise UnprocessableEntityError(
+                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
+            )
         try:
             _response_json = _response.json()
         except JSONDecodeError:
@@ -113,8 +111,8 @@ class AsyncDatasetClient:
         await client.dataset.get()
         """
         _response = await self._client_wrapper.httpx_client.request(
-            "GET",
-            urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
+            method="GET",
+            url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/datasets"),
             params=jsonable_encoder(
                 remove_none_from_dict(
                     {
@@ -138,14 +136,16 @@ class AsyncDatasetClient:
             ),
             timeout=request_options.get("timeout_in_seconds")
             if request_options is not None and request_options.get("timeout_in_seconds") is not None
-            else 60,
+            else self._client_wrapper.get_timeout(),
             retries=0,
             max_retries=request_options.get("max_retries") if request_options is not None else 0,  # type: ignore
         )
         if 200 <= _response.status_code < 300:
-            return pydantic.parse_obj_as(GetDatasetsResponse, _response.json())  # type: ignore
+            return pydantic_v1.parse_obj_as(GetDatasetsResponse, _response.json())  # type: ignore
         if _response.status_code == 422:
-            raise UnprocessableEntityError(pydantic.parse_obj_as(HttpValidationError, _response.json()))  # type: ignore
+            raise UnprocessableEntityError(
+                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
+            )
         try:
             _response_json = _response.json()
         except JSONDecodeError: