aidial-client 0.11.0.dev4__tar.gz → 0.12.0.dev1__tar.gz

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: aidial-client
-Version: 0.11.0.dev4
+Version: 0.12.0.dev1
 Summary: A Python client library for the AI DIAL API
 License-Expression: Apache-2.0
 License-File: LICENSE
@@ -39,6 +39,7 @@ Description-Content-Type: text/markdown
   - [Authentication](#authentication)
       - [API Keys](#api-keys)
       - [Bearer Token](#bearer-token)
+  - [Lifecycle Management](#lifecycle-management)
   - [Deployments](#deployments)
       - [List Deployments](#list-deployments)
       - [Get Deployment by Id](#get-deployment-by-id)
@@ -143,6 +144,45 @@ async_client = AsyncDial(
 )
 ```
 
+### Lifecycle Management
+
+For deterministic shutdown of underlying HTTP clients, both client types and
+client pools expose lifecycle APIs.
+
+```python
+from aidial_client import AsyncDial, AsyncDialClientPool, Dial, DialClientPool
+
+# Sync client
+with Dial(api_key="your_api_key", base_url="https://your-dial-instance.com") as client:
+    ...
+
+client = Dial(api_key="your_api_key", base_url="https://your-dial-instance.com")
+client.close()
+
+# Async client
+async with AsyncDial(
+    api_key="your_api_key", base_url="https://your-dial-instance.com"
+) as async_client:
+    ...
+
+async_client = AsyncDial(
+    api_key="your_api_key", base_url="https://your-dial-instance.com"
+)
+await async_client.aclose()
+
+# Sync pool
+with DialClientPool() as pool:
+    pooled_client = pool.create_client(
+        base_url="https://your-dial-instance.com", api_key="your-api-key"
+    )
+
+# Async pool
+async with AsyncDialClientPool() as async_pool:
+    pooled_async_client = async_pool.create_client(
+        base_url="https://your-dial-instance.com", api_key="your-api-key"
+    )
+```
+
 You can also pass `bearer_token` as a function without parameters, that returns a `string`:
 
 ```python
@@ -514,6 +554,16 @@ result = await async_client.files.download(
 )
 ```
 
+For large async downloads, use `stream_download()` to process bytes as they arrive without buffering the full response in memory:
+
+```python
+async with async_client.files.stream_download(
+    url=await async_client.my_files_home() / "relative_folder/my-file.txt"
+) as result:
+    async for bytes_chunk in result:
+        ...
+```
+
 As a result, you will receive an object of type `FileDownloadResponse`, that you can iterate by byte chunks:
 
 ```python
@@ -530,6 +580,13 @@ all_content = result.get_content()
 all_content = await result.aget_content()
 ```
 
+or access response metadata:
+
+```python
+headers = result.headers
+content_type = result.content_type
+```
+
 or write it to the file:
 
 ```python
@@ -950,7 +1007,7 @@ second_client = client_pool.create_client(
 #### Asynchronous Client Pool
 
 ```python
-from dial_client import (
+from aidial_client import (
     AsyncDialClientPool,
 )
 

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/README.md

@@ -17,6 +17,7 @@
   - [Authentication](#authentication)
       - [API Keys](#api-keys)
       - [Bearer Token](#bearer-token)
+  - [Lifecycle Management](#lifecycle-management)
   - [Deployments](#deployments)
       - [List Deployments](#list-deployments)
       - [Get Deployment by Id](#get-deployment-by-id)
@@ -121,6 +122,45 @@ async_client = AsyncDial(
 )
 ```
 
+### Lifecycle Management
+
+For deterministic shutdown of underlying HTTP clients, both client types and
+client pools expose lifecycle APIs.
+
+```python
+from aidial_client import AsyncDial, AsyncDialClientPool, Dial, DialClientPool
+
+# Sync client
+with Dial(api_key="your_api_key", base_url="https://your-dial-instance.com") as client:
+    ...
+
+client = Dial(api_key="your_api_key", base_url="https://your-dial-instance.com")
+client.close()
+
+# Async client
+async with AsyncDial(
+    api_key="your_api_key", base_url="https://your-dial-instance.com"
+) as async_client:
+    ...
+
+async_client = AsyncDial(
+    api_key="your_api_key", base_url="https://your-dial-instance.com"
+)
+await async_client.aclose()
+
+# Sync pool
+with DialClientPool() as pool:
+    pooled_client = pool.create_client(
+        base_url="https://your-dial-instance.com", api_key="your-api-key"
+    )
+
+# Async pool
+async with AsyncDialClientPool() as async_pool:
+    pooled_async_client = async_pool.create_client(
+        base_url="https://your-dial-instance.com", api_key="your-api-key"
+    )
+```
+
 You can also pass `bearer_token` as a function without parameters, that returns a `string`:
 
 ```python
@@ -492,6 +532,16 @@ result = await async_client.files.download(
 )
 ```
 
+For large async downloads, use `stream_download()` to process bytes as they arrive without buffering the full response in memory:
+
+```python
+async with async_client.files.stream_download(
+    url=await async_client.my_files_home() / "relative_folder/my-file.txt"
+) as result:
+    async for bytes_chunk in result:
+        ...
+```
+
 As a result, you will receive an object of type `FileDownloadResponse`, that you can iterate by byte chunks:
 
 ```python
@@ -508,6 +558,13 @@ all_content = result.get_content()
 all_content = await result.aget_content()
 ```
 
+or access response metadata:
+
+```python
+headers = result.headers
+content_type = result.content_type
+```
+
 or write it to the file:
 
 ```python
@@ -928,7 +985,7 @@ second_client = client_pool.create_client(
 #### Asynchronous Client Pool
 
 ```python
-from dial_client import (
+from aidial_client import (
     AsyncDialClientPool,
 )
 

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/_client.py

@@ -1,5 +1,6 @@
 from abc import ABC, abstractmethod
 from pathlib import PurePosixPath
+from types import TracebackType
 from typing import Generic, TypeVar
 from urllib.parse import urljoin
 
@@ -164,6 +165,20 @@ class Dial(BaseDialClient[SyncHTTPClient, SyncAuthValue]):
     def auth_headers(self) -> dict[str, str]:
         return self._http_client.auth_headers()
 
+    def close(self) -> None:
+        self._http_client.internal_http_client.close()
+
+    def __enter__(self) -> "Dial":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
 
 class AsyncDial(BaseDialClient[AsyncHTTPClient, AsyncAuthValue]):
     def _init_resources(self) -> None:
@@ -253,3 +268,17 @@ class AsyncDial(BaseDialClient[AsyncHTTPClient, AsyncAuthValue]):
 
     async def auth_headers(self) -> dict[str, str]:
         return await self._http_client.auth_headers()
+
+    async def aclose(self) -> None:
+        await self._http_client.internal_http_client.aclose()
+
+    async def __aenter__(self) -> "AsyncDial":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.aclose()

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/_client_pool.py

@@ -1,3 +1,5 @@
+from types import TracebackType
+
 import httpx
 
 from aidial_client._auth import AsyncAuthValue, SyncAuthValue
@@ -44,6 +46,20 @@ class DialClientPool:
             ),
         )
 
+    def close(self) -> None:
+        self._internal_http_client.close()
+
+    def __enter__(self) -> "DialClientPool":
+        return self
+
+    def __exit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        self.close()
+
 
 class AsyncDialClientPool:
     def __init__(
@@ -78,3 +94,17 @@ class AsyncDialClientPool:
                 internal_http_client=self._internal_http_client,
             ),
         )
+
+    async def aclose(self) -> None:
+        await self._internal_http_client.aclose()
+
+    async def __aenter__(self) -> "AsyncDialClientPool":
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_value: BaseException | None,
+        traceback: TracebackType | None,
+    ) -> None:
+        await self.aclose()

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/_http_client/_async.py

@@ -1,5 +1,5 @@
 import asyncio
-from collections.abc import AsyncIterator, Callable, Mapping
+from collections.abc import AsyncIterator, Mapping
 from contextlib import asynccontextmanager, suppress
 from http import HTTPStatus
 from typing import Any
@@ -8,7 +8,7 @@ import httpx
 
 from aidial_client._auth import AsyncAuthValue, aget_combined_auth_headers
 from aidial_client._exception import DialException
-from aidial_client._http_client._base import BaseHTTPClient
+from aidial_client._http_client._base import BaseHTTPClient, ErrorHandler
 from aidial_client._internal_types._defaults import NOT_GIVEN, NotGiven
 from aidial_client._internal_types._generic import ResponseT
 from aidial_client._internal_types._http_request import FinalRequestOptions
@@ -51,8 +51,7 @@ class AsyncHTTPClient(BaseHTTPClient[httpx.AsyncClient, AsyncAuthValue]):
         options: FinalRequestOptions,
         cast_to: type[ResponseT],
         remaining_retries: int | None = None,
-        on_http_error: Callable[[httpx.HTTPStatusError], DialException | None]
-        | None = None,
+        on_http_error: ErrorHandler | None = None,
     ) -> ResponseT:
         retries = self._remaining_retries(remaining_retries, options)
         auth_headers = await self.auth_headers()
@@ -101,16 +100,37 @@ class AsyncHTTPClient(BaseHTTPClient[httpx.AsyncClient, AsyncAuthValue]):
                     cast_to=cast_to,
                     remaining_retries=retries,
                 )
-            # Try to get a custom error from response status_code/code/message
-            custom_error = on_http_error(err) if on_http_error else None
-            # or fallback to default processing
-            raised_error = custom_error or self._make_dial_error_from_response(
-                err.response
-            )
-            raise raised_error from err
+            self._raise_for_status(response, on_http_error)
 
         return process_block_response(cast_to=cast_to, response=response)
 
+    @asynccontextmanager
+    async def stream(
+        self,
+        *,
+        options: FinalRequestOptions,
+        on_http_error: ErrorHandler | None = None,
+    ) -> AsyncIterator[httpx.Response]:
+        auth_headers = await self.auth_headers()
+        request = self._build_request(options, auth_headers)
+        try:
+            response = await self._internal_http_client.send(
+                request, stream=True
+            )
+        except httpx.TimeoutException as err:
+            raise DialException(
+                message="Request timed out",
+                status_code=HTTPStatus.REQUEST_TIMEOUT,
+            ) from err
+        except httpx.HTTPError as err:
+            raise DialException(message=f"Request failed: {err}") from err
+
+        try:
+            self._raise_for_status(response, on_http_error)
+            yield response
+        finally:
+            await response.aclose()
+
     @asynccontextmanager
     async def stream_sse(
         self,

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/_http_client/_base.py

@@ -1,4 +1,5 @@
 from abc import ABC, abstractmethod
+from collections.abc import Callable
 from http import HTTPStatus
 from random import uniform
 from typing import Generic, TypeVar
@@ -16,6 +17,8 @@ _HttpInternalClientT = TypeVar(
     "_HttpInternalClientT", bound=httpx.Client | httpx.AsyncClient
 )
 
+ErrorHandler = Callable[[httpx.HTTPStatusError], DialException | None]
+
 
 class BaseHTTPClient(ABC, Generic[_HttpInternalClientT, AuthValueT]):
     _internal_http_client: _HttpInternalClientT
@@ -106,6 +109,22 @@ class BaseHTTPClient(ABC, Generic[_HttpInternalClientT, AuthValueT]):
         timeout = sleep_seconds + uniform(-0.5, 0.5)  # noqa: S311
         return max(0, timeout)
 
+    def _raise_for_status(
+        self,
+        response: httpx.Response,
+        on_http_error: ErrorHandler | None,
+    ) -> None:
+        try:
+            response.raise_for_status()
+        except httpx.HTTPStatusError as err:
+            # Try to get a custom error from response status_code/code/message
+            custom_error = on_http_error(err) if on_http_error else None
+            # or fallback to default processing
+            raised_error = custom_error or self._make_dial_error_from_response(
+                err.response
+            )
+            raise raised_error from err
+
     def _make_dial_error_from_response(
         self,
         response: httpx.Response,

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/_http_client/_sync.py

@@ -1,5 +1,5 @@
 import time
-from collections.abc import Callable, Iterator, Mapping
+from collections.abc import Iterator, Mapping
 from contextlib import contextmanager, suppress
 from http import HTTPStatus
 from typing import Any
@@ -8,7 +8,7 @@ import httpx
 
 from aidial_client._auth import SyncAuthValue, get_combined_auth_headers
 from aidial_client._exception import DialException
-from aidial_client._http_client._base import BaseHTTPClient
+from aidial_client._http_client._base import BaseHTTPClient, ErrorHandler
 from aidial_client._internal_types._defaults import NOT_GIVEN, NotGiven
 from aidial_client._internal_types._generic import ResponseT
 from aidial_client._internal_types._http_request import FinalRequestOptions
@@ -50,8 +50,7 @@ class SyncHTTPClient(BaseHTTPClient[httpx.Client, SyncAuthValue]):
         cast_to: type[ResponseT],
         options: FinalRequestOptions,
         remaining_retries: int | None = None,
-        on_http_error: Callable[[httpx.HTTPStatusError], DialException | None]
-        | None = None,
+        on_http_error: ErrorHandler | None = None,
     ) -> ResponseT:
         retries = self._remaining_retries(remaining_retries, options)
         auth_headers = self.auth_headers()
@@ -101,13 +100,7 @@ class SyncHTTPClient(BaseHTTPClient[httpx.Client, SyncAuthValue]):
                     cast_to=cast_to,
                     remaining_retries=retries,
                 )
-            # Try to get a custom error from response status_code/code/message
-            custom_error = on_http_error(err) if on_http_error else None
-            # or fallback to default processing
-            raised_error = custom_error or self._make_dial_error_from_response(
-                err.response
-            )
-            raise raised_error from err
+            self._raise_for_status(response, on_http_error)
 
         return process_block_response(cast_to=cast_to, response=response)
 

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/helpers/storage_resource.py

@@ -5,6 +5,8 @@ from urllib.parse import urljoin, urlparse
 from aidial_client._compatibility.pydantic_v1 import BaseModel
 from aidial_client._constants import API_PREFIX
 from aidial_client._exception import InvalidDialURLError, NotDialURLError
+from aidial_client._internal_types._http_request import FinalRequestOptions
+from aidial_client._utils._dict import remove_none
 from aidial_client.helpers._url import enforce_trailing_slash
 
 StorageResourceType = Literal["files", "conversations", "prompts"]
@@ -156,3 +158,25 @@ class DialStorageResourceMixin(BaseModel):
         Get the display name of the resource from the URL
         """
         return self.get_storage_resource(url).bucket_path
+
+    def _prepare_download_request(
+        self,
+        url: str | PurePosixPath,
+        etag_if_match: str | None,
+    ) -> tuple[FinalRequestOptions, str]:
+        storage_resource = self.get_storage_resource(str(url))
+
+        if storage_resource.filename is None:
+            raise InvalidDialURLError("URL points to a directory, not a file")
+
+        options = FinalRequestOptions(
+            method="GET",
+            url=urljoin(API_PREFIX, storage_resource.api_path),
+            headers=remove_none(
+                {
+                    "If-Match": etag_if_match,
+                }
+            ),
+        )
+
+        return options, storage_resource.filename

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/resources/files.py

@@ -1,3 +1,5 @@
+from collections.abc import AsyncIterator
+from contextlib import asynccontextmanager
 from pathlib import PurePosixPath
 from typing import Literal
 from urllib.parse import urljoin
@@ -8,7 +10,6 @@ from aidial_client._constants import API_PREFIX
 from aidial_client._exception import (
     DialException,
     EtagMismatchError,
-    InvalidDialURLError,
     ResourceNotFoundError,
 )
 from aidial_client._internal_types._generic import NoneType
@@ -70,25 +71,13 @@ class Files(Resource, DialStorageResourceMixin):
         url: str | PurePosixPath,
         etag_if_match: str | None = None,
     ) -> FileDownloadResponse:
-        storage_resource = self.get_storage_resource(str(url))
-        if storage_resource.filename is None:
-            raise InvalidDialURLError("URL points to a directory, not a file")
+        options, filename = self._prepare_download_request(url, etag_if_match)
         response = self.http_client.request(
             cast_to=httpx.Response,
-            options=FinalRequestOptions(
-                method="GET",
-                url=urljoin(API_PREFIX, storage_resource.api_path),
-                headers=remove_none(
-                    {
-                        "If-Match": etag_if_match,
-                    }
-                ),
-            ),
+            options=options,
             on_http_error=_files_error_processor,
         )
-        return FileDownloadResponse(
-            response=response, filename=storage_resource.filename
-        )
+        return FileDownloadResponse(response=response, filename=filename)
 
     def delete(
         self,
@@ -188,25 +177,26 @@ class AsyncFiles(AsyncResource, DialStorageResourceMixin):
         url: str | PurePosixPath,
         etag_if_match: str | None = None,
     ) -> FileDownloadResponse:
-        storage_resource = self.get_storage_resource(str(url))
-        if storage_resource.filename is None:
-            raise InvalidDialURLError("URL points to a directory, not a file")
+        options, filename = self._prepare_download_request(url, etag_if_match)
         response = await self.http_client.request(
             cast_to=httpx.Response,
-            options=FinalRequestOptions(
-                method="GET",
-                url=urljoin(API_PREFIX, storage_resource.api_path),
-                headers=remove_none(
-                    {
-                        "If-Match": etag_if_match,
-                    }
-                ),
-            ),
+            options=options,
             on_http_error=_files_error_processor,
         )
-        return FileDownloadResponse(
-            response=response, filename=storage_resource.filename
-        )
+        return FileDownloadResponse(response=response, filename=filename)
+
+    @asynccontextmanager
+    async def stream_download(
+        self,
+        url: str | PurePosixPath,
+        etag_if_match: str | None = None,
+    ) -> AsyncIterator[FileDownloadResponse]:
+        options, filename = self._prepare_download_request(url, etag_if_match)
+        async with self.http_client.stream(
+            options=options,
+            on_http_error=_files_error_processor,
+        ) as response:
+            yield FileDownloadResponse(response=response, filename=filename)
 
     async def delete(
         self,

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/aidial_client/types/file.py

@@ -40,3 +40,11 @@ class FileDownloadResponse:
     @property
     def filename(self) -> str:
         return self._filename
+
+    @property
+    def headers(self) -> httpx.Headers:
+        return self._response.headers
+
+    @property
+    def content_type(self) -> str | None:
+        return self.headers.get("content-type")

{aidial_client-0.11.0.dev4 → aidial_client-0.12.0.dev1}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "aidial-client"
-version = "0.11.0.dev4"
+version = "0.12.0.dev1"
 description = "A Python client library for the AI DIAL API"
 readme = "README.md"
 license = "Apache-2.0"