athena-intelligence 0.1.127__py3-none-any.whl → 0.1.185__py3-none-any.whl

athena/agents/general/client.py

@@ -1,16 +1,12 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-from ...core.client_wrapper import SyncClientWrapper
-from ...types.general_agent_request import GeneralAgentRequest
+
+from ...core.client_wrapper import AsyncClientWrapper, SyncClientWrapper
 from ...core.request_options import RequestOptions
+from ...types.general_agent_request import GeneralAgentRequest
 from ...types.general_agent_response import GeneralAgentResponse
-from ...core.serialization import convert_and_respect_annotation_metadata
-from ...core.pydantic_utilities import parse_obj_as
-from ...errors.unprocessable_entity_error import UnprocessableEntityError
-from json.decoder import JSONDecodeError
-from ...core.api_error import ApiError
-from ...core.client_wrapper import AsyncClientWrapper
+from .raw_client import AsyncRawGeneralClient, RawGeneralClient
 
 # this is used as the default value for optional parameters
 OMIT = typing.cast(typing.Any, ...)
@@ -18,7 +14,18 @@ OMIT = typing.cast(typing.Any, ...)
 
 class GeneralClient:
     def __init__(self, *, client_wrapper: SyncClientWrapper):
-        self._client_wrapper = client_wrapper
+        self._raw_client = RawGeneralClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> RawGeneralClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        RawGeneralClient
+        """
+        return self._raw_client
 
     def batch(
         self, *, request: typing.Sequence[GeneralAgentRequest], request_options: typing.Optional[RequestOptions] = None
@@ -40,7 +47,7 @@ class GeneralClient:
 
         Examples
         --------
-        from athena import Athena, GeneralAgentConfig, GeneralAgentRequest
+        from athena import Athena, GeneralAgentConfig, GeneralAgentRequest, InputMessage
 
         client = Athena(
             api_key="YOUR_API_KEY",
@@ -49,50 +56,20 @@ class GeneralClient:
             request=[
                 GeneralAgentRequest(
                     config=GeneralAgentConfig(
-                        enabled_tools=["search"],
+                        enabled_tools=[],
                     ),
                     messages=[
-                        {
-                            "content": "Please call the search tool for AAPL news.",
-                            "type": "user",
-                        }
+                        InputMessage(
+                            content="Please call the search tool for AAPL news.",
+                            role="user",
+                        )
                     ],
                 )
             ],
         )
         """
-        _response = self._client_wrapper.httpx_client.request(
-            "api/v0/agents/general/batch",
-            method="POST",
-            json=convert_and_respect_annotation_metadata(
-                object_=request, annotation=typing.Sequence[GeneralAgentRequest], direction="write"
-            ),
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    typing.List[GeneralAgentResponse],
-                    parse_obj_as(
-                        type_=typing.List[GeneralAgentResponse],  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            if _response.status_code == 422:
-                raise UnprocessableEntityError(
-                    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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = self._raw_client.batch(request=request, request_options=request_options)
+        return _response.data
 
     def invoke(
         self, *, request: GeneralAgentRequest, request_options: typing.Optional[RequestOptions] = None
@@ -117,7 +94,7 @@ class GeneralClient:
 
         Examples
         --------
-        from athena import Athena, GeneralAgentConfig, GeneralAgentRequest
+        from athena import Athena, GeneralAgentConfig, GeneralAgentRequest, InputMessage
 
         client = Athena(
             api_key="YOUR_API_KEY",
@@ -125,49 +102,19 @@ class GeneralClient:
         client.agents.general.invoke(
             request=GeneralAgentRequest(
                 config=GeneralAgentConfig(
-                    enabled_tools=["search"],
+                    enabled_tools=[],
                 ),
                 messages=[
-                    {
-                        "content": "Please call the search tool for AAPL news.",
-                        "type": "user",
-                    }
+                    InputMessage(
+                        content="Please call the search tool for AAPL news.",
+                        role="user",
+                    )
                 ],
             ),
         )
         """
-        _response = self._client_wrapper.httpx_client.request(
-            "api/v0/agents/general/invoke",
-            method="POST",
-            json=convert_and_respect_annotation_metadata(
-                object_=request, annotation=GeneralAgentRequest, direction="write"
-            ),
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    GeneralAgentResponse,
-                    parse_obj_as(
-                        type_=GeneralAgentResponse,  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            if _response.status_code == 422:
-                raise UnprocessableEntityError(
-                    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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = self._raw_client.invoke(request=request, request_options=request_options)
+        return _response.data
 
     def stream_events(
         self, *, request: GeneralAgentRequest, request_options: typing.Optional[RequestOptions] = None
@@ -189,7 +136,7 @@ class GeneralClient:
 
         Examples
         --------
-        from athena import Athena, GeneralAgentConfig, GeneralAgentRequest
+        from athena import Athena, GeneralAgentConfig, GeneralAgentRequest, InputMessage
 
         client = Athena(
             api_key="YOUR_API_KEY",
@@ -197,54 +144,35 @@ class GeneralClient:
         client.agents.general.stream_events(
             request=GeneralAgentRequest(
                 config=GeneralAgentConfig(
-                    enabled_tools=["search"],
+                    enabled_tools=[],
                 ),
                 messages=[
-                    {
-                        "content": "Please call the search tool for AAPL news.",
-                        "type": "user",
-                    }
+                    InputMessage(
+                        content="Please call the search tool for AAPL news.",
+                        role="user",
+                    )
                 ],
             ),
         )
         """
-        _response = self._client_wrapper.httpx_client.request(
-            "api/v0/agents/general/stream_events",
-            method="POST",
-            json=convert_and_respect_annotation_metadata(
-                object_=request, annotation=GeneralAgentRequest, direction="write"
-            ),
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    GeneralAgentResponse,
-                    parse_obj_as(
-                        type_=GeneralAgentResponse,  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            if _response.status_code == 422:
-                raise UnprocessableEntityError(
-                    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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = self._raw_client.stream_events(request=request, request_options=request_options)
+        return _response.data
 
 
 class AsyncGeneralClient:
     def __init__(self, *, client_wrapper: AsyncClientWrapper):
-        self._client_wrapper = client_wrapper
+        self._raw_client = AsyncRawGeneralClient(client_wrapper=client_wrapper)
+
+    @property
+    def with_raw_response(self) -> AsyncRawGeneralClient:
+        """
+        Retrieves a raw implementation of this client that returns raw responses.
+
+        Returns
+        -------
+        AsyncRawGeneralClient
+        """
+        return self._raw_client
 
     async def batch(
         self, *, request: typing.Sequence[GeneralAgentRequest], request_options: typing.Optional[RequestOptions] = None
@@ -268,7 +196,12 @@ class AsyncGeneralClient:
         --------
         import asyncio
 
-        from athena import AsyncAthena, GeneralAgentConfig, GeneralAgentRequest
+        from athena import (
+            AsyncAthena,
+            GeneralAgentConfig,
+            GeneralAgentRequest,
+            InputMessage,
+        )
 
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
@@ -280,13 +213,13 @@ class AsyncGeneralClient:
                 request=[
                     GeneralAgentRequest(
                         config=GeneralAgentConfig(
-                            enabled_tools=["search"],
+                            enabled_tools=[],
                         ),
                         messages=[
-                            {
-                                "content": "Please call the search tool for AAPL news.",
-                                "type": "user",
-                            }
+                            InputMessage(
+                                content="Please call the search tool for AAPL news.",
+                                role="user",
+                            )
                         ],
                     )
                 ],
@@ -295,38 +228,8 @@ class AsyncGeneralClient:
 
         asyncio.run(main())
         """
-        _response = await self._client_wrapper.httpx_client.request(
-            "api/v0/agents/general/batch",
-            method="POST",
-            json=convert_and_respect_annotation_metadata(
-                object_=request, annotation=typing.Sequence[GeneralAgentRequest], direction="write"
-            ),
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    typing.List[GeneralAgentResponse],
-                    parse_obj_as(
-                        type_=typing.List[GeneralAgentResponse],  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            if _response.status_code == 422:
-                raise UnprocessableEntityError(
-                    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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = await self._raw_client.batch(request=request, request_options=request_options)
+        return _response.data
 
     async def invoke(
         self, *, request: GeneralAgentRequest, request_options: typing.Optional[RequestOptions] = None
@@ -353,7 +256,12 @@ class AsyncGeneralClient:
         --------
         import asyncio
 
-        from athena import AsyncAthena, GeneralAgentConfig, GeneralAgentRequest
+        from athena import (
+            AsyncAthena,
+            GeneralAgentConfig,
+            GeneralAgentRequest,
+            InputMessage,
+        )
 
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
@@ -364,13 +272,13 @@ class AsyncGeneralClient:
             await client.agents.general.invoke(
                 request=GeneralAgentRequest(
                     config=GeneralAgentConfig(
-                        enabled_tools=["search"],
+                        enabled_tools=[],
                     ),
                     messages=[
-                        {
-                            "content": "Please call the search tool for AAPL news.",
-                            "type": "user",
-                        }
+                        InputMessage(
+                            content="Please call the search tool for AAPL news.",
+                            role="user",
+                        )
                     ],
                 ),
             )
@@ -378,38 +286,8 @@ class AsyncGeneralClient:
 
         asyncio.run(main())
         """
-        _response = await self._client_wrapper.httpx_client.request(
-            "api/v0/agents/general/invoke",
-            method="POST",
-            json=convert_and_respect_annotation_metadata(
-                object_=request, annotation=GeneralAgentRequest, direction="write"
-            ),
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    GeneralAgentResponse,
-                    parse_obj_as(
-                        type_=GeneralAgentResponse,  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            if _response.status_code == 422:
-                raise UnprocessableEntityError(
-                    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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = await self._raw_client.invoke(request=request, request_options=request_options)
+        return _response.data
 
     async def stream_events(
         self, *, request: GeneralAgentRequest, request_options: typing.Optional[RequestOptions] = None
@@ -433,7 +311,12 @@ class AsyncGeneralClient:
         --------
         import asyncio
 
-        from athena import AsyncAthena, GeneralAgentConfig, GeneralAgentRequest
+        from athena import (
+            AsyncAthena,
+            GeneralAgentConfig,
+            GeneralAgentRequest,
+            InputMessage,
+        )
 
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
@@ -444,13 +327,13 @@ class AsyncGeneralClient:
             await client.agents.general.stream_events(
                 request=GeneralAgentRequest(
                     config=GeneralAgentConfig(
-                        enabled_tools=["search"],
+                        enabled_tools=[],
                     ),
                     messages=[
-                        {
-                            "content": "Please call the search tool for AAPL news.",
-                            "type": "user",
-                        }
+                        InputMessage(
+                            content="Please call the search tool for AAPL news.",
+                            role="user",
+                        )
                     ],
                 ),
             )
@@ -458,35 +341,5 @@ class AsyncGeneralClient:
 
         asyncio.run(main())
         """
-        _response = await self._client_wrapper.httpx_client.request(
-            "api/v0/agents/general/stream_events",
-            method="POST",
-            json=convert_and_respect_annotation_metadata(
-                object_=request, annotation=GeneralAgentRequest, direction="write"
-            ),
-            request_options=request_options,
-            omit=OMIT,
-        )
-        try:
-            if 200 <= _response.status_code < 300:
-                return typing.cast(
-                    GeneralAgentResponse,
-                    parse_obj_as(
-                        type_=GeneralAgentResponse,  # type: ignore
-                        object_=_response.json(),
-                    ),
-                )
-            if _response.status_code == 422:
-                raise UnprocessableEntityError(
-                    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)
-        raise ApiError(status_code=_response.status_code, body=_response_json)
+        _response = await self._raw_client.stream_events(request=request, request_options=request_options)
+        return _response.data