diff --git a/README.md b/README.md index 98ffb61..f94ad61 100644 --- a/README.md +++ b/README.md @@ -508,6 +508,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 diff --git a/aidial_client/types/file.py b/aidial_client/types/file.py index ad7941c..fd81f9e 100644 --- a/aidial_client/types/file.py +++ b/aidial_client/types/file.py @@ -1,5 +1,5 @@ from pathlib import Path -from typing import Union +from typing import Optional, Union import aiofiles import httpx @@ -43,3 +43,11 @@ async def aget_content(self) -> bytes: @property def filename(self) -> str: return self._filename + + @property + def headers(self) -> httpx.Headers: + return self._response.headers + + @property + def content_type(self) -> Optional[str]: + return self.headers.get("content-type") diff --git a/tests/test_types.py b/tests/test_types.py index 0fcea3a..f5f0b95 100644 --- a/tests/test_types.py +++ b/tests/test_types.py @@ -1,7 +1,9 @@ +import httpx import pytest from pydantic import ValidationError from aidial_client.types.chat.response import Attachment +from aidial_client.types.file import FileDownloadResponse from aidial_client.types.metadata import BaseMetadata @@ -54,3 +56,27 @@ def test_metadata_population(): assert getattr(metadata_by_name, field) == getattr( metadata_by_alias, field ) + + +def test_file_download_response_metadata(): + response = httpx.Response( + 200, + content=b"test content", + headers={"content-type": "text/plain", "x-test-header": "test"}, + ) + + download_response = FileDownloadResponse( + response=response, filename="test.txt" + ) + + assert download_response.headers["x-test-header"] == "test" + assert download_response.content_type == "text/plain" + + +def test_file_download_response_metadata_without_content_type(): + response = httpx.Response(200, content=b"test content") + download_response = FileDownloadResponse( + response=response, filename="test.txt" + ) + + assert download_response.content_type is None