Skip to content

Commit 6a4ef93

Browse files
sergey-zinchenkosergey-zinchenko
authored
feat: allow both api-key and access token in DIAL client (#48)
1 parent 992f4e2 commit 6a4ef93

9 files changed

Lines changed: 203 additions & 116 deletions

File tree

‎aidial_client/__init__.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from aidial_client._auth import AsyncAuthValue, AuthType, SyncAuthValue
1+
from aidial_client._auth import AsyncAuthValue, SyncAuthValue
22
from aidial_client._client import AsyncDial, Dial
33
from aidial_client._client_pool import AsyncDialClientPool, DialClientPool
44
from aidial_client._exception import (
@@ -15,7 +15,6 @@
1515
"AsyncDial",
1616
"DialClientPool",
1717
"AsyncDialClientPool",
18-
"AuthType",
1918
"SyncAuthValue",
2019
"AsyncAuthValue",
2120
# Exceptions

‎aidial_client/_auth.py‎

Lines changed: 55 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -1,23 +1,5 @@
1-
from enum import Enum
21
from inspect import isawaitable
3-
from typing import (
4-
Awaitable,
5-
Callable,
6-
Dict,
7-
Optional,
8-
Tuple,
9-
TypeVar,
10-
Union,
11-
overload,
12-
)
13-
14-
from typing_extensions import assert_never
15-
16-
17-
class AuthType(Enum):
18-
API_KEY = "API_KEY"
19-
BEARER = "BEARER"
20-
2+
from typing import Awaitable, Callable, Dict, Optional, TypeVar, Union
213

224
SyncAuthValue = Union[str, Callable[[], str]]
235
AsyncAuthValue = Union[SyncAuthValue, Callable[[], Awaitable[str]]]
@@ -28,72 +10,78 @@ class AuthType(Enum):
2810
)
2911

3012

31-
@overload
32-
def get_auth_value(auth_value: SyncAuthValue) -> str: ...
33-
13+
def get_auth_value(auth_value: SyncAuthValue) -> str:
14+
if isinstance(auth_value, str):
15+
return auth_value
16+
if callable(auth_value):
17+
return auth_value()
18+
from typing import TYPE_CHECKING, assert_never
3419

35-
@overload
36-
def get_auth_value(
37-
auth_value: AsyncAuthValue,
38-
) -> Union[str, Awaitable[str]]: ...
20+
if TYPE_CHECKING:
21+
assert_never(auth_value)
22+
raise TypeError(
23+
f"auth_value must be a string or a callable returning a string, got {type(auth_value).__name__}"
24+
)
3925

4026

41-
def get_auth_value(
42-
auth_value: Union[SyncAuthValue, AsyncAuthValue]
43-
) -> Union[str, Awaitable[str]]:
27+
async def aget_auth_value(auth_value: AsyncAuthValue) -> str:
4428
if isinstance(auth_value, str):
4529
return auth_value
46-
elif callable(auth_value):
47-
return auth_value()
48-
else:
30+
if callable(auth_value):
31+
result = auth_value()
32+
return await result if isawaitable(result) else result
33+
from typing import TYPE_CHECKING, assert_never
34+
35+
if TYPE_CHECKING:
4936
assert_never(auth_value)
37+
raise TypeError(
38+
f"auth_value must be a string or a callable, got {type(auth_value).__name__}"
39+
)
5040

5141

52-
async def aget_auth_value(auth_value: AsyncAuthValue) -> str:
53-
processed_auth_value = get_auth_value(auth_value)
54-
if isawaitable(processed_auth_value):
55-
return await processed_auth_value
56-
return processed_auth_value
42+
def get_combined_auth_headers(
43+
*,
44+
api_key: Optional[SyncAuthValue] = None,
45+
bearer_token: Optional[SyncAuthValue] = None,
46+
) -> Dict[str, str]:
47+
headers: Dict[str, str] = {}
48+
49+
if api_key is not None:
50+
headers["api-key"] = get_auth_value(api_key)
5751

52+
if bearer_token is not None:
53+
bearer_str = get_auth_value(bearer_token)
54+
headers["Authorization"] = f"Bearer {bearer_str}"
5855

59-
def _get_auth_headers(auth_type: AuthType, auth_value: str) -> Dict[str, str]:
60-
if auth_type == AuthType.API_KEY:
61-
return {"api-key": auth_value}
62-
elif auth_type == AuthType.BEARER:
63-
return {"Authorization": f"Bearer {auth_value}"}
64-
else:
65-
assert_never(auth_type)
56+
return headers
6657

6758

68-
def get_auth_headers(
59+
async def aget_combined_auth_headers(
6960
*,
70-
auth_value: SyncAuthValue,
71-
auth_type: AuthType,
61+
api_key: Optional[AsyncAuthValue] = None,
62+
bearer_token: Optional[AsyncAuthValue] = None,
7263
) -> Dict[str, str]:
73-
processed_auth_value = get_auth_value(auth_value)
74-
return _get_auth_headers(auth_type, processed_auth_value)
64+
"""Get combined authentication headers from both api_key and bearer_token (async)."""
65+
headers: Dict[str, str] = {}
7566

67+
if api_key is not None:
68+
processed_api_key = await aget_auth_value(api_key)
69+
headers["api-key"] = processed_api_key
7670

77-
async def aget_auth_headers(
78-
auth_value: AsyncAuthValue,
79-
auth_type: AuthType,
80-
) -> Dict[str, str]:
81-
processed_auth_value = await aget_auth_value(auth_value)
82-
return _get_auth_headers(auth_type, processed_auth_value)
71+
if bearer_token is not None:
72+
processed_bearer_token = await aget_auth_value(bearer_token)
73+
headers["Authorization"] = f"Bearer {processed_bearer_token}"
74+
75+
return headers
8376

8477

85-
def process_auth(
78+
def validate_auth(
8679
*,
87-
api_key: Optional[AuthValueT] = None,
88-
bearer_token: Optional[AuthValueT] = None,
89-
) -> Tuple[AuthType, AuthValueT]:
90-
if api_key and bearer_token:
80+
api_key: Optional[AsyncAuthValue] = None,
81+
bearer_token: Optional[AsyncAuthValue] = None,
82+
) -> None:
83+
"""Validate that at least one authentication method is provided."""
84+
if not api_key and not bearer_token:
9185
raise ValueError(
92-
"Either api_key or bearer_token must be provided, but not both"
86+
"At least one of api_key or bearer_token must be provided"
9387
)
94-
elif api_key:
95-
return AuthType.API_KEY, api_key
96-
elif bearer_token:
97-
return AuthType.BEARER, bearer_token
98-
else:
99-
raise ValueError("Either api_key or bearer_token must be provided")

‎aidial_client/_client.py‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -9,10 +9,9 @@
99
import aidial_client.resources as resources
1010
from aidial_client._auth import (
1111
AsyncAuthValue,
12-
AuthType,
1312
AuthValueT,
1413
SyncAuthValue,
15-
process_auth,
14+
validate_auth,
1615
)
1716
from aidial_client._constants import (
1817
API_PREFIX,
@@ -31,8 +30,8 @@
3130

3231

3332
class BaseDialClient(Generic[_HttpClientT, AuthValueT], ABC):
34-
_auth_type: AuthType
35-
_auth_value: AuthValueT
33+
_api_key: Optional[AuthValueT]
34+
_bearer_token: Optional[AuthValueT]
3635
_base_url: str
3736
_http_client: _HttpClientT
3837
_auth_headers: Dict[str, str]
@@ -50,9 +49,9 @@ def __init__(
5049
api_version: Optional[str] = None,
5150
http_client: Optional[_HttpClientT] = None,
5251
):
53-
self._auth_type, self._auth_value = process_auth(
54-
api_key=api_key, bearer_token=bearer_token
55-
)
52+
validate_auth(api_key=api_key, bearer_token=bearer_token)
53+
self._api_key = api_key
54+
self._bearer_token = bearer_token
5655
self._max_retries = max_retries
5756
self._timeout = timeout
5857
self._base_url = enforce_trailing_slash(base_url)
@@ -85,6 +84,7 @@ def api_version(self) -> Optional[str]:
8584

8685

8786
class Dial(BaseDialClient[SyncHTTPClient, SyncAuthValue]):
87+
8888
def _init_resources(self) -> None:
8989
openai_client = openai.AzureOpenAI(
9090
api_key="-",
@@ -113,8 +113,8 @@ def _init_resources(self) -> None:
113113
def _create_http_client(self) -> SyncHTTPClient:
114114
return SyncHTTPClient(
115115
self._base_url,
116-
self._auth_value,
117-
self._auth_type,
116+
self._api_key,
117+
self._bearer_token,
118118
self._max_retries,
119119
self._timeout,
120120
)
@@ -156,10 +156,11 @@ def auth_headers(self) -> Dict[str, str]:
156156

157157

158158
class AsyncDial(BaseDialClient[AsyncHTTPClient, AsyncAuthValue]):
159+
159160
def _init_resources(self) -> None:
160161
openai_client = openai.AsyncAzureOpenAI(
161162
# set empty string, we will override
162-
# it with our client values during request
163+
# it with our client values during a request
163164
api_key="",
164165
api_version="",
165166
base_url=urljoin(self._base_url, OPENAI_PREFIX),
@@ -192,8 +193,8 @@ def _init_resources(self) -> None:
192193
def _create_http_client(self) -> AsyncHTTPClient:
193194
return AsyncHTTPClient(
194195
self._base_url,
195-
self._auth_value,
196-
self._auth_type,
196+
self._api_key,
197+
self._bearer_token,
197198
self._max_retries,
198199
self._timeout,
199200
)

‎aidial_client/_client_pool.py‎

Lines changed: 5 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,7 +2,7 @@
22

33
import httpx
44

5-
from aidial_client._auth import AsyncAuthValue, SyncAuthValue, process_auth
5+
from aidial_client._auth import AsyncAuthValue, SyncAuthValue
66
from aidial_client._client import AsyncDial, Dial
77
from aidial_client._constants import (
88
DEFAULT_CONNECTION_LIMITS,
@@ -32,17 +32,14 @@ def create_client(
3232
max_retries: int = DEFAULT_MAX_RETRIES,
3333
timeout: Union[httpx.Timeout, float] = DEFAULT_TIMEOUT,
3434
) -> Dial:
35-
auth_type, auth_value = process_auth(
36-
api_key=api_key, bearer_token=bearer_token
37-
)
3835
return Dial(
3936
base_url=base_url,
4037
api_key=api_key,
4138
bearer_token=bearer_token,
4239
http_client=SyncHTTPClient(
4340
base_url=base_url,
44-
auth_value=auth_value,
45-
auth_type=auth_type,
41+
api_key=api_key,
42+
bearer_token=bearer_token,
4643
max_retries=max_retries,
4744
timeout=timeout,
4845
internal_http_client=self._internal_http_client,
@@ -70,17 +67,14 @@ def create_client(
7067
max_retries: int = DEFAULT_MAX_RETRIES,
7168
timeout: Union[httpx.Timeout, float] = DEFAULT_TIMEOUT,
7269
) -> AsyncDial:
73-
auth_type, auth_value = process_auth(
74-
api_key=api_key, bearer_token=bearer_token
75-
)
7670
return AsyncDial(
7771
base_url=base_url,
7872
api_key=api_key,
7973
bearer_token=bearer_token,
8074
http_client=AsyncHTTPClient(
8175
base_url=base_url,
82-
auth_value=auth_value,
83-
auth_type=auth_type,
76+
api_key=api_key,
77+
bearer_token=bearer_token,
8478
max_retries=max_retries,
8579
timeout=timeout,
8680
internal_http_client=self._internal_http_client,

‎aidial_client/_http_client/_async.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import httpx
66

7-
from aidial_client._auth import AsyncAuthValue, aget_auth_headers
7+
from aidial_client._auth import AsyncAuthValue, aget_combined_auth_headers
88
from aidial_client._exception import DialException
99
from aidial_client._http_client._base import BaseHTTPClient
1010
from aidial_client._internal_types._generic import ResponseT
@@ -20,8 +20,8 @@ def _create_internal_client(self) -> httpx.AsyncClient:
2020
)
2121

2222
async def auth_headers(self) -> Dict[str, str]:
23-
return await aget_auth_headers(
24-
auth_value=self._auth_value, auth_type=self._auth_type
23+
return await aget_combined_auth_headers(
24+
api_key=self._api_key, bearer_token=self._bearer_token
2525
)
2626

2727
async def _retry_request(
@@ -99,7 +99,7 @@ async def request(
9999
cast_to=cast_to,
100100
remaining_retries=retries,
101101
)
102-
# Try to get custom error from response status_code/code/message
102+
# Try to get a custom error from response status_code/code/message
103103
custom_error = on_http_error(err) if on_http_error else None
104104
# or fallback to default processing
105105
raised_error = custom_error or self._make_dial_error_from_response(

‎aidial_client/_http_client/_base.py‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55

66
import httpx
77

8-
from aidial_client._auth import AuthType, AuthValueT
8+
from aidial_client._auth import AuthValueT
99
from aidial_client._constants import INITIAL_RETRY_DELAY, MAX_RETRY_DELAY
1010
from aidial_client._exception import DialException
1111
from aidial_client._internal_types._http_request import FinalRequestOptions
@@ -19,21 +19,21 @@
1919

2020
class BaseHTTPClient(ABC, Generic[_HttpInternalClientT, AuthValueT]):
2121
_internal_http_client: _HttpInternalClientT
22-
_auth_value: AuthValueT
23-
_auth_type: AuthType
22+
_api_key: Optional[AuthValueT]
23+
_bearer_token: Optional[AuthValueT]
2424

2525
def __init__(
2626
self,
2727
base_url: str,
28-
auth_value: AuthValueT,
29-
auth_type: AuthType,
28+
api_key: Optional[AuthValueT],
29+
bearer_token: Optional[AuthValueT],
3030
max_retries: int,
3131
timeout: Union[float, httpx.Timeout, None],
3232
internal_http_client: Optional[_HttpInternalClientT] = None,
3333
):
3434
self.base_url = httpx.URL(enforce_trailing_slash(base_url))
35-
self._auth_value = auth_value
36-
self._auth_type = auth_type
35+
self._api_key = api_key
36+
self._bearer_token = bearer_token
3737
self._max_retries = max_retries
3838
self._timeout = timeout
3939
self._internal_http_client = (

‎aidial_client/_http_client/_sync.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44

55
import httpx
66

7-
from aidial_client._auth import SyncAuthValue, get_auth_headers
7+
from aidial_client._auth import SyncAuthValue, get_combined_auth_headers
88
from aidial_client._exception import DialException
99
from aidial_client._http_client._base import BaseHTTPClient
1010
from aidial_client._internal_types._generic import ResponseT
@@ -37,8 +37,8 @@ def _retry_request(
3737
)
3838

3939
def auth_headers(self) -> Dict[str, str]:
40-
return get_auth_headers(
41-
auth_value=self._auth_value, auth_type=self._auth_type
40+
return get_combined_auth_headers(
41+
api_key=self._api_key, bearer_token=self._bearer_token
4242
)
4343

4444
def request(
@@ -99,7 +99,7 @@ def request(
9999
cast_to=cast_to,
100100
remaining_retries=retries,
101101
)
102-
# Try to get custom error from response status_code/code/message
102+
# Try to get a custom error from response status_code/code/message
103103
custom_error = on_http_error(err) if on_http_error else None
104104
# or fallback to default processing
105105
raised_error = custom_error or self._make_dial_error_from_response(

0 commit comments

Comments
 (0)