Skip to content

Commit c2aeb00

Browse files
committed
feat: support passing an API token
1 parent f099998 commit c2aeb00

File tree

3 files changed

+23
-11
lines changed

3 files changed

+23
-11
lines changed

README.md

Lines changed: 11 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -17,7 +17,10 @@ Python v3.13 or greater.
1717
```python
1818
from crowdstrike_aidr import AIGuard
1919

20-
client = AIGuard(base_url_template="https://api.crowdstrike.com/aidr/{SERVICE_NAME}")
20+
client = AIGuard(
21+
base_url_template="https://api.crowdstrike.com/aidr/{SERVICE_NAME}",
22+
token="my API token"
23+
)
2124

2225
response = client.guard_chat_completions(
2326
guard_input={
@@ -46,13 +49,15 @@ from crowdstrike_aidr import AIGuard
4649
# Using a float (total timeout in seconds).
4750
client = AIGuard(
4851
base_url_template="https://api.crowdstrike.com/aidr/{SERVICE_NAME}",
49-
timeout=30.0
52+
token="my API token",
53+
timeout=30.0,
5054
)
5155

5256
# Using httpx.Timeout for more granular control.
5357
client = AIGuard(
5458
base_url_template="https://api.crowdstrike.com/aidr/{SERVICE_NAME}",
55-
timeout=httpx.Timeout(timeout=60.0, connect=10.0)
59+
token="my API token",
60+
timeout=httpx.Timeout(timeout=60.0, connect=10.0),
5661
)
5762
```
5863

@@ -77,17 +82,14 @@ response = client.guard_chat_completions(
7782
## Retries
7883

7984
The SDK automatically retries failed requests with exponential backoff. By
80-
default, the client will retry up to 2 times.
81-
82-
### Client-level retries
83-
84-
Set the maximum number of retries for all requests:
85+
default, the client will retry up to 2 times. Set `max_retries` during client
86+
creation to change this.
8587

8688
```python
8789
from crowdstrike_aidr import AIGuard
8890

8991
client = AIGuard(
9092
base_url_template="https://api.crowdstrike.com/aidr/{SERVICE_NAME}",
91-
max_retries=5 # Retry up to 5 times
93+
max_retries=5 # Retry up to 5 times.
9294
)
9395
```

src/crowdstrike_aidr/_client.py

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,7 +5,7 @@
55
import time
66
from collections.abc import Mapping
77
from random import random
8-
from typing import TYPE_CHECKING, Any, cast, get_origin
8+
from typing import TYPE_CHECKING, Any, cast, get_origin, override
99

1010
import httpx
1111
from httpx import URL, Timeout
@@ -237,11 +237,13 @@ def __del__(self) -> None:
237237

238238
class SyncAPIClient(BaseClient[httpx.Client]):
239239
_client: httpx.Client
240+
_token: str
240241

241242
def __init__(
242243
self,
243244
*,
244245
base_url_template: str,
246+
token: str,
245247
max_retries: int = DEFAULT_MAX_RETRIES,
246248
timeout: float | Timeout = DEFAULT_TIMEOUT,
247249
http_client: httpx.Client | None = None,
@@ -260,9 +262,17 @@ def __init__(
260262
custom_headers=custom_headers,
261263
custom_query=custom_query,
262264
)
265+
266+
self._token = token
267+
263268
resolved_base_url = self.base_url
264269
self._client = http_client or SyncHttpxClientWrapper(base_url=resolved_base_url, timeout=self.timeout)
265270

271+
@property
272+
@override
273+
def auth_headers(self) -> dict[str, str]:
274+
return {"Authorization": f"Bearer {self._token}"}
275+
266276
def _post(
267277
self,
268278
path: str,

tests/test_ai_guard.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@
1515

1616
@pytest.fixture(scope="session")
1717
def client(request: pytest.FixtureRequest) -> Iterator[AIGuard]:
18-
yield AIGuard(base_url_template=base_url_template)
18+
yield AIGuard(base_url_template=base_url_template, token="my API token")
1919

2020

2121
def test_guard_chat_completions(client: AIGuard) -> None:

0 commit comments

Comments
 (0)