athena-intelligence 0.1.81__py3-none-any.whl → 0.1.83__py3-none-any.whl

athena/tools/client.py

@@ -13,9 +13,11 @@ from ..core.request_options import RequestOptions
 from ..errors.internal_server_error import InternalServerError
 from ..errors.not_found_error import NotFoundError
 from ..errors.unprocessable_entity_error import UnprocessableEntityError
+from ..errors.unsupported_media_type_error import UnsupportedMediaTypeError
 from ..types.convert_pdf_to_sheet_out import ConvertPdfToSheetOut
 from ..types.data_frame_parsing_error import DataFrameParsingError
 from ..types.data_frame_request_out import DataFrameRequestOut
+from ..types.data_frame_unknown_format_error import DataFrameUnknownFormatError
 from ..types.excecute_tool_first_workflow_out import ExcecuteToolFirstWorkflowOut
 from ..types.file_fetch_error import FileFetchError
 from ..types.firecrawl_scrape_url_data_reponse_dto import FirecrawlScrapeUrlDataReponseDto
@@ -44,13 +46,22 @@ class ToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> FirecrawlScrapeUrlDataReponseDto:
         """
-        Parameters:
-            - url: str.
+        Parameters
+        ----------
+        url : str
 
-            - params: typing.Optional[typing.Dict[str, typing.Any]].
+        params : typing.Optional[typing.Dict[str, typing.Any]]
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        FirecrawlScrapeUrlDataReponseDto
+            Successful Response
+
+        Examples
+        --------
         from athena.client import Athena
 
         client = Athena(
@@ -58,6 +69,7 @@ class ToolsClient:
         )
         client.tools.scrape_url(
             url="https://www.athenaintelligence.ai",
+            params={"key": "value"},
         )
         """
         _request: typing.Dict[str, typing.Any] = {"url": url}
@@ -110,15 +122,24 @@ class ToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> LangchainDocumentsRequestOut:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
+
+        pagination_limit : typing.Optional[int]
 
-            - pagination_limit: typing.Optional[int].
+        pagination_offset : typing.Optional[int]
 
-            - pagination_offset: typing.Optional[int].
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Returns
+        -------
+        LangchainDocumentsRequestOut
+            Successful Response
+
+        Examples
+        --------
         from athena.client import Athena
 
         client = Athena(
@@ -189,21 +210,33 @@ class ToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DataFrameRequestOut:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
+
+        row_limit : typing.Optional[int]
 
-            - row_limit: typing.Optional[int].
+        index_column : typing.Optional[int]
 
-            - index_column: typing.Optional[int].
+        columns : typing.Optional[typing.Union[ToolsDataFrameRequestColumnsItem, typing.Sequence[ToolsDataFrameRequestColumnsItem]]]
+            should be a list of strings or a list of integers
 
-            - columns: typing.Optional[typing.Union[ToolsDataFrameRequestColumnsItem, typing.Sequence[ToolsDataFrameRequestColumnsItem]]]. should be a list of strings or a list of integers
+        sheet_name : typing.Optional[str]
+            only for excel files
 
-            - sheet_name: typing.Optional[str]. only for excel files
+        separator : typing.Optional[str]
+            only for csv files
 
-            - separator: typing.Optional[str]. only for csv files
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Returns
+        -------
+        DataFrameRequestOut
+            Successful Response
+
+        Examples
+        --------
         from athena.client import Athena
 
         client = Athena(
@@ -251,6 +284,10 @@ class ToolsClient:
             return pydantic_v1.parse_obj_as(DataFrameRequestOut, _response.json())  # type: ignore
         if _response.status_code == 404:
             raise NotFoundError(pydantic_v1.parse_obj_as(FileFetchError, _response.json()))  # type: ignore
+        if _response.status_code == 415:
+            raise UnsupportedMediaTypeError(
+                pydantic_v1.parse_obj_as(DataFrameUnknownFormatError, _response.json())  # type: ignore
+            )
         if _response.status_code == 422:
             raise UnprocessableEntityError(
                 pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
@@ -263,23 +300,34 @@ class ToolsClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
-    def raw_data(self, *, document_id: str, request_options: typing.Optional[RequestOptions] = None) -> None:
+    def raw_data(
+        self, *, document_id: str, request_options: typing.Optional[RequestOptions] = None
+    ) -> typing.Iterator[bytes]:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Yields
+        ------
+        typing.Iterator[bytes]
+            Stream the file in original format.
+
+        Examples
+        --------
         from athena.client import Athena
 
         client = Athena(
             api_key="YOUR_API_KEY",
         )
         client.tools.raw_data(
-            document_id="document_id",
+            document_id="string",
         )
         """
-        _response = self._client_wrapper.httpx_client.request(
+        with self._client_wrapper.httpx_client.stream(
             method="GET",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/tools/file/raw-data"),
             params=jsonable_encoder(
@@ -307,42 +355,54 @@ class ToolsClient:
             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
-        if _response.status_code == 404:
-            raise NotFoundError(pydantic_v1.parse_obj_as(FileFetchError, _response.json()))  # type: ignore
-        if _response.status_code == 422:
-            raise UnprocessableEntityError(
-                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
-            )
-        try:
-            _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)
+        ) as _response:
+            if 200 <= _response.status_code < 300:
+                for _chunk in _response.iter_bytes():
+                    yield _chunk
+                return
+            _response.read()
+            if _response.status_code == 404:
+                raise NotFoundError(pydantic_v1.parse_obj_as(FileFetchError, _response.json()))  # type: ignore
+            if _response.status_code == 422:
+                raise UnprocessableEntityError(
+                    pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
+                )
+            try:
+                _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)
 
     def tool_first_workflow(
         self,
         *,
-        model: typing.Optional[LlmModel] = OMIT,
         tool_name: str,
         content: str,
+        model: typing.Optional[LlmModel] = OMIT,
         tool_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> ExcecuteToolFirstWorkflowOut:
         """
-        Parameters:
-            - model: typing.Optional[LlmModel].
+        Parameters
+        ----------
+        tool_name : str
+
+        content : str
+
+        model : typing.Optional[LlmModel]
 
-            - tool_name: str.
+        tool_kwargs : typing.Optional[typing.Dict[str, typing.Any]]
 
-            - content: str.
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - tool_kwargs: typing.Optional[typing.Dict[str, typing.Any]].
+        Returns
+        -------
+        ExcecuteToolFirstWorkflowOut
+            Successful Response
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Examples
+        --------
         from athena import LlmModel
         from athena.client import Athena
 
@@ -411,21 +471,30 @@ class ToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> ResearcherOut:
         """
-        Parameters:
-            - query: str.
+        Parameters
+        ----------
+        query : str
+
+        max_sections : int
+
+        guidelines : typing.Sequence[str]
 
-            - max_sections: int.
+        publish_formats : typing.Optional[PublishFormats]
 
-            - guidelines: typing.Sequence[str].
+        source : typing.Optional[str]
 
-            - publish_formats: typing.Optional[PublishFormats].
+        athena_document_ids : typing.Optional[typing.Sequence[str]]
 
-            - source: typing.Optional[str].
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - athena_document_ids: typing.Optional[typing.Sequence[str]].
+        Returns
+        -------
+        ResearcherOut
+            Successful Response
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Examples
+        --------
         from athena import PublishFormats
         from athena.client import Athena
 
@@ -502,11 +571,20 @@ class ToolsClient:
         self, *, document_id: str, request_options: typing.Optional[RequestOptions] = None
     ) -> ConvertPdfToSheetOut:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        ConvertPdfToSheetOut
+            Successful Response
+
+        Examples
+        --------
         from athena.client import Athena
 
         client = Athena(
@@ -562,13 +640,22 @@ class ToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> SemanticQueryOut:
         """
-        Parameters:
-            - query: QueryModel.
+        Parameters
+        ----------
+        query : QueryModel
+
+        table_name : typing.Optional[str]
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - table_name: typing.Optional[str].
+        Returns
+        -------
+        SemanticQueryOut
+            Successful Response
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Examples
+        --------
         from athena import FilterModel, FilterOperator, QueryModel, TimeDimensionModel
         from athena.client import Athena
 
@@ -650,13 +737,22 @@ class AsyncToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> FirecrawlScrapeUrlDataReponseDto:
         """
-        Parameters:
-            - url: str.
+        Parameters
+        ----------
+        url : str
 
-            - params: typing.Optional[typing.Dict[str, typing.Any]].
+        params : typing.Optional[typing.Dict[str, typing.Any]]
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        FirecrawlScrapeUrlDataReponseDto
+            Successful Response
+
+        Examples
+        --------
         from athena.client import AsyncAthena
 
         client = AsyncAthena(
@@ -664,6 +760,7 @@ class AsyncToolsClient:
         )
         await client.tools.scrape_url(
             url="https://www.athenaintelligence.ai",
+            params={"key": "value"},
         )
         """
         _request: typing.Dict[str, typing.Any] = {"url": url}
@@ -716,15 +813,24 @@ class AsyncToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> LangchainDocumentsRequestOut:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
+
+        pagination_limit : typing.Optional[int]
 
-            - pagination_limit: typing.Optional[int].
+        pagination_offset : typing.Optional[int]
 
-            - pagination_offset: typing.Optional[int].
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Returns
+        -------
+        LangchainDocumentsRequestOut
+            Successful Response
+
+        Examples
+        --------
         from athena.client import AsyncAthena
 
         client = AsyncAthena(
@@ -795,21 +901,33 @@ class AsyncToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> DataFrameRequestOut:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
+
+        row_limit : typing.Optional[int]
 
-            - row_limit: typing.Optional[int].
+        index_column : typing.Optional[int]
 
-            - index_column: typing.Optional[int].
+        columns : typing.Optional[typing.Union[ToolsDataFrameRequestColumnsItem, typing.Sequence[ToolsDataFrameRequestColumnsItem]]]
+            should be a list of strings or a list of integers
 
-            - columns: typing.Optional[typing.Union[ToolsDataFrameRequestColumnsItem, typing.Sequence[ToolsDataFrameRequestColumnsItem]]]. should be a list of strings or a list of integers
+        sheet_name : typing.Optional[str]
+            only for excel files
 
-            - sheet_name: typing.Optional[str]. only for excel files
+        separator : typing.Optional[str]
+            only for csv files
 
-            - separator: typing.Optional[str]. only for csv files
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Returns
+        -------
+        DataFrameRequestOut
+            Successful Response
+
+        Examples
+        --------
         from athena.client import AsyncAthena
 
         client = AsyncAthena(
@@ -857,6 +975,10 @@ class AsyncToolsClient:
             return pydantic_v1.parse_obj_as(DataFrameRequestOut, _response.json())  # type: ignore
         if _response.status_code == 404:
             raise NotFoundError(pydantic_v1.parse_obj_as(FileFetchError, _response.json()))  # type: ignore
+        if _response.status_code == 415:
+            raise UnsupportedMediaTypeError(
+                pydantic_v1.parse_obj_as(DataFrameUnknownFormatError, _response.json())  # type: ignore
+            )
         if _response.status_code == 422:
             raise UnprocessableEntityError(
                 pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
@@ -869,23 +991,34 @@ class AsyncToolsClient:
             raise ApiError(status_code=_response.status_code, body=_response.text)
         raise ApiError(status_code=_response.status_code, body=_response_json)
 
-    async def raw_data(self, *, document_id: str, request_options: typing.Optional[RequestOptions] = None) -> None:
+    async def raw_data(
+        self, *, document_id: str, request_options: typing.Optional[RequestOptions] = None
+    ) -> typing.AsyncIterator[bytes]:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Yields
+        ------
+        typing.AsyncIterator[bytes]
+            Stream the file in original format.
+
+        Examples
+        --------
         from athena.client import AsyncAthena
 
         client = AsyncAthena(
             api_key="YOUR_API_KEY",
         )
         await client.tools.raw_data(
-            document_id="document_id",
+            document_id="string",
         )
         """
-        _response = await self._client_wrapper.httpx_client.request(
+        async with self._client_wrapper.httpx_client.stream(
             method="GET",
             url=urllib.parse.urljoin(f"{self._client_wrapper.get_base_url()}/", "api/v0/tools/file/raw-data"),
             params=jsonable_encoder(
@@ -913,42 +1046,54 @@ class AsyncToolsClient:
             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
-        if _response.status_code == 404:
-            raise NotFoundError(pydantic_v1.parse_obj_as(FileFetchError, _response.json()))  # type: ignore
-        if _response.status_code == 422:
-            raise UnprocessableEntityError(
-                pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
-            )
-        try:
-            _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)
+        ) as _response:
+            if 200 <= _response.status_code < 300:
+                async for _chunk in _response.aiter_bytes():
+                    yield _chunk
+                return
+            await _response.aread()
+            if _response.status_code == 404:
+                raise NotFoundError(pydantic_v1.parse_obj_as(FileFetchError, _response.json()))  # type: ignore
+            if _response.status_code == 422:
+                raise UnprocessableEntityError(
+                    pydantic_v1.parse_obj_as(HttpValidationError, _response.json())  # type: ignore
+                )
+            try:
+                _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)
 
     async def tool_first_workflow(
         self,
         *,
-        model: typing.Optional[LlmModel] = OMIT,
         tool_name: str,
         content: str,
+        model: typing.Optional[LlmModel] = OMIT,
         tool_kwargs: typing.Optional[typing.Dict[str, typing.Any]] = OMIT,
         request_options: typing.Optional[RequestOptions] = None,
     ) -> ExcecuteToolFirstWorkflowOut:
         """
-        Parameters:
-            - model: typing.Optional[LlmModel].
+        Parameters
+        ----------
+        tool_name : str
+
+        content : str
+
+        model : typing.Optional[LlmModel]
 
-            - tool_name: str.
+        tool_kwargs : typing.Optional[typing.Dict[str, typing.Any]]
 
-            - content: str.
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - tool_kwargs: typing.Optional[typing.Dict[str, typing.Any]].
+        Returns
+        -------
+        ExcecuteToolFirstWorkflowOut
+            Successful Response
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Examples
+        --------
         from athena import LlmModel
         from athena.client import AsyncAthena
 
@@ -1017,21 +1162,30 @@ class AsyncToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> ResearcherOut:
         """
-        Parameters:
-            - query: str.
+        Parameters
+        ----------
+        query : str
+
+        max_sections : int
+
+        guidelines : typing.Sequence[str]
 
-            - max_sections: int.
+        publish_formats : typing.Optional[PublishFormats]
 
-            - guidelines: typing.Sequence[str].
+        source : typing.Optional[str]
 
-            - publish_formats: typing.Optional[PublishFormats].
+        athena_document_ids : typing.Optional[typing.Sequence[str]]
 
-            - source: typing.Optional[str].
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - athena_document_ids: typing.Optional[typing.Sequence[str]].
+        Returns
+        -------
+        ResearcherOut
+            Successful Response
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Examples
+        --------
         from athena import PublishFormats
         from athena.client import AsyncAthena
 
@@ -1108,11 +1262,20 @@ class AsyncToolsClient:
         self, *, document_id: str, request_options: typing.Optional[RequestOptions] = None
     ) -> ConvertPdfToSheetOut:
         """
-        Parameters:
-            - document_id: str.
+        Parameters
+        ----------
+        document_id : str
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
+
+        Returns
+        -------
+        ConvertPdfToSheetOut
+            Successful Response
+
+        Examples
+        --------
         from athena.client import AsyncAthena
 
         client = AsyncAthena(
@@ -1168,13 +1331,22 @@ class AsyncToolsClient:
         request_options: typing.Optional[RequestOptions] = None,
     ) -> SemanticQueryOut:
         """
-        Parameters:
-            - query: QueryModel.
+        Parameters
+        ----------
+        query : QueryModel
+
+        table_name : typing.Optional[str]
+
+        request_options : typing.Optional[RequestOptions]
+            Request-specific configuration.
 
-            - table_name: typing.Optional[str].
+        Returns
+        -------
+        SemanticQueryOut
+            Successful Response
 
-            - request_options: typing.Optional[RequestOptions]. Request-specific configuration.
-        ---
+        Examples
+        --------
         from athena import FilterModel, FilterOperator, QueryModel, TimeDimensionModel
         from athena.client import AsyncAthena