truerecord-sdk 0.1.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- truerecord_sdk-0.1.2/LICENSE +21 -0
- truerecord_sdk-0.1.2/PKG-INFO +132 -0
- truerecord_sdk-0.1.2/README.md +112 -0
- truerecord_sdk-0.1.2/pyproject.toml +34 -0
- truerecord_sdk-0.1.2/src/truerecord_sdk/__init__.py +20 -0
- truerecord_sdk-0.1.2/src/truerecord_sdk/client.py +220 -0
- truerecord_sdk-0.1.2/src/truerecord_sdk/encryption.py +52 -0
- truerecord_sdk-0.1.2/src/truerecord_sdk/exceptions.py +27 -0
- truerecord_sdk-0.1.2/src/truerecord_sdk/py.typed +0 -0
- truerecord_sdk-0.1.2/tests/__init__.py +0 -0
- truerecord_sdk-0.1.2/tests/test_client.py +307 -0
- truerecord_sdk-0.1.2/tests/test_encryption.py +75 -0
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Exploring Data B.V.
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
|
@@ -0,0 +1,132 @@
|
|
|
1
|
+
Metadata-Version: 2.5
|
|
2
|
+
Name: truerecord-sdk
|
|
3
|
+
Version: 0.1.2
|
|
4
|
+
Summary: Python SDK for the TrueRecord LLM Governance Platform
|
|
5
|
+
Author-email: "Exploring Data B.V." <support@truerecord.eu>
|
|
6
|
+
License-Expression: MIT
|
|
7
|
+
License-File: LICENSE
|
|
8
|
+
Classifier: Development Status :: 3 - Alpha
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Programming Language :: Python :: 3
|
|
12
|
+
Classifier: Typing :: Typed
|
|
13
|
+
Requires-Python: >=3.9
|
|
14
|
+
Requires-Dist: cryptography>=41.0.0
|
|
15
|
+
Requires-Dist: httpx>=0.24.0
|
|
16
|
+
Provides-Extra: dev
|
|
17
|
+
Requires-Dist: pytest-asyncio>=0.21; extra == 'dev'
|
|
18
|
+
Requires-Dist: pytest>=7.0; extra == 'dev'
|
|
19
|
+
Description-Content-Type: text/markdown
|
|
20
|
+
|
|
21
|
+
# TrueRecord SDK for Python
|
|
22
|
+
|
|
23
|
+
[TrueRecord](https://truerecord.eu) keeps a record of every interaction with an AI system in your app. This SDK sends those records. Each message is encrypted before it leaves your process, so only TrueRecord can process it. See the [TrueRecord SDK docs](https://truerecord.eu/docs/sdk) for setup and guides.
|
|
24
|
+
|
|
25
|
+
Requires Python 3.9+.
|
|
26
|
+
|
|
27
|
+
## Install
|
|
28
|
+
|
|
29
|
+
```
|
|
30
|
+
pip install truerecord-sdk
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
## Quick start
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
from truerecord_sdk import TrueRecordClient, TrueRecordError
|
|
37
|
+
|
|
38
|
+
client = TrueRecordClient(
|
|
39
|
+
api_key="tr_sk_...",
|
|
40
|
+
public_key=open("truerecord_public.pem").read(),
|
|
41
|
+
application_guid="...",
|
|
42
|
+
interaction_guid="...",
|
|
43
|
+
)
|
|
44
|
+
|
|
45
|
+
try:
|
|
46
|
+
# Send the message your user sends to the LLM
|
|
47
|
+
client.send(user_message, "user", conversation_ref="session-42")
|
|
48
|
+
|
|
49
|
+
# Send a message from the LLM to the user
|
|
50
|
+
client.send(llm_message, "llm", conversation_ref="session-42")
|
|
51
|
+
except TrueRecordError as e:
|
|
52
|
+
print(e)
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## What to send
|
|
56
|
+
|
|
57
|
+
Each call to `send` creates a record of one message, either what the user sent (`"user"`) or what the LLM sent (`"llm"`). Use the same `conversation_ref` for both so TrueRecord can pair them.
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
client.send(message, origin, message_datetime=None, conversation_ref=None)
|
|
61
|
+
```
|
|
62
|
+
|
|
63
|
+
| Argument | Required | Description |
|
|
64
|
+
|---|---|---|
|
|
65
|
+
| `message` | Yes | The message a user sent to the AI or the AI sent to the user. Can be any value that converts to JSON (string, dict, list, etc.). |
|
|
66
|
+
| `origin` | Yes | `"user"` if this is a message sent to the LLM, `"llm"` if it's the LLM's response, `"agent"` if it comes from an agent. |
|
|
67
|
+
| `message_datetime` | No | When the message happened, as an ISO 8601 string (e.g. `"2025-01-15T10:30:00Z"`). When not set defaults to the current time. |
|
|
68
|
+
| `conversation_ref` | No | A reference you choose (like a session ID or chat ID) to link related messages together. Pass the same value for a user message and its AI response so TrueRecord can pair them. |
|
|
69
|
+
|
|
70
|
+
## Implementation details
|
|
71
|
+
|
|
72
|
+
### Non-blocking (recommended)
|
|
73
|
+
|
|
74
|
+
Sending data to TrueRecord runs in the background so the rest of your code continues without waiting for TrueRecord to respond.
|
|
75
|
+
|
|
76
|
+
#### Synchronous apps
|
|
77
|
+
|
|
78
|
+
For apps using Flask, Django or scripts and CLI tools. Use a thread pool to send in the background.
|
|
79
|
+
|
|
80
|
+
```python
|
|
81
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
82
|
+
|
|
83
|
+
executor = ThreadPoolExecutor(max_workers=4)
|
|
84
|
+
|
|
85
|
+
executor.submit(client.send, user_message, "user", conversation_ref="session-42")
|
|
86
|
+
executor.submit(client.send, llm_message, "llm", conversation_ref="session-42")
|
|
87
|
+
```
|
|
88
|
+
|
|
89
|
+
#### Asynchronous apps
|
|
90
|
+
|
|
91
|
+
For apps using FastAPI, Starlette, aiohttp or functions defined with `async`. Use `asyncio.create_task` to send in the background.
|
|
92
|
+
|
|
93
|
+
```python
|
|
94
|
+
import asyncio
|
|
95
|
+
|
|
96
|
+
asyncio.create_task(client.asend(user_message, "user", conversation_ref="session-42"))
|
|
97
|
+
asyncio.create_task(client.asend(llm_message, "llm", conversation_ref="session-42"))
|
|
98
|
+
```
|
|
99
|
+
|
|
100
|
+
### Blocking
|
|
101
|
+
|
|
102
|
+
Sending data to TrueRecord waits for TrueRecord to respond before the rest of your code continues.
|
|
103
|
+
|
|
104
|
+
#### Synchronous apps
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
client.send(user_message, "user", conversation_ref="session-42")
|
|
108
|
+
client.send(llm_message, "llm", conversation_ref="session-42")
|
|
109
|
+
```
|
|
110
|
+
|
|
111
|
+
#### Asynchronous apps
|
|
112
|
+
|
|
113
|
+
```python
|
|
114
|
+
await client.asend(user_message, "user", conversation_ref="session-42")
|
|
115
|
+
await client.asend(llm_message, "llm", conversation_ref="session-42")
|
|
116
|
+
```
|
|
117
|
+
|
|
118
|
+
## Errors
|
|
119
|
+
|
|
120
|
+
All errors extend `TrueRecordError`, so you can catch them in one place.
|
|
121
|
+
|
|
122
|
+
| Error | When |
|
|
123
|
+
|---|---|
|
|
124
|
+
| `ValidationError` | A required argument is missing or the message can't be serialised to JSON. |
|
|
125
|
+
| `EncryptionError` | The public key isn't a valid RSA key. |
|
|
126
|
+
| `NetworkError` | TrueRecord couldn't be reached after retries. |
|
|
127
|
+
| `AuthenticationError` | The API key was rejected. |
|
|
128
|
+
| `APIError` | Any other server error. Has `status_code` and `detail`. |
|
|
129
|
+
|
|
130
|
+
## License
|
|
131
|
+
|
|
132
|
+
MIT
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
# TrueRecord SDK for Python
|
|
2
|
+
|
|
3
|
+
[TrueRecord](https://truerecord.eu) keeps a record of every interaction with an AI system in your app. This SDK sends those records. Each message is encrypted before it leaves your process, so only TrueRecord can process it. See the [TrueRecord SDK docs](https://truerecord.eu/docs/sdk) for setup and guides.
|
|
4
|
+
|
|
5
|
+
Requires Python 3.9+.
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
```
|
|
10
|
+
pip install truerecord-sdk
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Quick start
|
|
14
|
+
|
|
15
|
+
```python
|
|
16
|
+
from truerecord_sdk import TrueRecordClient, TrueRecordError
|
|
17
|
+
|
|
18
|
+
client = TrueRecordClient(
|
|
19
|
+
api_key="tr_sk_...",
|
|
20
|
+
public_key=open("truerecord_public.pem").read(),
|
|
21
|
+
application_guid="...",
|
|
22
|
+
interaction_guid="...",
|
|
23
|
+
)
|
|
24
|
+
|
|
25
|
+
try:
|
|
26
|
+
# Send the message your user sends to the LLM
|
|
27
|
+
client.send(user_message, "user", conversation_ref="session-42")
|
|
28
|
+
|
|
29
|
+
# Send a message from the LLM to the user
|
|
30
|
+
client.send(llm_message, "llm", conversation_ref="session-42")
|
|
31
|
+
except TrueRecordError as e:
|
|
32
|
+
print(e)
|
|
33
|
+
```
|
|
34
|
+
|
|
35
|
+
## What to send
|
|
36
|
+
|
|
37
|
+
Each call to `send` creates a record of one message, either what the user sent (`"user"`) or what the LLM sent (`"llm"`). Use the same `conversation_ref` for both so TrueRecord can pair them.
|
|
38
|
+
|
|
39
|
+
```python
|
|
40
|
+
client.send(message, origin, message_datetime=None, conversation_ref=None)
|
|
41
|
+
```
|
|
42
|
+
|
|
43
|
+
| Argument | Required | Description |
|
|
44
|
+
|---|---|---|
|
|
45
|
+
| `message` | Yes | The message a user sent to the AI or the AI sent to the user. Can be any value that converts to JSON (string, dict, list, etc.). |
|
|
46
|
+
| `origin` | Yes | `"user"` if this is a message sent to the LLM, `"llm"` if it's the LLM's response, `"agent"` if it comes from an agent. |
|
|
47
|
+
| `message_datetime` | No | When the message happened, as an ISO 8601 string (e.g. `"2025-01-15T10:30:00Z"`). When not set defaults to the current time. |
|
|
48
|
+
| `conversation_ref` | No | A reference you choose (like a session ID or chat ID) to link related messages together. Pass the same value for a user message and its AI response so TrueRecord can pair them. |
|
|
49
|
+
|
|
50
|
+
## Implementation details
|
|
51
|
+
|
|
52
|
+
### Non-blocking (recommended)
|
|
53
|
+
|
|
54
|
+
Sending data to TrueRecord runs in the background so the rest of your code continues without waiting for TrueRecord to respond.
|
|
55
|
+
|
|
56
|
+
#### Synchronous apps
|
|
57
|
+
|
|
58
|
+
For apps using Flask, Django or scripts and CLI tools. Use a thread pool to send in the background.
|
|
59
|
+
|
|
60
|
+
```python
|
|
61
|
+
from concurrent.futures import ThreadPoolExecutor
|
|
62
|
+
|
|
63
|
+
executor = ThreadPoolExecutor(max_workers=4)
|
|
64
|
+
|
|
65
|
+
executor.submit(client.send, user_message, "user", conversation_ref="session-42")
|
|
66
|
+
executor.submit(client.send, llm_message, "llm", conversation_ref="session-42")
|
|
67
|
+
```
|
|
68
|
+
|
|
69
|
+
#### Asynchronous apps
|
|
70
|
+
|
|
71
|
+
For apps using FastAPI, Starlette, aiohttp or functions defined with `async`. Use `asyncio.create_task` to send in the background.
|
|
72
|
+
|
|
73
|
+
```python
|
|
74
|
+
import asyncio
|
|
75
|
+
|
|
76
|
+
asyncio.create_task(client.asend(user_message, "user", conversation_ref="session-42"))
|
|
77
|
+
asyncio.create_task(client.asend(llm_message, "llm", conversation_ref="session-42"))
|
|
78
|
+
```
|
|
79
|
+
|
|
80
|
+
### Blocking
|
|
81
|
+
|
|
82
|
+
Sending data to TrueRecord waits for TrueRecord to respond before the rest of your code continues.
|
|
83
|
+
|
|
84
|
+
#### Synchronous apps
|
|
85
|
+
|
|
86
|
+
```python
|
|
87
|
+
client.send(user_message, "user", conversation_ref="session-42")
|
|
88
|
+
client.send(llm_message, "llm", conversation_ref="session-42")
|
|
89
|
+
```
|
|
90
|
+
|
|
91
|
+
#### Asynchronous apps
|
|
92
|
+
|
|
93
|
+
```python
|
|
94
|
+
await client.asend(user_message, "user", conversation_ref="session-42")
|
|
95
|
+
await client.asend(llm_message, "llm", conversation_ref="session-42")
|
|
96
|
+
```
|
|
97
|
+
|
|
98
|
+
## Errors
|
|
99
|
+
|
|
100
|
+
All errors extend `TrueRecordError`, so you can catch them in one place.
|
|
101
|
+
|
|
102
|
+
| Error | When |
|
|
103
|
+
|---|---|
|
|
104
|
+
| `ValidationError` | A required argument is missing or the message can't be serialised to JSON. |
|
|
105
|
+
| `EncryptionError` | The public key isn't a valid RSA key. |
|
|
106
|
+
| `NetworkError` | TrueRecord couldn't be reached after retries. |
|
|
107
|
+
| `AuthenticationError` | The API key was rejected. |
|
|
108
|
+
| `APIError` | Any other server error. Has `status_code` and `detail`. |
|
|
109
|
+
|
|
110
|
+
## License
|
|
111
|
+
|
|
112
|
+
MIT
|
|
@@ -0,0 +1,34 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "truerecord-sdk"
|
|
7
|
+
version = "0.1.2"
|
|
8
|
+
description = "Python SDK for the TrueRecord LLM Governance Platform"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = "MIT"
|
|
12
|
+
license-files = ["LICENSE"]
|
|
13
|
+
authors = [{ name = "Exploring Data B.V.", email = "support@truerecord.eu" }]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Development Status :: 3 - Alpha",
|
|
16
|
+
"Intended Audience :: Developers",
|
|
17
|
+
"License :: OSI Approved :: MIT License",
|
|
18
|
+
"Programming Language :: Python :: 3",
|
|
19
|
+
"Typing :: Typed",
|
|
20
|
+
]
|
|
21
|
+
dependencies = [
|
|
22
|
+
"httpx>=0.24.0",
|
|
23
|
+
"cryptography>=41.0.0",
|
|
24
|
+
]
|
|
25
|
+
|
|
26
|
+
[project.optional-dependencies]
|
|
27
|
+
dev = ["pytest>=7.0", "pytest-asyncio>=0.21"]
|
|
28
|
+
|
|
29
|
+
[tool.hatch.build.targets.wheel]
|
|
30
|
+
packages = ["src/truerecord_sdk"]
|
|
31
|
+
|
|
32
|
+
[tool.pytest.ini_options]
|
|
33
|
+
testpaths = ["tests"]
|
|
34
|
+
asyncio_mode = "auto"
|
|
@@ -0,0 +1,20 @@
|
|
|
1
|
+
from truerecord_sdk.client import TrueRecordClient, __version__
|
|
2
|
+
from truerecord_sdk.exceptions import (
|
|
3
|
+
APIError,
|
|
4
|
+
AuthenticationError,
|
|
5
|
+
EncryptionError,
|
|
6
|
+
NetworkError,
|
|
7
|
+
TrueRecordError,
|
|
8
|
+
ValidationError,
|
|
9
|
+
)
|
|
10
|
+
|
|
11
|
+
__all__ = [
|
|
12
|
+
"TrueRecordClient",
|
|
13
|
+
"TrueRecordError",
|
|
14
|
+
"APIError",
|
|
15
|
+
"AuthenticationError",
|
|
16
|
+
"EncryptionError",
|
|
17
|
+
"NetworkError",
|
|
18
|
+
"ValidationError",
|
|
19
|
+
"__version__",
|
|
20
|
+
]
|
|
@@ -0,0 +1,220 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import json
|
|
5
|
+
import math
|
|
6
|
+
import random
|
|
7
|
+
import time
|
|
8
|
+
import uuid
|
|
9
|
+
from datetime import datetime, timezone
|
|
10
|
+
from importlib.metadata import PackageNotFoundError, version
|
|
11
|
+
from typing import Any, Literal
|
|
12
|
+
|
|
13
|
+
import httpx
|
|
14
|
+
|
|
15
|
+
from truerecord_sdk.encryption import Encryptor, load_public_key
|
|
16
|
+
from truerecord_sdk.exceptions import (
|
|
17
|
+
APIError,
|
|
18
|
+
AuthenticationError,
|
|
19
|
+
NetworkError,
|
|
20
|
+
ValidationError,
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
try:
|
|
24
|
+
__version__ = version("truerecord-sdk")
|
|
25
|
+
except PackageNotFoundError: # running from source, not installed
|
|
26
|
+
__version__ = "0.0.0"
|
|
27
|
+
|
|
28
|
+
API_URL = "https://api.truerecord.eu/evaluate"
|
|
29
|
+
|
|
30
|
+
# Responses worth retrying: the API is busy, restarting, or asked us to wait.
|
|
31
|
+
RETRY_STATUSES = frozenset({408, 409, 429, 500, 502, 503, 504})
|
|
32
|
+
TIMEOUT = 15.0 # seconds per attempt
|
|
33
|
+
RETRIES = 2 # attempts after the first
|
|
34
|
+
BACKOFF_BASE = 0.5 # seconds; doubles each retry
|
|
35
|
+
BACKOFF_CAP = 10.0 # seconds
|
|
36
|
+
RETRY_AFTER_CAP = 30.0 # seconds; the longest Retry-After we honour
|
|
37
|
+
DETAIL_MAX = 200 # characters of a response body quoted in an error
|
|
38
|
+
|
|
39
|
+
Origin = Literal["user", "llm", "agent"]
|
|
40
|
+
|
|
41
|
+
|
|
42
|
+
class TrueRecordClient:
|
|
43
|
+
"""Sends encrypted LLM interaction records to TrueRecord.
|
|
44
|
+
|
|
45
|
+
client = TrueRecordClient(api_key, public_key, application_guid, interaction_guid)
|
|
46
|
+
client.send({"prompt": "hello"}, "user")
|
|
47
|
+
client.send({"text": "hi there"}, "llm")
|
|
48
|
+
|
|
49
|
+
Every method raises only subclasses of `TrueRecordError`.
|
|
50
|
+
|
|
51
|
+
Retries: a send that fails because the API is unreachable, times out, or
|
|
52
|
+
answers 408/409/429/5xx is retried twice with exponential
|
|
53
|
+
backoff and jitter (about 0.5s, then 1s, then 2s...). A `Retry-After`
|
|
54
|
+
header from the API overrides the wait, up to 30s. Other 4xx responses
|
|
55
|
+
are not retried. Redirects are not followed: anything but a 2xx is an error.
|
|
56
|
+
"""
|
|
57
|
+
|
|
58
|
+
def __init__(
|
|
59
|
+
self,
|
|
60
|
+
api_key: str,
|
|
61
|
+
public_key: str,
|
|
62
|
+
application_guid: str,
|
|
63
|
+
interaction_guid: str,
|
|
64
|
+
) -> None:
|
|
65
|
+
"""
|
|
66
|
+
Args:
|
|
67
|
+
api_key: Your TrueRecord API key. Sent as a bearer token.
|
|
68
|
+
public_key: Your TrueRecord RSA public key, as a PEM string or
|
|
69
|
+
base64-encoded DER. Messages are encrypted with it before
|
|
70
|
+
leaving your process.
|
|
71
|
+
application_guid: Identifies the application the records belong to.
|
|
72
|
+
interaction_guid: Identifies the specific interaction within the application.
|
|
73
|
+
|
|
74
|
+
Raises:
|
|
75
|
+
ValidationError: An argument is empty.
|
|
76
|
+
EncryptionError: The public key is not a valid RSA key of at least 2048 bits.
|
|
77
|
+
"""
|
|
78
|
+
if not api_key:
|
|
79
|
+
raise ValidationError("api_key is required")
|
|
80
|
+
if not public_key:
|
|
81
|
+
raise ValidationError("public_key is required")
|
|
82
|
+
if not application_guid:
|
|
83
|
+
raise ValidationError("application_guid is required")
|
|
84
|
+
if not interaction_guid:
|
|
85
|
+
raise ValidationError("interaction_guid is required")
|
|
86
|
+
|
|
87
|
+
self._public_key = load_public_key(public_key)
|
|
88
|
+
self._application_guid = application_guid
|
|
89
|
+
self._interaction_guid = interaction_guid
|
|
90
|
+
self._headers = {
|
|
91
|
+
"Authorization": f"Bearer {api_key}",
|
|
92
|
+
"User-Agent": f"truerecord-sdk/{__version__}",
|
|
93
|
+
}
|
|
94
|
+
|
|
95
|
+
def send(
|
|
96
|
+
self,
|
|
97
|
+
message: Any,
|
|
98
|
+
origin: Origin,
|
|
99
|
+
message_datetime: str | None = None,
|
|
100
|
+
conversation_ref: str | None = None,
|
|
101
|
+
) -> dict[str, Any]:
|
|
102
|
+
"""Encrypt an LLM message and send it to TrueRecord.
|
|
103
|
+
|
|
104
|
+
Args:
|
|
105
|
+
message: The message content. Any JSON-serialisable value. Encrypted.
|
|
106
|
+
origin: ``"user"`` for a user message, ``"llm"`` for the LLM's
|
|
107
|
+
response, ``"agent"`` for a message from an agent.
|
|
108
|
+
message_datetime: When the message occurred, ISO 8601.
|
|
109
|
+
Defaults to now. Not encrypted.
|
|
110
|
+
conversation_ref: Optional ID linking interactions in one conversation.
|
|
111
|
+
|
|
112
|
+
Returns:
|
|
113
|
+
The server's JSON response as a dict.
|
|
114
|
+
|
|
115
|
+
Raises:
|
|
116
|
+
ValidationError: A required argument is missing or a value is not JSON-serialisable.
|
|
117
|
+
NetworkError: The server could not be reached or timed out, after all retries.
|
|
118
|
+
AuthenticationError: The API key was rejected.
|
|
119
|
+
APIError: The server returned any other error, after all retries where applicable.
|
|
120
|
+
"""
|
|
121
|
+
body = self._prepare(message, origin, message_datetime, conversation_ref)
|
|
122
|
+
with httpx.Client(headers=self._headers, timeout=TIMEOUT) as http:
|
|
123
|
+
for attempt in range(RETRIES + 1):
|
|
124
|
+
last = attempt == RETRIES
|
|
125
|
+
try:
|
|
126
|
+
result = http.post(API_URL, json=body)
|
|
127
|
+
except httpx.HTTPError as exc:
|
|
128
|
+
if last:
|
|
129
|
+
raise NetworkError(str(exc)) from exc
|
|
130
|
+
time.sleep(_backoff(attempt, None))
|
|
131
|
+
continue
|
|
132
|
+
if result.status_code in RETRY_STATUSES and not last:
|
|
133
|
+
time.sleep(_backoff(attempt, result.headers.get("Retry-After")))
|
|
134
|
+
continue
|
|
135
|
+
return _parse(result.status_code, result.text)
|
|
136
|
+
raise AssertionError("unreachable")
|
|
137
|
+
|
|
138
|
+
async def asend(
|
|
139
|
+
self,
|
|
140
|
+
message: Any,
|
|
141
|
+
origin: Origin,
|
|
142
|
+
message_datetime: str | None = None,
|
|
143
|
+
conversation_ref: str | None = None,
|
|
144
|
+
) -> dict[str, Any]:
|
|
145
|
+
"""Async version of `send`. Same arguments, return value, retries and errors."""
|
|
146
|
+
body = self._prepare(message, origin, message_datetime, conversation_ref)
|
|
147
|
+
async with httpx.AsyncClient(headers=self._headers, timeout=TIMEOUT) as http:
|
|
148
|
+
for attempt in range(RETRIES + 1):
|
|
149
|
+
last = attempt == RETRIES
|
|
150
|
+
try:
|
|
151
|
+
result = await http.post(API_URL, json=body)
|
|
152
|
+
except httpx.HTTPError as exc:
|
|
153
|
+
if last:
|
|
154
|
+
raise NetworkError(str(exc)) from exc
|
|
155
|
+
await asyncio.sleep(_backoff(attempt, None))
|
|
156
|
+
continue
|
|
157
|
+
if result.status_code in RETRY_STATUSES and not last:
|
|
158
|
+
await asyncio.sleep(_backoff(attempt, result.headers.get("Retry-After")))
|
|
159
|
+
continue
|
|
160
|
+
return _parse(result.status_code, result.text)
|
|
161
|
+
raise AssertionError("unreachable")
|
|
162
|
+
|
|
163
|
+
def _prepare(
|
|
164
|
+
self,
|
|
165
|
+
message: Any,
|
|
166
|
+
origin: Origin,
|
|
167
|
+
message_datetime: str | None,
|
|
168
|
+
conversation_ref: str | None,
|
|
169
|
+
) -> dict[str, Any]:
|
|
170
|
+
if message is None:
|
|
171
|
+
raise ValidationError("message is required")
|
|
172
|
+
if origin not in ("user", "llm", "agent"):
|
|
173
|
+
raise ValidationError("origin must be 'user', 'llm' or 'agent'")
|
|
174
|
+
|
|
175
|
+
enc = Encryptor(self._public_key)
|
|
176
|
+
try:
|
|
177
|
+
encrypted = enc.encrypt(message)
|
|
178
|
+
except (TypeError, ValueError) as exc:
|
|
179
|
+
raise ValidationError(f"message is not JSON-serialisable: {exc}") from exc
|
|
180
|
+
|
|
181
|
+
return {
|
|
182
|
+
"idempotency_key": str(uuid.uuid4()),
|
|
183
|
+
"interaction_guid": self._interaction_guid,
|
|
184
|
+
"application_guid": self._application_guid,
|
|
185
|
+
"conversation_ref": conversation_ref,
|
|
186
|
+
"sent_at_datetime": datetime.now(timezone.utc).isoformat(),
|
|
187
|
+
"record": {
|
|
188
|
+
"encryption_key": enc.wrapped_key,
|
|
189
|
+
"content": encrypted,
|
|
190
|
+
"type": origin,
|
|
191
|
+
"occurred_at": message_datetime or datetime.now(timezone.utc).isoformat(),
|
|
192
|
+
},
|
|
193
|
+
}
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
def _backoff(attempt: int, retry_after: str | None) -> float:
|
|
197
|
+
if retry_after:
|
|
198
|
+
try:
|
|
199
|
+
seconds = float(retry_after)
|
|
200
|
+
if math.isfinite(seconds):
|
|
201
|
+
return min(max(0.0, seconds), RETRY_AFTER_CAP)
|
|
202
|
+
except ValueError:
|
|
203
|
+
pass
|
|
204
|
+
return min(BACKOFF_BASE * 2**attempt, BACKOFF_CAP) * random.uniform(1.0, 1.5)
|
|
205
|
+
|
|
206
|
+
|
|
207
|
+
def _parse(status_code: int, text: str) -> dict[str, Any]:
|
|
208
|
+
if status_code in (401, 403):
|
|
209
|
+
raise AuthenticationError("API key rejected")
|
|
210
|
+
try:
|
|
211
|
+
body = json.loads(text)
|
|
212
|
+
except (ValueError, RecursionError):
|
|
213
|
+
body = None
|
|
214
|
+
# Only 2xx means the record landed. A 3xx is not followed and not a success.
|
|
215
|
+
if not 200 <= status_code < 300:
|
|
216
|
+
detail = body.get("detail") if isinstance(body, dict) else None
|
|
217
|
+
raise APIError(status_code, str(detail or text or "Unexpected status")[:DETAIL_MAX])
|
|
218
|
+
if not isinstance(body, dict):
|
|
219
|
+
raise APIError(status_code, f"Expected JSON object, got: {text[:DETAIL_MAX]}")
|
|
220
|
+
return body
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import base64
|
|
4
|
+
import json
|
|
5
|
+
import os
|
|
6
|
+
from typing import Any
|
|
7
|
+
|
|
8
|
+
from cryptography.hazmat.primitives import hashes, serialization
|
|
9
|
+
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
|
10
|
+
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
11
|
+
|
|
12
|
+
from truerecord_sdk.exceptions import EncryptionError
|
|
13
|
+
|
|
14
|
+
_OAEP = padding.OAEP(
|
|
15
|
+
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
|
16
|
+
algorithm=hashes.SHA256(),
|
|
17
|
+
label=None,
|
|
18
|
+
)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def load_public_key(key_string: str) -> rsa.RSAPublicKey:
|
|
22
|
+
try:
|
|
23
|
+
if "-----BEGIN" in key_string:
|
|
24
|
+
key = serialization.load_pem_public_key(key_string.encode())
|
|
25
|
+
else:
|
|
26
|
+
key = serialization.load_der_public_key(base64.b64decode(key_string))
|
|
27
|
+
except (ValueError, TypeError) as exc:
|
|
28
|
+
raise EncryptionError(f"Invalid public key: {exc}") from exc
|
|
29
|
+
if not isinstance(key, rsa.RSAPublicKey):
|
|
30
|
+
raise EncryptionError(f"Expected an RSA public key, got {type(key).__name__}")
|
|
31
|
+
if key.key_size < 2048:
|
|
32
|
+
raise EncryptionError(f"RSA key is {key.key_size} bits, need at least 2048")
|
|
33
|
+
return key
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
class Encryptor:
|
|
37
|
+
"""One fresh AES-256 key. `wrapped_key` is that key encrypted with the RSA public key."""
|
|
38
|
+
|
|
39
|
+
def __init__(self, public_key: rsa.RSAPublicKey) -> None:
|
|
40
|
+
key = AESGCM.generate_key(bit_length=256)
|
|
41
|
+
self._aes = AESGCM(key)
|
|
42
|
+
self.wrapped_key = _b64(public_key.encrypt(key, _OAEP))
|
|
43
|
+
|
|
44
|
+
def encrypt(self, value: Any) -> str:
|
|
45
|
+
"""JSON-encode `value`, encrypt with AES-GCM. Returns base64(nonce + ciphertext)."""
|
|
46
|
+
nonce = os.urandom(12)
|
|
47
|
+
plaintext = json.dumps(value, separators=(",", ":"), ensure_ascii=False).encode("utf-8")
|
|
48
|
+
return _b64(nonce + self._aes.encrypt(nonce, plaintext, None))
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def _b64(raw: bytes) -> str:
|
|
52
|
+
return base64.b64encode(raw).decode("ascii")
|
|
@@ -0,0 +1,27 @@
|
|
|
1
|
+
class TrueRecordError(Exception):
|
|
2
|
+
"""Base class for every error raised by this SDK."""
|
|
3
|
+
|
|
4
|
+
|
|
5
|
+
class ValidationError(TrueRecordError):
|
|
6
|
+
"""A required argument was missing or invalid. No request was sent."""
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class EncryptionError(TrueRecordError):
|
|
10
|
+
"""The public key could not be loaded or the record could not be encrypted."""
|
|
11
|
+
|
|
12
|
+
|
|
13
|
+
class NetworkError(TrueRecordError):
|
|
14
|
+
"""The server could not be reached, or the request timed out."""
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class AuthenticationError(TrueRecordError):
|
|
18
|
+
"""The server rejected the API key (HTTP 401 or 403)."""
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class APIError(TrueRecordError):
|
|
22
|
+
"""The server returned an error. Check `status_code` and `detail`."""
|
|
23
|
+
|
|
24
|
+
def __init__(self, status_code: int, detail: str) -> None:
|
|
25
|
+
self.status_code = status_code
|
|
26
|
+
self.detail = detail
|
|
27
|
+
super().__init__(f"HTTP {status_code}: {detail}")
|
|
File without changes
|
|
File without changes
|
|
@@ -0,0 +1,307 @@
|
|
|
1
|
+
import base64
|
|
2
|
+
import json
|
|
3
|
+
from datetime import datetime
|
|
4
|
+
|
|
5
|
+
import httpx
|
|
6
|
+
import pytest
|
|
7
|
+
from cryptography.hazmat.primitives import hashes, serialization
|
|
8
|
+
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
|
9
|
+
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
10
|
+
|
|
11
|
+
from truerecord_sdk import (
|
|
12
|
+
APIError,
|
|
13
|
+
AuthenticationError,
|
|
14
|
+
EncryptionError,
|
|
15
|
+
NetworkError,
|
|
16
|
+
TrueRecordClient,
|
|
17
|
+
ValidationError,
|
|
18
|
+
)
|
|
19
|
+
from truerecord_sdk import client as client_module
|
|
20
|
+
from truerecord_sdk.client import API_URL
|
|
21
|
+
|
|
22
|
+
OAEP = padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
@pytest.fixture(scope="module")
|
|
26
|
+
def private_key():
|
|
27
|
+
return rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
|
28
|
+
|
|
29
|
+
|
|
30
|
+
@pytest.fixture(scope="module")
|
|
31
|
+
def pem(private_key):
|
|
32
|
+
return private_key.public_key().public_bytes(
|
|
33
|
+
serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo
|
|
34
|
+
).decode()
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
@pytest.fixture
|
|
38
|
+
def client(pem):
|
|
39
|
+
return TrueRecordClient(
|
|
40
|
+
api_key="tr_sk_test", public_key=pem,
|
|
41
|
+
application_guid="app-1", interaction_guid="int-1",
|
|
42
|
+
)
|
|
43
|
+
|
|
44
|
+
|
|
45
|
+
def fake_server(monkeypatch, *responses):
|
|
46
|
+
"""Route every httpx request to a scripted sequence of responses.
|
|
47
|
+
|
|
48
|
+
Each entry is an httpx.Response or an exception to raise.
|
|
49
|
+
The last entry repeats. Returns captured requests and sleep durations."""
|
|
50
|
+
responses = list(responses) or [httpx.Response(200, json={"ok": True})]
|
|
51
|
+
captured = {"requests": [], "sleeps": []}
|
|
52
|
+
|
|
53
|
+
def handler(request):
|
|
54
|
+
captured["requests"].append(request)
|
|
55
|
+
item = responses[min(len(captured["requests"]) - 1, len(responses) - 1)]
|
|
56
|
+
if isinstance(item, Exception):
|
|
57
|
+
raise item
|
|
58
|
+
return item
|
|
59
|
+
|
|
60
|
+
transport = httpx.MockTransport(handler)
|
|
61
|
+
real_client, real_async = httpx.Client, httpx.AsyncClient
|
|
62
|
+
monkeypatch.setattr(httpx, "Client", lambda **kw: real_client(transport=transport, **kw))
|
|
63
|
+
monkeypatch.setattr(httpx, "AsyncClient", lambda **kw: real_async(transport=transport, **kw))
|
|
64
|
+
monkeypatch.setattr(client_module.time, "sleep", lambda s: captured["sleeps"].append(s))
|
|
65
|
+
|
|
66
|
+
async def fake_async_sleep(s):
|
|
67
|
+
captured["sleeps"].append(s)
|
|
68
|
+
|
|
69
|
+
monkeypatch.setattr(client_module.asyncio, "sleep", fake_async_sleep)
|
|
70
|
+
return captured
|
|
71
|
+
|
|
72
|
+
|
|
73
|
+
def decrypt(private_key, body):
|
|
74
|
+
key = private_key.decrypt(base64.b64decode(body["record"]["encryption_key"]), OAEP)
|
|
75
|
+
raw = base64.b64decode(body["record"]["content"])
|
|
76
|
+
return json.loads(AESGCM(key).decrypt(raw[:12], raw[12:], None))
|
|
77
|
+
|
|
78
|
+
|
|
79
|
+
# -- constructor --------------------------------------------------------------
|
|
80
|
+
|
|
81
|
+
@pytest.mark.parametrize("missing", ["api_key", "public_key", "application_guid", "interaction_guid"])
|
|
82
|
+
def test_init_requires_args(pem, missing):
|
|
83
|
+
kwargs = {"api_key": "k", "public_key": pem, "application_guid": "app-1", "interaction_guid": "int-1", missing: ""}
|
|
84
|
+
with pytest.raises(ValidationError, match=missing):
|
|
85
|
+
TrueRecordClient(**kwargs)
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def test_init_rejects_bad_key():
|
|
89
|
+
with pytest.raises(EncryptionError):
|
|
90
|
+
TrueRecordClient(api_key="k", public_key="bad", application_guid="app-1", interaction_guid="int-1")
|
|
91
|
+
|
|
92
|
+
|
|
93
|
+
# -- send ---------------------------------------------------------------------
|
|
94
|
+
|
|
95
|
+
def test_send_user_message(client, private_key, monkeypatch):
|
|
96
|
+
captured = fake_server(monkeypatch, httpx.Response(200, json={"id": "rec-1"}))
|
|
97
|
+
|
|
98
|
+
result = client.send(
|
|
99
|
+
{"prompt": "héllo 🎉"}, "user",
|
|
100
|
+
message_datetime="2026-09-17T10:00:00Z",
|
|
101
|
+
conversation_ref="conv-1",
|
|
102
|
+
)
|
|
103
|
+
|
|
104
|
+
assert result == {"id": "rec-1"}
|
|
105
|
+
[req] = captured["requests"]
|
|
106
|
+
assert req.method == "POST"
|
|
107
|
+
assert str(req.url) == API_URL
|
|
108
|
+
assert req.headers["Authorization"] == "Bearer tr_sk_test"
|
|
109
|
+
assert req.headers["User-Agent"].startswith("truerecord-sdk/")
|
|
110
|
+
|
|
111
|
+
body = json.loads(req.content)
|
|
112
|
+
assert body["application_guid"] == "app-1"
|
|
113
|
+
assert body["interaction_guid"] == "int-1"
|
|
114
|
+
assert body["conversation_ref"] == "conv-1"
|
|
115
|
+
datetime.fromisoformat(body["sent_at_datetime"])
|
|
116
|
+
assert body["record"]["type"] == "user"
|
|
117
|
+
assert body["record"]["occurred_at"] == "2026-09-17T10:00:00Z"
|
|
118
|
+
assert "prompt" not in body["record"]["content"]
|
|
119
|
+
assert decrypt(private_key, body) == {"prompt": "héllo 🎉"}
|
|
120
|
+
|
|
121
|
+
|
|
122
|
+
def test_send_llm_message(client, private_key, monkeypatch):
|
|
123
|
+
captured = fake_server(monkeypatch, httpx.Response(200, json={"id": "rec-2"}))
|
|
124
|
+
|
|
125
|
+
result = client.send(
|
|
126
|
+
{"text": "world"}, "llm",
|
|
127
|
+
message_datetime="2026-09-17T10:00:01Z",
|
|
128
|
+
)
|
|
129
|
+
|
|
130
|
+
assert result == {"id": "rec-2"}
|
|
131
|
+
body = json.loads(captured["requests"][0].content)
|
|
132
|
+
assert body["record"]["type"] == "llm"
|
|
133
|
+
assert body["record"]["occurred_at"] == "2026-09-17T10:00:01Z"
|
|
134
|
+
assert decrypt(private_key, body) == {"text": "world"}
|
|
135
|
+
|
|
136
|
+
|
|
137
|
+
def test_send_defaults_message_datetime(client, monkeypatch):
|
|
138
|
+
captured = fake_server(monkeypatch)
|
|
139
|
+
client.send("just a string", "user")
|
|
140
|
+
body = json.loads(captured["requests"][0].content)
|
|
141
|
+
assert body["conversation_ref"] is None
|
|
142
|
+
datetime.fromisoformat(body["record"]["occurred_at"])
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def test_send_agent_message(client, monkeypatch):
|
|
146
|
+
captured = fake_server(monkeypatch)
|
|
147
|
+
client.send("from an agent", "agent")
|
|
148
|
+
body = json.loads(captured["requests"][0].content)
|
|
149
|
+
assert body["record"]["type"] == "agent"
|
|
150
|
+
|
|
151
|
+
|
|
152
|
+
def test_send_requires_message(client):
|
|
153
|
+
with pytest.raises(ValidationError, match="message is required"):
|
|
154
|
+
client.send(None, "user")
|
|
155
|
+
|
|
156
|
+
|
|
157
|
+
def test_send_requires_valid_origin(client):
|
|
158
|
+
with pytest.raises(ValidationError, match="origin"):
|
|
159
|
+
client.send({}, "invalid")
|
|
160
|
+
with pytest.raises(ValidationError, match="origin"):
|
|
161
|
+
client.send({}, None)
|
|
162
|
+
|
|
163
|
+
|
|
164
|
+
def test_send_rejects_unserialisable(client):
|
|
165
|
+
with pytest.raises(ValidationError, match="not JSON-serialisable"):
|
|
166
|
+
client.send({"when": datetime.now()}, "user")
|
|
167
|
+
|
|
168
|
+
|
|
169
|
+
# -- errors -------------------------------------------------------------------
|
|
170
|
+
|
|
171
|
+
@pytest.mark.parametrize("status", [401, 403])
|
|
172
|
+
def test_auth_error_not_retried(client, monkeypatch, status):
|
|
173
|
+
captured = fake_server(monkeypatch, httpx.Response(status))
|
|
174
|
+
with pytest.raises(AuthenticationError):
|
|
175
|
+
client.send({}, "user")
|
|
176
|
+
assert len(captured["requests"]) == 1
|
|
177
|
+
|
|
178
|
+
|
|
179
|
+
def test_api_error_with_detail(client, monkeypatch):
|
|
180
|
+
fake_server(monkeypatch, httpx.Response(422, json={"detail": "bad guid"}))
|
|
181
|
+
with pytest.raises(APIError, match="HTTP 422: bad guid") as exc:
|
|
182
|
+
client.send({}, "user")
|
|
183
|
+
assert exc.value.status_code == 422
|
|
184
|
+
assert exc.value.detail == "bad guid"
|
|
185
|
+
|
|
186
|
+
|
|
187
|
+
def test_api_error_non_json_body(client, monkeypatch):
|
|
188
|
+
fake_server(monkeypatch, httpx.Response(400, text="<html>Bad Request</html>"))
|
|
189
|
+
with pytest.raises(APIError, match="400.*Bad Request"):
|
|
190
|
+
client.send({}, "user")
|
|
191
|
+
|
|
192
|
+
|
|
193
|
+
def test_success_non_json_body(client, monkeypatch):
|
|
194
|
+
fake_server(monkeypatch, httpx.Response(200, text="OK"))
|
|
195
|
+
with pytest.raises(APIError, match="Expected JSON object"):
|
|
196
|
+
client.send({}, "user")
|
|
197
|
+
|
|
198
|
+
|
|
199
|
+
def test_redirect_is_an_error_not_a_success(client, monkeypatch):
|
|
200
|
+
redirect = httpx.Response(302, headers={"Location": "https://elsewhere.example/"}, json={"ok": True})
|
|
201
|
+
captured = fake_server(monkeypatch, redirect)
|
|
202
|
+
with pytest.raises(APIError) as exc:
|
|
203
|
+
client.send({}, "user")
|
|
204
|
+
assert exc.value.status_code == 302
|
|
205
|
+
assert len(captured["requests"]) == 1
|
|
206
|
+
|
|
207
|
+
|
|
208
|
+
def test_error_detail_is_truncated(client, monkeypatch):
|
|
209
|
+
fake_server(monkeypatch, httpx.Response(400, text="x" * 5000))
|
|
210
|
+
with pytest.raises(APIError) as exc:
|
|
211
|
+
client.send({}, "user")
|
|
212
|
+
assert len(exc.value.detail) == client_module.DETAIL_MAX
|
|
213
|
+
|
|
214
|
+
|
|
215
|
+
# -- retries ------------------------------------------------------------------
|
|
216
|
+
|
|
217
|
+
def test_retries_then_succeeds(client, monkeypatch):
|
|
218
|
+
captured = fake_server(
|
|
219
|
+
monkeypatch,
|
|
220
|
+
httpx.Response(503), httpx.Response(503), httpx.Response(200, json={"ok": True}),
|
|
221
|
+
)
|
|
222
|
+
assert client.send({}, "user") == {"ok": True}
|
|
223
|
+
assert len(captured["requests"]) == 3
|
|
224
|
+
first, second = captured["sleeps"]
|
|
225
|
+
assert 0.5 <= first <= 0.75
|
|
226
|
+
assert 1.0 <= second <= 1.5
|
|
227
|
+
|
|
228
|
+
|
|
229
|
+
def test_gives_up_after_retries(client, monkeypatch):
|
|
230
|
+
captured = fake_server(monkeypatch, httpx.Response(503, json={"detail": "busy"}))
|
|
231
|
+
with pytest.raises(APIError, match="HTTP 503: busy"):
|
|
232
|
+
client.send({}, "user")
|
|
233
|
+
assert len(captured["requests"]) == 3 # 1 + 2 retries
|
|
234
|
+
|
|
235
|
+
|
|
236
|
+
def test_client_error_not_retried(client, monkeypatch):
|
|
237
|
+
captured = fake_server(monkeypatch, httpx.Response(422, json={"detail": "nope"}))
|
|
238
|
+
with pytest.raises(APIError):
|
|
239
|
+
client.send({}, "user")
|
|
240
|
+
assert len(captured["requests"]) == 1
|
|
241
|
+
assert captured["sleeps"] == []
|
|
242
|
+
|
|
243
|
+
|
|
244
|
+
def test_retry_after_header_honoured(client, monkeypatch):
|
|
245
|
+
captured = fake_server(
|
|
246
|
+
monkeypatch,
|
|
247
|
+
httpx.Response(429, headers={"Retry-After": "3"}), httpx.Response(200, json={}),
|
|
248
|
+
)
|
|
249
|
+
client.send({}, "user")
|
|
250
|
+
assert captured["sleeps"] == [3.0]
|
|
251
|
+
|
|
252
|
+
|
|
253
|
+
def test_retry_after_is_capped(client, monkeypatch):
|
|
254
|
+
captured = fake_server(
|
|
255
|
+
monkeypatch,
|
|
256
|
+
httpx.Response(503, headers={"Retry-After": "1e9"}), httpx.Response(200, json={}),
|
|
257
|
+
)
|
|
258
|
+
client.send({}, "user")
|
|
259
|
+
assert captured["sleeps"] == [client_module.RETRY_AFTER_CAP]
|
|
260
|
+
|
|
261
|
+
|
|
262
|
+
@pytest.mark.parametrize("value", ["inf", "nan", "soon"])
|
|
263
|
+
def test_retry_after_garbage_falls_back_to_backoff(client, monkeypatch, value):
|
|
264
|
+
captured = fake_server(
|
|
265
|
+
monkeypatch,
|
|
266
|
+
httpx.Response(503, headers={"Retry-After": value}), httpx.Response(200, json={}),
|
|
267
|
+
)
|
|
268
|
+
client.send({}, "user")
|
|
269
|
+
assert 0.5 <= captured["sleeps"][0] <= 0.75
|
|
270
|
+
|
|
271
|
+
|
|
272
|
+
@pytest.mark.parametrize("error", [httpx.ConnectError("refused"), httpx.ReadTimeout("slow")])
|
|
273
|
+
def test_network_error_retried_then_raised(client, monkeypatch, error):
|
|
274
|
+
captured = fake_server(monkeypatch, error)
|
|
275
|
+
with pytest.raises(NetworkError):
|
|
276
|
+
client.send({}, "user")
|
|
277
|
+
assert len(captured["requests"]) == 3
|
|
278
|
+
|
|
279
|
+
|
|
280
|
+
# -- asend --------------------------------------------------------------------
|
|
281
|
+
|
|
282
|
+
async def test_asend_user_message(client, private_key, monkeypatch):
|
|
283
|
+
captured = fake_server(monkeypatch, httpx.Response(200, json={"id": "rec-2"}))
|
|
284
|
+
result = await client.asend({"q": 1}, "user")
|
|
285
|
+
assert result == {"id": "rec-2"}
|
|
286
|
+
body = json.loads(captured["requests"][0].content)
|
|
287
|
+
assert decrypt(private_key, body) == {"q": 1}
|
|
288
|
+
|
|
289
|
+
|
|
290
|
+
async def test_asend_retries_then_succeeds(client, monkeypatch):
|
|
291
|
+
captured = fake_server(monkeypatch, httpx.Response(503), httpx.Response(200, json={}))
|
|
292
|
+
await client.asend({}, "user")
|
|
293
|
+
assert len(captured["requests"]) == 2
|
|
294
|
+
assert 0.5 <= captured["sleeps"][0] <= 0.75
|
|
295
|
+
|
|
296
|
+
|
|
297
|
+
async def test_asend_errors(client, monkeypatch):
|
|
298
|
+
fake_server(monkeypatch, httpx.Response(401))
|
|
299
|
+
with pytest.raises(AuthenticationError):
|
|
300
|
+
await client.asend({}, "user")
|
|
301
|
+
|
|
302
|
+
|
|
303
|
+
async def test_asend_redirect_is_an_error(client, monkeypatch):
|
|
304
|
+
fake_server(monkeypatch, httpx.Response(307, headers={"Location": "https://elsewhere.example/"}))
|
|
305
|
+
with pytest.raises(APIError) as exc:
|
|
306
|
+
await client.asend({}, "user")
|
|
307
|
+
assert exc.value.status_code == 307
|
|
@@ -0,0 +1,75 @@
|
|
|
1
|
+
import base64
|
|
2
|
+
import json
|
|
3
|
+
|
|
4
|
+
import pytest
|
|
5
|
+
from cryptography.hazmat.primitives import hashes, serialization
|
|
6
|
+
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
|
7
|
+
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
|
|
8
|
+
|
|
9
|
+
from truerecord_sdk.encryption import Encryptor, load_public_key
|
|
10
|
+
from truerecord_sdk.exceptions import EncryptionError
|
|
11
|
+
|
|
12
|
+
OAEP = padding.OAEP(mgf=padding.MGF1(hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def pem_for(private_key):
|
|
16
|
+
return private_key.public_key().public_bytes(
|
|
17
|
+
serialization.Encoding.PEM, serialization.PublicFormat.SubjectPublicKeyInfo
|
|
18
|
+
).decode()
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
def der_b64_for(private_key):
|
|
22
|
+
return base64.b64encode(private_key.public_key().public_bytes(
|
|
23
|
+
serialization.Encoding.DER, serialization.PublicFormat.SubjectPublicKeyInfo
|
|
24
|
+
)).decode()
|
|
25
|
+
|
|
26
|
+
|
|
27
|
+
@pytest.fixture(scope="module")
|
|
28
|
+
def private_key():
|
|
29
|
+
return rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
@pytest.fixture(scope="module")
|
|
33
|
+
def public_key(private_key):
|
|
34
|
+
return load_public_key(pem_for(private_key))
|
|
35
|
+
|
|
36
|
+
|
|
37
|
+
def decrypt(private_key, wrapped_key: str, blob: str):
|
|
38
|
+
key = private_key.decrypt(base64.b64decode(wrapped_key), OAEP)
|
|
39
|
+
raw = base64.b64decode(blob)
|
|
40
|
+
return json.loads(AESGCM(key).decrypt(raw[:12], raw[12:], None))
|
|
41
|
+
|
|
42
|
+
|
|
43
|
+
def test_rejects_garbage():
|
|
44
|
+
with pytest.raises(EncryptionError, match="Invalid public key"):
|
|
45
|
+
load_public_key("not a valid key")
|
|
46
|
+
|
|
47
|
+
|
|
48
|
+
def test_rejects_small_key():
|
|
49
|
+
small = rsa.generate_private_key(public_exponent=65537, key_size=1024)
|
|
50
|
+
with pytest.raises(EncryptionError, match="at least 2048"):
|
|
51
|
+
load_public_key(pem_for(small))
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def test_round_trip(private_key, public_key):
|
|
55
|
+
value = {"text": "héllo 日本 🎉 \"quotes\" \\ \x00\n", "n": [1, 2.5, None, True]}
|
|
56
|
+
enc = Encryptor(public_key)
|
|
57
|
+
assert decrypt(private_key, enc.wrapped_key, enc.encrypt(value)) == value
|
|
58
|
+
|
|
59
|
+
|
|
60
|
+
def test_round_trip_der_b64(private_key):
|
|
61
|
+
public_key = load_public_key(der_b64_for(private_key))
|
|
62
|
+
value = {"msg": "hello from DER"}
|
|
63
|
+
enc = Encryptor(public_key)
|
|
64
|
+
assert decrypt(private_key, enc.wrapped_key, enc.encrypt(value)) == value
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def test_same_key_different_nonce(private_key, public_key):
|
|
68
|
+
enc = Encryptor(public_key)
|
|
69
|
+
a, b = enc.encrypt("x"), enc.encrypt("x")
|
|
70
|
+
assert a != b
|
|
71
|
+
assert decrypt(private_key, enc.wrapped_key, a) == decrypt(private_key, enc.wrapped_key, b) == "x"
|
|
72
|
+
|
|
73
|
+
|
|
74
|
+
def test_fresh_key_per_encryptor(public_key):
|
|
75
|
+
assert Encryptor(public_key).wrapped_key != Encryptor(public_key).wrapped_key
|