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

athena/client.py

@@ -7,8 +7,9 @@ import typing
 import warnings
 
 import httpx
-from typing import cast, List, Union
+from typing import cast, List, Tuple, Union, Optional
 from typing_extensions import TypeVar, ParamSpec
+from langserve import RemoteRunnable
 
 from . import core
 from .base_client import BaseAthena, AsyncBaseAthena
@@ -24,23 +25,23 @@ if typing.TYPE_CHECKING:
 
 P = ParamSpec("P")
 T = TypeVar("T")
-U = TypeVar('U')
+U = TypeVar("U")
 
 
 def _inherit_signature_and_doc(
-    f: typing.Callable[P, T],
-    replace_in_doc: typing.Dict[str, str]
+    f: typing.Callable[P, T], replace_in_doc: typing.Dict[str, str]
 ) -> typing.Callable[..., typing.Callable[P, U]]:
     def decorator(decorated):
         for old, new in replace_in_doc.items():
             assert old in f.__doc__
             decorated.__doc__ = f.__doc__.replace(old, new)
         return decorated
+
     return decorator
 
 
 class SpecialEnvironments(enum.Enum):
-    AUTODETECT_ENVIRONMENT = 'AUTO'
+    AUTODETECT_ENVIRONMENT = "AUTO"
 
 
 @dataclasses.dataclass
@@ -53,7 +54,7 @@ class AthenaAsset:
         if self.media_type == "application/sql":
             # it is safe to import IPython in `_repr_mimebundle_`
             # as this is only intended to be invoked by IPython.
-            from IPython import display    # type: ignore[import]
+            from IPython import display  # type: ignore[import]
 
             code = display.Code(
                 data=self.data.decode(),
@@ -84,17 +85,19 @@ class WrappedToolsClient(ToolsClient):
         file_io = client.tools.get_file(asset_id="asset_id")
         pl.read_csv(file_io)
         """
-        file_bytes = b''.join(self.raw_data(asset_id=asset_id))
+        file_bytes = b"".join(self.raw_data(asset_id=asset_id))
         bytes_io = io.BytesIO(file_bytes)
         return bytes_io
 
-    @_inherit_signature_and_doc(ToolsClient.data_frame, {'DataFrameRequestOut': 'pd.DataFrame'})
-    def data_frame(self, *, asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        ToolsClient.data_frame, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    def data_frame(self, *, asset_id: str, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         model = super().data_frame(asset_id=asset_id, **kwargs)
         return _read_json_frame(model)
 
-    def read_data_frame(self, asset_id: str, *args, **kwargs) -> 'pd.DataFrame':
+    def read_data_frame(self, asset_id: str, *args, **kwargs) -> "pd.DataFrame":
         """
         Parameters
         ----------
@@ -116,11 +119,10 @@ class WrappedToolsClient(ToolsClient):
         client.tools.read_data_frame(asset_id="asset_id")
         """
         _check_pandas_installed()
-        file_bytes = self.get_file(asset_id)
-        return _to_pandas_df(file_bytes, *args, **kwargs)
+        file_bytes, media_type = self._get_file_and_media_type(asset_id=asset_id)
+        return _to_pandas_df(file_bytes, *args, media_type=media_type, **kwargs)
 
-
-    def save_asset(    # type: ignore[override]
+    def save_asset(  # type: ignore[override]
         self,
         asset_object: Union["pd.DataFrame", "pd.Series", core.File],
         *,
@@ -131,7 +133,7 @@ class WrappedToolsClient(ToolsClient):
         """
         Parameters
         ----------
-        asset_object : pd.DataFrame | pd.Series | core.File
+        asset_object : pd.DataFrame | pd.Series | matplotlib.figure.Figure | core.File
             A pandas data frame, series, matplotlib figure, or core.File
 
         parent_folder_id : typing.Optional[str]
@@ -143,6 +145,8 @@ class WrappedToolsClient(ToolsClient):
         request_options : typing.Optional[RequestOptions]
             Request-specific configuration.
 
+        **kwargs : passed down to conversion methods
+
         Returns
         -------
         SaveAssetRequestOut
@@ -155,27 +159,14 @@ class WrappedToolsClient(ToolsClient):
         client = Athena(api_key="YOUR_API_KEY")
         client.tools.save_asset(df)
         """
-        asset_object = _convert_asset_object(asset_object=asset_object, name=name)
-        return super().save_asset(
-            file=asset_object, parent_folder_id=parent_folder_id, **kwargs
+        asset_object = _convert_asset_object(
+            asset_object=asset_object, name=name, **kwargs
         )
+        return super().save_asset(file=asset_object, parent_folder_id=parent_folder_id)
 
-    def get_asset(self, asset_id: str) -> Union["pd.DataFrame", AthenaAsset]:
+    def _get_file_and_media_type(self, asset_id: str) -> Tuple[io.BytesIO, str]:
         """
-        Parameters
-        ----------
-        asset_id : str
-
-        Returns
-        -------
-        pd.DataFrame or AthenaAsset
-
-        Examples
-        --------
-        from athena.client import Athena
-
-        client = Athena(api_key="YOUR_API_KEY")
-        client.tools.get_asset(asset_id="asset_id")
+        Gets the file togehter with media type returned by server
         """
         # while we wait for https://github.com/fern-api/fern/issues/4316
         result = self._client_wrapper.httpx_client.request(
@@ -196,6 +187,27 @@ class WrappedToolsClient(ToolsClient):
             # fallback to `libmagic` inference
             media_type = _infer_media_type(bytes_io=file_bytes)
 
+        return file_bytes, media_type
+
+    def get_asset(self, asset_id: str) -> Union["pd.DataFrame", AthenaAsset]:
+        """
+        Parameters
+        ----------
+        asset_id : str
+
+        Returns
+        -------
+        pd.DataFrame or AthenaAsset
+
+        Examples
+        --------
+        from athena.client import Athena
+
+        client = Athena(api_key="YOUR_API_KEY")
+        client.tools.get_asset(asset_id="asset_id")
+        """
+        file_bytes, media_type = self._get_file_and_media_type(asset_id=asset_id)
+
         media_type_aliases = {"image/jpg": "image/jpeg"}
         media_type = media_type_aliases.get(media_type, media_type)
 
@@ -226,22 +238,26 @@ class WrappedToolsClient(ToolsClient):
             return _to_pandas_df(file_bytes, media_type=media_type)
 
         raise NotImplementedError("Assets of `{media_type}` type are not yet supported")
-    
+
 
 class WrappedQueryClient(QueryClient):
 
-    @_inherit_signature_and_doc(QueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
-    def execute(self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        QueryClient.execute, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    def execute(
+        self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs
+    ) -> "pd.DataFrame":
         _check_pandas_installed()
         model = super().execute(
-            sql_command=sql_command,
-            database_asset_ids=database_asset_ids,
-            **kwargs
+            sql_command=sql_command, database_asset_ids=database_asset_ids, **kwargs
         )
         return _read_json_frame(model)
 
-    @_inherit_signature_and_doc(QueryClient.execute_snippet, {'DataFrameRequestOut': 'pd.DataFrame'})
-    def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        QueryClient.execute_snippet, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         model = super().execute_snippet(snippet_asset_id=snippet_asset_id, **kwargs)
         return _read_json_frame(model)
@@ -251,12 +267,11 @@ def _add_docs_for_async_variant(obj):
     def decorator(decorated):
         doc = obj.__doc__
         name = obj.__name__
-        decorated.__doc__ = (
-            doc
-            .replace('client = Athena', 'client = AsyncAthena')
-            .replace(f'client.tools.{name}', f'await client.tools.{name}')
-        )
+        decorated.__doc__ = doc.replace(
+            "client = Athena", "client = AsyncAthena"
+        ).replace(f"client.tools.{name}", f"await client.tools.{name}")
         return decorated
+
     return decorator
 
 
@@ -264,18 +279,20 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
 
     @_add_docs_for_async_variant(WrappedToolsClient.get_file)
     async def get_file(self, asset_id: str) -> io.BytesIO:
-        file_bytes = b''.join([gen async for gen in self.raw_data(asset_id=asset_id)])
+        file_bytes = b"".join([gen async for gen in self.raw_data(asset_id=asset_id)])
         bytes_io = io.BytesIO(file_bytes)
         return bytes_io
 
-    @_inherit_signature_and_doc(AsyncToolsClient.data_frame, {'DataFrameRequestOut': 'pd.DataFrame'})
-    async def data_frame(self, *, asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        AsyncToolsClient.data_frame, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    async def data_frame(self, *, asset_id: str, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         model = await super().data_frame(asset_id=asset_id, **kwargs)
         return _read_json_frame(model)
 
     @_add_docs_for_async_variant(WrappedToolsClient.read_data_frame)
-    async def read_data_frame(self, asset_id: str, *args, **kwargs) -> 'pd.DataFrame':
+    async def read_data_frame(self, asset_id: str, *args, **kwargs) -> "pd.DataFrame":
         _check_pandas_installed()
         file_bytes = await self.get_file(asset_id)
         return _to_pandas_df(file_bytes, *args, **kwargs)
@@ -289,31 +306,119 @@ class WrappedAsyncToolsClient(AsyncToolsClient):
         name: Union[str, None] = None,
         **kwargs,
     ) -> SaveAssetRequestOut:
-        asset_object = _convert_asset_object(asset_object=asset_object, name=name)
+        asset_object = _convert_asset_object(
+            asset_object=asset_object, name=name, **kwargs
+        )
         return await super().save_asset(
-            file=asset_object, parent_folder_id=parent_folder_id, **kwargs
+            file=asset_object, parent_folder_id=parent_folder_id
         )
 
 
 class WrappedAsyncQueryClient(AsyncQueryClient):
 
-    @_inherit_signature_and_doc(AsyncQueryClient.execute, {'DataFrameRequestOut': 'pd.DataFrame'})
-    async def execute(self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        AsyncQueryClient.execute, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    async def execute(
+        self, *, sql_command: str, database_asset_ids: Union[str, List[str]], **kwargs
+    ) -> "pd.DataFrame":
         _check_pandas_installed()
         model = await super().execute(
-            sql_command=sql_command,
-            database_asset_ids=database_asset_ids,
-            **kwargs
+            sql_command=sql_command, database_asset_ids=database_asset_ids, **kwargs
         )
         return _read_json_frame(model)
 
-    @_inherit_signature_and_doc(AsyncQueryClient.execute_snippet, {'DataFrameRequestOut': 'pd.DataFrame'})
-    async def execute_snippet(self, *, snippet_asset_id: str, **kwargs) -> 'pd.DataFrame':
+    @_inherit_signature_and_doc(
+        AsyncQueryClient.execute_snippet, {"DataFrameRequestOut": "pd.DataFrame"}
+    )
+    async def execute_snippet(
+        self, *, snippet_asset_id: str, **kwargs
+    ) -> "pd.DataFrame":
         _check_pandas_installed()
-        model = await super().execute_snippet(snippet_asset_id=snippet_asset_id, **kwargs)
+        model = await super().execute_snippet(
+            snippet_asset_id=snippet_asset_id, **kwargs
+        )
         return _read_json_frame(model)
 
 
+class AthenaModel(RemoteRunnable):
+    """Use Athena's models directly with a Langchain-compatible client.
+
+    The Langchain Runnable interface is supported:
+    - `invoke`: Invoke the model with a string or a Langchain message.
+    - `batch`: Batch invoke the model on multiple inputs.
+    - `astream_events`: Streaming for real-time applications.
+
+    See Langchain documentation for more details.
+
+    Examples
+    --------
+    from src.athena.client import Athena
+    import asyncio
+    from langchain_core.messages import HumanMessage
+
+    client = Athena()
+    llm = client.llm
+
+    # sync invoke -- use strings or langchain messages
+    result = llm.invoke("Hello")
+    print(result.content)
+
+    result = llm.invoke("Hello")
+    print(result.content)
+
+    # choose the model explicitly
+    claude = llm.with_config(configurable={"model": "claude_3_7_sonnet"})
+    print(claude.invoke("Who are you?").content)
+
+    # batch (for multiple parallel requests)
+    results = llm.batch(["Hello", "World"] * 5)
+    print([r.content for r in results])
+
+    # handle stream events
+    async def stream_response():
+        separator = "---------------------------------------"
+        print(separator)
+        print("Starting stream")
+        print(separator)
+        async for event in llm.astream_events("Hello.  Please respond in 5 sentences."):
+            data = event["data"]
+            if "chunk" in data:
+                print(data["chunk"].content)
+            elif "output" in data:
+                print(separator)
+                print("Final response")
+                print(separator)
+                print(data["output"].content)
+            else:
+                print("(other event) ", event.keys())
+
+
+    asyncio.run(stream_response())
+    from athena.client import Athena
+    client = Athena()
+    llm = client.llm
+
+    # sync invoke -- use strings or langchain messages
+    result = llm.invoke("Hello")
+    print(result.content)
+
+    # batch
+    results = llm.batch(["Hello", "World"])
+    print(r.content for r in results)
+    """
+
+    def __init__(self, base_url: str, api_key: str, timeout: Optional[float] = None):
+        self.base_url = base_url
+        self.api_key = api_key
+        self.timeout = timeout
+        super().__init__(
+            base_url + "/api/models/general",
+            headers={"X-API-KEY": api_key},
+            timeout=timeout,
+        )
+
+
 class Athena(BaseAthena):
     """
     Use this class to access the different functions within the SDK. You can instantiate any number of clients with different configuration that will propogate to these functions.
@@ -350,37 +455,40 @@ class Athena(BaseAthena):
         self,
         *,
         base_url: typing.Optional[str] = None,
-        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,    # type: ignore[arg-type]
+        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,  # type: ignore[arg-type]
         api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
-        httpx_client: typing.Optional[httpx.Client] = None
+        httpx_client: typing.Optional[httpx.Client] = None,
     ):
         if api_key is None:
             try:
-                api_key = os.environ['ATHENA_API_KEY']
+                api_key = os.environ["ATHENA_API_KEY"]
             except KeyError:
                 raise TypeError(
                     "Athena() missing 1 required keyword-only argument: 'api_key'"
-                    ' (ATHENA_API_KEY environment variable not found)'
+                    " (ATHENA_API_KEY environment variable not found)"
                 )
         if environment == SpecialEnvironments.AUTODETECT_ENVIRONMENT:
-            if 'ATHENA_API_URL' in os.environ:
+            if "ATHENA_API_URL" in os.environ:
 
                 class AutodetectedEnvironments(enum.Enum):
-                    CURRENT = os.environ['ATHENA_API_URL']
+                    CURRENT = os.environ["ATHENA_API_URL"]
 
                 environment = cast(AthenaEnvironment, AutodetectedEnvironments.CURRENT)
             else:
                 environment = AthenaEnvironment.PRODUCTION
         super().__init__(
             base_url=base_url,
-            environment=environment,    # type: ignore[arg-type]
+            environment=environment,  # type: ignore[arg-type]
             api_key=api_key,
             timeout=timeout,
             httpx_client=httpx_client,
         )
         self.tools = WrappedToolsClient(client_wrapper=self._client_wrapper)
         self.query = WrappedQueryClient(client_wrapper=self._client_wrapper)
+        self.llm = AthenaModel(
+            base_url=base_url or environment.value, api_key=api_key, timeout=timeout
+        )
 
 
 class AsyncAthena(AsyncBaseAthena):
@@ -420,31 +528,31 @@ class AsyncAthena(AsyncBaseAthena):
         self,
         *,
         base_url: typing.Optional[str] = None,
-        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,    # type: ignore[arg-type]
+        environment: Union[AthenaEnvironment, SpecialEnvironments] = SpecialEnvironments.AUTODETECT_ENVIRONMENT,  # type: ignore[arg-type]
         api_key: typing.Optional[str] = None,
         timeout: typing.Optional[float] = 60,
-        httpx_client: typing.Optional[httpx.AsyncClient] = None
+        httpx_client: typing.Optional[httpx.AsyncClient] = None,
     ):
         if api_key is None:
             try:
-                api_key = os.environ['ATHENA_API_KEY']
+                api_key = os.environ["ATHENA_API_KEY"]
             except KeyError:
                 raise TypeError(
                     "AsyncAthena() missing 1 required keyword-only argument: 'api_key'"
-                    ' (ATHENA_API_KEY environment variable not found)'
+                    " (ATHENA_API_KEY environment variable not found)"
                 )
         if environment == SpecialEnvironments.AUTODETECT_ENVIRONMENT:
-            if 'ATHENA_API_URL' in os.environ:
+            if "ATHENA_API_URL" in os.environ:
 
                 class AutodetectedEnvironments(enum.Enum):
-                    CURRENT = os.environ['ATHENA_API_URL']
+                    CURRENT = os.environ["ATHENA_API_URL"]
 
                 environment = cast(AthenaEnvironment, AutodetectedEnvironments.CURRENT)
             else:
                 environment = AthenaEnvironment.PRODUCTION
         super().__init__(
             base_url=base_url,
-            environment=environment,    # type: ignore[arg-type]
+            environment=environment,  # type: ignore[arg-type]
             api_key=api_key,
             timeout=timeout,
             httpx_client=httpx_client,
@@ -453,19 +561,20 @@ class AsyncAthena(AsyncBaseAthena):
         self.query = WrappedAsyncQueryClient(client_wrapper=self._client_wrapper)
 
 
-def _read_json_frame(model: DataFrameRequestOut) -> 'pd.DataFrame':
+def _read_json_frame(model: DataFrameRequestOut) -> "pd.DataFrame":
     import pandas as pd
 
     string_io = io.StringIO(model.json())
 
     with warnings.catch_warnings():
         # Filter warnings due to https://github.com/pandas-dev/pandas/issues/59511
-        warnings.simplefilter(action='ignore', category=FutureWarning)
-        return pd.read_json(string_io, orient='split')
+        warnings.simplefilter(action="ignore", category=FutureWarning)
+        return pd.read_json(string_io, orient="split")
 
 
 def _check_pandas_installed():
     import pandas
+
     assert pandas
 
 
@@ -514,10 +623,12 @@ def _to_pandas_df(
 def _convert_asset_object(
     asset_object: Union["pd.DataFrame", "pd.Series", core.File],
     name: Union[str, None] = None,
+    **kwargs,
 ) -> core.File:
     import pandas as pd
+
     try:
-        from IPython.core.formatters import format_display_data    # type: ignore[import]
+        from IPython.core.formatters import format_display_data  # type: ignore[import]
     except ImportError:
         format_display_data = None
 
@@ -526,7 +637,7 @@ def _convert_asset_object(
     if isinstance(asset_object, pd.DataFrame):
         return (
             name or "Uploaded data frame",
-            asset_object.to_parquet(path=None),
+            asset_object.to_parquet(path=None, **kwargs),
             "application/vnd.apache.parquet",
         )
     if format_display_data:
@@ -548,4 +659,4 @@ def _convert_asset_object(
                     image_bytes,
                     media_type,
                 )
-    return asset_object    # type: ignore[return-value]
+    return asset_object  # type: ignore[return-value]

athena/core/__init__.py

@@ -3,28 +3,45 @@
 from .api_error import ApiError
 from .client_wrapper import AsyncClientWrapper, BaseClientWrapper, SyncClientWrapper
 from .datetime_utils import serialize_datetime
-from .file import File, convert_file_dict_to_httpx_tuples
+from .file import File, convert_file_dict_to_httpx_tuples, with_content_type
 from .http_client import AsyncHttpClient, HttpClient
 from .jsonable_encoder import jsonable_encoder
-from .pydantic_utilities import deep_union_pydantic_dicts, pydantic_v1
+from .pydantic_utilities import (
+    IS_PYDANTIC_V2,
+    UniversalBaseModel,
+    UniversalRootModel,
+    parse_obj_as,
+    universal_field_validator,
+    universal_root_validator,
+    update_forward_refs,
+)
 from .query_encoder import encode_query
 from .remove_none_from_dict import remove_none_from_dict
 from .request_options import RequestOptions
+from .serialization import FieldMetadata, convert_and_respect_annotation_metadata
 
 __all__ = [
     "ApiError",
     "AsyncClientWrapper",
     "AsyncHttpClient",
     "BaseClientWrapper",
+    "FieldMetadata",
     "File",
     "HttpClient",
+    "IS_PYDANTIC_V2",
     "RequestOptions",
     "SyncClientWrapper",
+    "UniversalBaseModel",
+    "UniversalRootModel",
+    "convert_and_respect_annotation_metadata",
     "convert_file_dict_to_httpx_tuples",
-    "deep_union_pydantic_dicts",
     "encode_query",
     "jsonable_encoder",
-    "pydantic_v1",
+    "parse_obj_as",
     "remove_none_from_dict",
     "serialize_datetime",
+    "universal_field_validator",
+    "universal_root_validator",
+    "update_forward_refs",
+    "with_content_type",
 ]

athena/core/client_wrapper.py

@@ -1,10 +1,9 @@
 # This file was auto-generated by Fern from our API Definition.
 
 import typing
-
 import httpx
-
-from .http_client import AsyncHttpClient, HttpClient
+from .http_client import HttpClient
+from .http_client import AsyncHttpClient
 
 
 class BaseClientWrapper:
@@ -17,7 +16,7 @@ class BaseClientWrapper:
         headers: typing.Dict[str, str] = {
             "X-Fern-Language": "Python",
             "X-Fern-SDK-Name": "athena-intelligence",
-            "X-Fern-SDK-Version": "0.1.124",
+            "X-Fern-SDK-Version": "0.1.126",
         }
         headers["X-API-KEY"] = self.api_key
         return headers
@@ -36,9 +35,9 @@ class SyncClientWrapper(BaseClientWrapper):
         super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
         self.httpx_client = HttpClient(
             httpx_client=httpx_client,
-            base_headers=self.get_headers(),
-            base_timeout=self.get_timeout(),
-            base_url=self.get_base_url(),
+            base_headers=self.get_headers,
+            base_timeout=self.get_timeout,
+            base_url=self.get_base_url,
         )
 
 
@@ -49,7 +48,7 @@ class AsyncClientWrapper(BaseClientWrapper):
         super().__init__(api_key=api_key, base_url=base_url, timeout=timeout)
         self.httpx_client = AsyncHttpClient(
             httpx_client=httpx_client,
-            base_headers=self.get_headers(),
-            base_timeout=self.get_timeout(),
-            base_url=self.get_base_url(),
+            base_headers=self.get_headers,
+            base_timeout=self.get_timeout,
+            base_url=self.get_base_url,
         )

athena/core/file.py

@@ -1,25 +1,30 @@
 # This file was auto-generated by Fern from our API Definition.
 
-import typing
+from typing import IO, Dict, List, Mapping, Optional, Tuple, Union, cast
 
 # File typing inspired by the flexibility of types within the httpx library
 # https://github.com/encode/httpx/blob/master/httpx/_types.py
-FileContent = typing.Union[typing.IO[bytes], bytes, str]
-File = typing.Union[
+FileContent = Union[IO[bytes], bytes, str]
+File = Union[
     # file (or bytes)
     FileContent,
     # (filename, file (or bytes))
-    typing.Tuple[typing.Optional[str], FileContent],
+    Tuple[Optional[str], FileContent],
     # (filename, file (or bytes), content_type)
-    typing.Tuple[typing.Optional[str], FileContent, typing.Optional[str]],
+    Tuple[Optional[str], FileContent, Optional[str]],
     # (filename, file (or bytes), content_type, headers)
-    typing.Tuple[typing.Optional[str], FileContent, typing.Optional[str], typing.Mapping[str, str]],
+    Tuple[
+        Optional[str],
+        FileContent,
+        Optional[str],
+        Mapping[str, str],
+    ],
 ]
 
 
 def convert_file_dict_to_httpx_tuples(
-    d: typing.Dict[str, typing.Union[File, typing.List[File]]]
-) -> typing.List[typing.Tuple[str, File]]:
+    d: Dict[str, Union[File, List[File]]],
+) -> List[Tuple[str, File]]:
     """
     The format we use is a list of tuples, where the first element is the
     name of the file and the second is the file object. Typically HTTPX wants
@@ -36,3 +41,27 @@ def convert_file_dict_to_httpx_tuples(
         else:
             httpx_tuples.append((key, file_like))
     return httpx_tuples
+
+
+def with_content_type(*, file: File, default_content_type: str) -> File:
+    """
+    This function resolves to the file's content type, if provided, and defaults
+    to the default_content_type value if not.
+    """
+    if isinstance(file, tuple):
+        if len(file) == 2:
+            filename, content = cast(Tuple[Optional[str], FileContent], file)  # type: ignore
+            return (filename, content, default_content_type)
+        elif len(file) == 3:
+            filename, content, file_content_type = cast(Tuple[Optional[str], FileContent, Optional[str]], file)  # type: ignore
+            out_content_type = file_content_type or default_content_type
+            return (filename, content, out_content_type)
+        elif len(file) == 4:
+            filename, content, file_content_type, headers = cast(  # type: ignore
+                Tuple[Optional[str], FileContent, Optional[str], Mapping[str, str]], file
+            )
+            out_content_type = file_content_type or default_content_type
+            return (filename, content, out_content_type, headers)
+        else:
+            raise ValueError(f"Unexpected tuple length: {len(file)}")
+    return (None, file, default_content_type)