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

athena/tools/tasks/client.py

@@ -1,19 +1,19 @@
 # This file was auto-generated by Fern from our API Definition.
 
+from ...core.client_wrapper import SyncClientWrapper
 import typing
+from ...core.request_options import RequestOptions
+from ...core.pydantic_utilities import parse_obj_as
 from json.decoder import JSONDecodeError
-
 from ...core.api_error import ApiError
-from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
-from ...core.pydantic_utilities import pydantic_v1
-from ...core.request_options import RequestOptions
+from ...core.client_wrapper import AsyncClientWrapper
 
 
 class TasksClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
         self._client_wrapper = client_wrapper
 
-    def run_task(self, *, request_options: typing.Optional[RequestOptions] = None) -> typing.Any:
+    def run_task(self, *, request_options: typing.Optional[RequestOptions] = None) -> typing.Optional[typing.Any]:
         """
         Coming soon! Run a [task](https://resources.athenaintel.com/docs/task-studio/home).
 
@@ -24,12 +24,12 @@ class TasksClient:
 
         Returns
         -------
-        typing.Any
+        typing.Optional[typing.Any]
             Successful Response
 
         Examples
         --------
-        from athena.client import Athena
+        from athena import Athena
 
         client = Athena(
             api_key="YOUR_API_KEY",
@@ -37,11 +37,19 @@ class TasksClient:
         client.tools.tasks.run_task()
         """
         _response = self._client_wrapper.httpx_client.request(
-            "api/v0/tools/tasks/run", method="POST", request_options=request_options
+            "api/v0/tools/tasks/run",
+            method="POST",
+            request_options=request_options,
         )
-        if 200 <= _response.status_code < 300:
-            return pydantic_v1.parse_obj_as(typing.Any, _response.json())  # type: ignore
         try:
+            if 200 <= _response.status_code < 300:
+                return typing.cast(
+                    typing.Optional[typing.Any],
+                    parse_obj_as(
+                        type_=typing.Optional[typing.Any],  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
             _response_json = _response.json()
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, body=_response.text)
@@ -52,7 +60,7 @@ class AsyncTasksClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
         self._client_wrapper = client_wrapper
 
-    async def run_task(self, *, request_options: typing.Optional[RequestOptions] = None) -> typing.Any:
+    async def run_task(self, *, request_options: typing.Optional[RequestOptions] = None) -> typing.Optional[typing.Any]:
         """
         Coming soon! Run a [task](https://resources.athenaintel.com/docs/task-studio/home).
 
@@ -63,24 +71,40 @@ class AsyncTasksClient:
 
         Returns
         -------
-        typing.Any
+        typing.Optional[typing.Any]
             Successful Response
 
         Examples
         --------
-        from athena.client import AsyncAthena
+        import asyncio
+
+        from athena import AsyncAthena
 
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
         )
-        await client.tools.tasks.run_task()
+
+
+        async def main() -> None:
+            await client.tools.tasks.run_task()
+
+
+        asyncio.run(main())
         """
         _response = await self._client_wrapper.httpx_client.request(
-            "api/v0/tools/tasks/run", method="POST", request_options=request_options
+            "api/v0/tools/tasks/run",
+            method="POST",
+            request_options=request_options,
         )
-        if 200 <= _response.status_code < 300:
-            return pydantic_v1.parse_obj_as(typing.Any, _response.json())  # type: ignore
         try:
+            if 200 <= _response.status_code < 300:
+                return typing.cast(
+                    typing.Optional[typing.Any],
+                    parse_obj_as(
+                        type_=typing.Optional[typing.Any],  # type: ignore
+                        object_=_response.json(),
+                    ),
+                )
             _response_json = _response.json()
         except JSONDecodeError:
             raise ApiError(status_code=_response.status_code, body=_response.text)

athena/types/asset_node.py

@@ -1,42 +1,32 @@
 # This file was auto-generated by Fern from our API Definition.
 
 from __future__ import annotations
-
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
-
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
+from ..core.pydantic_utilities import update_forward_refs
 
 
-class AssetNode(pydantic_v1.BaseModel):
+class AssetNode(UniversalBaseModel):
     """
     Model representing a node in the folder tree.
     """
 
-    children: typing.Optional[typing.Dict[str, typing.Optional[AssetNode]]] = None
+    children: typing.Optional[typing.Dict[str, typing.Optional["AssetNode"]]] = None
     id: str
     media_type: str
     name: str
     type: str
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow
 
 
-AssetNode.update_forward_refs()
+update_forward_refs(AssetNode)

athena/types/asset_not_found_error.py

@@ -1,29 +1,19 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import typing
+import pydantic
 
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 
-
-class AssetNotFoundError(pydantic_v1.BaseModel):
+class AssetNotFoundError(UniversalBaseModel):
     message: str
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/chunk.py

@@ -1,14 +1,13 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
-
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 from .chunk_content_item import ChunkContentItem
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
 
 
-class Chunk(pydantic_v1.BaseModel):
+class Chunk(UniversalBaseModel):
     """
     A chunk of content to extract data from.
     """
@@ -17,20 +16,11 @@ class Chunk(pydantic_v1.BaseModel):
     content: typing.List[ChunkContentItem]
     metadata: typing.Optional[typing.Dict[str, typing.Optional[str]]] = None
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/chunk_content_item.py

@@ -1,58 +1,38 @@
 # This file was auto-generated by Fern from our API Definition.
 
 from __future__ import annotations
-
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
-
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
 
 
-class ChunkContentItem_Text(pydantic_v1.BaseModel):
-    text: str
+class ChunkContentItem_Text(UniversalBaseModel):
     type: typing.Literal["text"] = "text"
+    text: str
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow
 
 
-class ChunkContentItem_ImageUrl(pydantic_v1.BaseModel):
-    image_url: typing.Dict[str, str]
+class ChunkContentItem_ImageUrl(UniversalBaseModel):
     type: typing.Literal["image_url"] = "image_url"
+    image_url: typing.Dict[str, str]
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow
 
 
 ChunkContentItem = typing.Union[ChunkContentItem_Text, ChunkContentItem_ImageUrl]

athena/types/chunk_result.py

@@ -1,35 +1,25 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
-import typing
-
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
+from ..core.pydantic_utilities import UniversalBaseModel
 from .chunk_result_chunk_id import ChunkResultChunkId
+import typing
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
 
 
-class ChunkResult(pydantic_v1.BaseModel):
+class ChunkResult(UniversalBaseModel):
     """
     The result of a chunk extraction.
     """
 
     chunk_id: ChunkResultChunkId
-    chunk_result: typing.Optional[typing.Dict[str, typing.Any]] = None
-
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
+    chunk_result: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
 
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/custom_agent_response.py

@@ -1,32 +1,22 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
 
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 
-
-class CustomAgentResponse(pydantic_v1.BaseModel):
-    result: typing.Dict[str, typing.Any] = pydantic_v1.Field()
+class CustomAgentResponse(UniversalBaseModel):
+    result: typing.Dict[str, typing.Optional[typing.Any]] = pydantic.Field()
     """
     The agent's response. Format depends on the specific agent implementation.
     """
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/data_frame_request_out.py

@@ -1,16 +1,15 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
-
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 from .data_frame_request_out_columns_item import DataFrameRequestOutColumnsItem
 from .data_frame_request_out_data_item_item import DataFrameRequestOutDataItemItem
 from .data_frame_request_out_index_item import DataFrameRequestOutIndexItem
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
 
 
-class DataFrameRequestOut(pydantic_v1.BaseModel):
+class DataFrameRequestOut(UniversalBaseModel):
     """
     Response model with JSON dataframe representation.
     """
@@ -19,20 +18,11 @@ class DataFrameRequestOut(pydantic_v1.BaseModel):
     data: typing.List[typing.List[typing.Optional[DataFrameRequestOutDataItemItem]]]
     index: typing.Optional[typing.List[typing.Optional[DataFrameRequestOutIndexItem]]] = None
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/data_frame_unknown_format_error.py

@@ -1,31 +1,21 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
 import typing
+import pydantic
 
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 
-
-class DataFrameUnknownFormatError(pydantic_v1.BaseModel):
+class DataFrameUnknownFormatError(UniversalBaseModel):
     asset_id: str
     media_type: str
     message: str
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/document_chunk.py

@@ -1,34 +1,24 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
+import pydantic
 
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 
-
-class DocumentChunk(pydantic_v1.BaseModel):
+class DocumentChunk(UniversalBaseModel):
     """
     A document chunk.
     """
 
-    metadata: typing.Optional[typing.Dict[str, typing.Any]] = None
+    metadata: typing.Optional[typing.Dict[str, typing.Optional[typing.Any]]] = None
     page_content: str
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow

athena/types/drive_agent_response.py

@@ -1,32 +1,22 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import datetime as dt
+from ..core.pydantic_utilities import UniversalBaseModel
 import typing
+import pydantic
+from ..core.pydantic_utilities import IS_PYDANTIC_V2
 
-from ..core.datetime_utils import serialize_datetime
-from ..core.pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
 
-
-class DriveAgentResponse(pydantic_v1.BaseModel):
-    result: typing.Dict[str, typing.Any] = pydantic_v1.Field()
+class DriveAgentResponse(UniversalBaseModel):
+    result: typing.Dict[str, typing.Optional[typing.Any]] = pydantic.Field()
     """
     Results of the drive operation
     """
 
-    def json(self, **kwargs: typing.Any) -> str:
-        kwargs_with_defaults: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        return super().json(**kwargs_with_defaults)
-
-    def dict(self, **kwargs: typing.Any) -> typing.Dict[str, typing.Any]:
-        kwargs_with_defaults_exclude_unset: typing.Any = {"by_alias": True, "exclude_unset": True, **kwargs}
-        kwargs_with_defaults_exclude_none: typing.Any = {"by_alias": True, "exclude_none": True, **kwargs}
-
-        return deep_union_pydantic_dicts(
-            super().dict(**kwargs_with_defaults_exclude_unset), super().dict(**kwargs_with_defaults_exclude_none)
-        )
+    if IS_PYDANTIC_V2:
+        model_config: typing.ClassVar[pydantic.ConfigDict] = pydantic.ConfigDict(extra="allow", frozen=True)  # type: ignore # Pydantic v2
+    else:
 
-    class Config:
-        frozen = True
-        smart_union = True
-        extra = pydantic_v1.Extra.allow
-        json_encoders = {dt.datetime: serialize_datetime}
+        class Config:
+            frozen = True
+            smart_union = True
+            extra = pydantic.Extra.allow