swarmauri_transport_aws_eventstream 0.1.0.dev1__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.
- swarmauri_transport_aws_eventstream-0.1.0.dev1/LICENSE +13 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/PKG-INFO +100 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/README.md +58 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/pyproject.toml +62 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/swarmauri_transport_aws_eventstream/AwsEventStreamClient.py +168 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/swarmauri_transport_aws_eventstream/AwsEventStreamCodec.py +45 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/swarmauri_transport_aws_eventstream/AwsEventStreamConnection.py +79 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/swarmauri_transport_aws_eventstream/AwsEventStreamError.py +2 -0
- swarmauri_transport_aws_eventstream-0.1.0.dev1/swarmauri_transport_aws_eventstream/__init__.py +11 -0
|
@@ -0,0 +1,13 @@
|
|
|
1
|
+
Copyright 2024 Jacob Stewart
|
|
2
|
+
|
|
3
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
4
|
+
you may not use this file except in compliance with the License.
|
|
5
|
+
You may obtain a copy of the License at
|
|
6
|
+
|
|
7
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
|
8
|
+
|
|
9
|
+
Unless required by applicable law or agreed to in writing, software
|
|
10
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
11
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
12
|
+
See the License for the specific language governing permissions and
|
|
13
|
+
limitations under the License.
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: swarmauri_transport_aws_eventstream
|
|
3
|
+
Version: 0.1.0.dev1
|
|
4
|
+
Summary: Amazon Bedrock bidirectional EventStream transport for Swarmauri live models.
|
|
5
|
+
License: Copyright 2024 Jacob Stewart
|
|
6
|
+
|
|
7
|
+
Licensed under the Apache License, Version 2.0 (the "License");
|
|
8
|
+
you may not use this file except in compliance with the License.
|
|
9
|
+
You may obtain a copy of the License at
|
|
10
|
+
|
|
11
|
+
http://www.apache.org/licenses/LICENSE-2.0
|
|
12
|
+
|
|
13
|
+
Unless required by applicable law or agreed to in writing, software
|
|
14
|
+
distributed under the License is distributed on an "AS IS" BASIS,
|
|
15
|
+
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
16
|
+
See the License for the specific language governing permissions and
|
|
17
|
+
limitations under the License.
|
|
18
|
+
License-File: LICENSE
|
|
19
|
+
Keywords: swarmauri,transport,aws,amazon bedrock,eventstream,bidirectional,nova sonic,smithy,http2
|
|
20
|
+
Author: Jacob Stewart
|
|
21
|
+
Author-email: jacob@swarmauri.com
|
|
22
|
+
Requires-Python: >=3.10,<3.13
|
|
23
|
+
Classifier: Development Status :: 1 - Planning
|
|
24
|
+
Classifier: Intended Audience :: Developers
|
|
25
|
+
Classifier: License :: OSI Approved :: Apache Software License
|
|
26
|
+
Classifier: Programming Language :: Python :: 3
|
|
27
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
28
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
29
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
30
|
+
Classifier: Topic :: Internet :: WWW/HTTP
|
|
31
|
+
Classifier: Topic :: Software Development :: Libraries
|
|
32
|
+
Requires-Dist: aws-sdk-bedrock-runtime (>=0.2.0,<0.3.0) ; python_version >= "3.12"
|
|
33
|
+
Requires-Dist: botocore (>=1.35)
|
|
34
|
+
Requires-Dist: swarmauri_base (>=0.11.0.dev5)
|
|
35
|
+
Requires-Dist: swarmauri_core (>=0.11.0.dev5)
|
|
36
|
+
Project-URL: Documentation, https://github.com/swarmauri/swarmauri-sdk/tree/master/pkgs/standards/swarmauri_transport_aws_eventstream#readme
|
|
37
|
+
Project-URL: Homepage, https://github.com/swarmauri/swarmauri-sdk
|
|
38
|
+
Project-URL: Repository, https://github.com/swarmauri/swarmauri-sdk
|
|
39
|
+
Project-URL: Source, https://github.com/swarmauri/swarmauri-sdk/tree/master/pkgs/standards/swarmauri_transport_aws_eventstream
|
|
40
|
+
Description-Content-Type: text/markdown
|
|
41
|
+
|
|
42
|
+

|
|
43
|
+
|
|
44
|
+
[](https://pypi.org/project/swarmauri_transport_aws_eventstream/) [](https://pypi.org/project/swarmauri_transport_aws_eventstream/) [](LICENSE) [](https://pepy.tech/project/swarmauri_transport_aws_eventstream) [](https://hits.sh/github.com/swarmauri/swarmauri-sdk/)
|
|
45
|
+
|
|
46
|
+
# swarmauri_transport_aws_eventstream
|
|
47
|
+
|
|
48
|
+
`swarmauri_transport_aws_eventstream` adapts Amazon Bedrock's Smithy
|
|
49
|
+
bidirectional EventStream operation to Swarmauri's
|
|
50
|
+
`IBidirectionalStreamClient` and `IBidirectionalStreamConnection` contracts.
|
|
51
|
+
The AWS SDK owns HTTP/2 framing, credential-chain resolution, and SigV4 request
|
|
52
|
+
and message signing; this package owns the Swarmauri lifecycle and JSON event
|
|
53
|
+
mapping.
|
|
54
|
+
|
|
55
|
+
## Features
|
|
56
|
+
|
|
57
|
+
- Sends and receives JSON events over `InvokeModelWithBidirectionalStream`.
|
|
58
|
+
- Supports explicit temporary credentials, the default AWS credential chain,
|
|
59
|
+
and named profiles.
|
|
60
|
+
- Exposes synchronous and asynchronous Swarmauri connection APIs.
|
|
61
|
+
- Converts AWS stream exception events into `AwsEventStreamError`.
|
|
62
|
+
- Keeps credentials out of serialized component state.
|
|
63
|
+
|
|
64
|
+
The codec and package metadata support Python 3.10 through 3.12. Amazon's
|
|
65
|
+
official `aws-sdk-bedrock-runtime` currently requires Python 3.12, so opening a
|
|
66
|
+
live Bedrock connection requires Python 3.12.
|
|
67
|
+
|
|
68
|
+
## Installation
|
|
69
|
+
|
|
70
|
+
```bash
|
|
71
|
+
uv add swarmauri_transport_aws_eventstream
|
|
72
|
+
```
|
|
73
|
+
|
|
74
|
+
```bash
|
|
75
|
+
pip install swarmauri_transport_aws_eventstream
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
## Usage
|
|
79
|
+
|
|
80
|
+
Use the client directly when you already have a Bedrock model ID:
|
|
81
|
+
|
|
82
|
+
```python
|
|
83
|
+
from swarmauri_transport_aws_eventstream import AwsEventStreamClient
|
|
84
|
+
|
|
85
|
+
client = AwsEventStreamClient(region="us-east-1")
|
|
86
|
+
connection = await client.aconnect(
|
|
87
|
+
"https://bedrock-runtime.us-east-1.amazonaws.com",
|
|
88
|
+
model_id="amazon.nova-sonic-v1:0",
|
|
89
|
+
)
|
|
90
|
+
|
|
91
|
+
await connection.asend({"event": {"sessionStart": {}}})
|
|
92
|
+
async for event in connection.aiter_events():
|
|
93
|
+
print(event)
|
|
94
|
+
await connection.aclose()
|
|
95
|
+
```
|
|
96
|
+
|
|
97
|
+
`NovaSonicLiveModel` configures this transport automatically. Pass an
|
|
98
|
+
`AmazonNovaProvider` when you need a named profile, explicit temporary
|
|
99
|
+
credentials, a different region, or a custom Bedrock runtime endpoint.
|
|
100
|
+
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+

|
|
2
|
+
|
|
3
|
+
[](https://pypi.org/project/swarmauri_transport_aws_eventstream/) [](https://pypi.org/project/swarmauri_transport_aws_eventstream/) [](LICENSE) [](https://pepy.tech/project/swarmauri_transport_aws_eventstream) [](https://hits.sh/github.com/swarmauri/swarmauri-sdk/)
|
|
4
|
+
|
|
5
|
+
# swarmauri_transport_aws_eventstream
|
|
6
|
+
|
|
7
|
+
`swarmauri_transport_aws_eventstream` adapts Amazon Bedrock's Smithy
|
|
8
|
+
bidirectional EventStream operation to Swarmauri's
|
|
9
|
+
`IBidirectionalStreamClient` and `IBidirectionalStreamConnection` contracts.
|
|
10
|
+
The AWS SDK owns HTTP/2 framing, credential-chain resolution, and SigV4 request
|
|
11
|
+
and message signing; this package owns the Swarmauri lifecycle and JSON event
|
|
12
|
+
mapping.
|
|
13
|
+
|
|
14
|
+
## Features
|
|
15
|
+
|
|
16
|
+
- Sends and receives JSON events over `InvokeModelWithBidirectionalStream`.
|
|
17
|
+
- Supports explicit temporary credentials, the default AWS credential chain,
|
|
18
|
+
and named profiles.
|
|
19
|
+
- Exposes synchronous and asynchronous Swarmauri connection APIs.
|
|
20
|
+
- Converts AWS stream exception events into `AwsEventStreamError`.
|
|
21
|
+
- Keeps credentials out of serialized component state.
|
|
22
|
+
|
|
23
|
+
The codec and package metadata support Python 3.10 through 3.12. Amazon's
|
|
24
|
+
official `aws-sdk-bedrock-runtime` currently requires Python 3.12, so opening a
|
|
25
|
+
live Bedrock connection requires Python 3.12.
|
|
26
|
+
|
|
27
|
+
## Installation
|
|
28
|
+
|
|
29
|
+
```bash
|
|
30
|
+
uv add swarmauri_transport_aws_eventstream
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
```bash
|
|
34
|
+
pip install swarmauri_transport_aws_eventstream
|
|
35
|
+
```
|
|
36
|
+
|
|
37
|
+
## Usage
|
|
38
|
+
|
|
39
|
+
Use the client directly when you already have a Bedrock model ID:
|
|
40
|
+
|
|
41
|
+
```python
|
|
42
|
+
from swarmauri_transport_aws_eventstream import AwsEventStreamClient
|
|
43
|
+
|
|
44
|
+
client = AwsEventStreamClient(region="us-east-1")
|
|
45
|
+
connection = await client.aconnect(
|
|
46
|
+
"https://bedrock-runtime.us-east-1.amazonaws.com",
|
|
47
|
+
model_id="amazon.nova-sonic-v1:0",
|
|
48
|
+
)
|
|
49
|
+
|
|
50
|
+
await connection.asend({"event": {"sessionStart": {}}})
|
|
51
|
+
async for event in connection.aiter_events():
|
|
52
|
+
print(event)
|
|
53
|
+
await connection.aclose()
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
`NovaSonicLiveModel` configures this transport automatically. Pass an
|
|
57
|
+
`AmazonNovaProvider` when you need a named profile, explicit temporary
|
|
58
|
+
credentials, a different region, or a custom Bedrock runtime endpoint.
|
|
@@ -0,0 +1,62 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "swarmauri_transport_aws_eventstream"
|
|
3
|
+
version = "0.1.0.dev1"
|
|
4
|
+
description = "Amazon Bedrock bidirectional EventStream transport for Swarmauri live models."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.10,<3.13"
|
|
7
|
+
authors = [{ name = "Jacob Stewart", email = "jacob@swarmauri.com" }]
|
|
8
|
+
license = { file = "LICENSE" }
|
|
9
|
+
dependencies = [
|
|
10
|
+
"botocore>=1.35",
|
|
11
|
+
"swarmauri_base>=0.11.0.dev5",
|
|
12
|
+
"swarmauri_core>=0.11.0.dev5",
|
|
13
|
+
"aws-sdk-bedrock-runtime~=0.2.0; python_version >= '3.12'",
|
|
14
|
+
]
|
|
15
|
+
keywords = [
|
|
16
|
+
"swarmauri",
|
|
17
|
+
"transport",
|
|
18
|
+
"aws",
|
|
19
|
+
"amazon bedrock",
|
|
20
|
+
"eventstream",
|
|
21
|
+
"bidirectional",
|
|
22
|
+
"nova sonic",
|
|
23
|
+
"smithy",
|
|
24
|
+
"http2",
|
|
25
|
+
]
|
|
26
|
+
classifiers = [
|
|
27
|
+
"Development Status :: 1 - Planning",
|
|
28
|
+
"Intended Audience :: Developers",
|
|
29
|
+
"License :: OSI Approved :: Apache Software License",
|
|
30
|
+
"Programming Language :: Python :: 3",
|
|
31
|
+
"Programming Language :: Python :: 3.10",
|
|
32
|
+
"Programming Language :: Python :: 3.11",
|
|
33
|
+
"Programming Language :: Python :: 3.12",
|
|
34
|
+
"Topic :: Internet :: WWW/HTTP",
|
|
35
|
+
"Topic :: Software Development :: Libraries",
|
|
36
|
+
]
|
|
37
|
+
|
|
38
|
+
[tool.uv.sources]
|
|
39
|
+
swarmauri_base = { workspace = true }
|
|
40
|
+
swarmauri_core = { workspace = true }
|
|
41
|
+
|
|
42
|
+
[project.urls]
|
|
43
|
+
Homepage = "https://github.com/swarmauri/swarmauri-sdk"
|
|
44
|
+
Repository = "https://github.com/swarmauri/swarmauri-sdk"
|
|
45
|
+
Source = "https://github.com/swarmauri/swarmauri-sdk/tree/master/pkgs/standards/swarmauri_transport_aws_eventstream"
|
|
46
|
+
Documentation = "https://github.com/swarmauri/swarmauri-sdk/tree/master/pkgs/standards/swarmauri_transport_aws_eventstream#readme"
|
|
47
|
+
|
|
48
|
+
[project.entry-points."swarmauri.transports"]
|
|
49
|
+
AwsEventStreamClient = "swarmauri_transport_aws_eventstream:AwsEventStreamClient"
|
|
50
|
+
|
|
51
|
+
[build-system]
|
|
52
|
+
requires = ["poetry-core>=1.8.0"]
|
|
53
|
+
build-backend = "poetry.core.masonry.api"
|
|
54
|
+
|
|
55
|
+
[dependency-groups]
|
|
56
|
+
dev = ["pytest>=8.0", "pytest-asyncio>=0.24", "ruff>=0.9.9", "swarmauri_tests_component"]
|
|
57
|
+
|
|
58
|
+
[tool.ruff]
|
|
59
|
+
extend = "../../pyproject.toml"
|
|
60
|
+
|
|
61
|
+
[tool.ruff.lint]
|
|
62
|
+
ignore = ["N999"]
|
|
@@ -0,0 +1,168 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import warnings
|
|
5
|
+
from collections.abc import Mapping
|
|
6
|
+
from typing import Any, Literal
|
|
7
|
+
from urllib.parse import unquote, urlsplit
|
|
8
|
+
|
|
9
|
+
from pydantic import Field
|
|
10
|
+
from swarmauri_base.ComponentBase import ComponentBase, ResourceTypes
|
|
11
|
+
from swarmauri_base.transports.bidirectional.BidirectionalStreamClientBase import ( # noqa: E501
|
|
12
|
+
BidirectionalStreamClientBase,
|
|
13
|
+
)
|
|
14
|
+
|
|
15
|
+
from .AwsEventStreamConnection import AwsEventStreamConnection
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
@ComponentBase.register_type(
|
|
19
|
+
BidirectionalStreamClientBase, "AwsEventStreamClient"
|
|
20
|
+
)
|
|
21
|
+
class AwsEventStreamClient(BidirectionalStreamClientBase):
|
|
22
|
+
"""Open Bedrock bidirectional streams through AWS's Smithy SDK."""
|
|
23
|
+
|
|
24
|
+
region: str = "us-east-1"
|
|
25
|
+
sdk_client_factory: Any = Field(default=None, exclude=True)
|
|
26
|
+
type: Literal["AwsEventStreamClient"] = "AwsEventStreamClient"
|
|
27
|
+
resource: str = Field(default=ResourceTypes.TRANSPORT.value, frozen=True)
|
|
28
|
+
|
|
29
|
+
@staticmethod
|
|
30
|
+
def _run_sync(awaitable: Any) -> Any:
|
|
31
|
+
try:
|
|
32
|
+
asyncio.get_running_loop()
|
|
33
|
+
except RuntimeError:
|
|
34
|
+
return asyncio.run(awaitable)
|
|
35
|
+
close = getattr(awaitable, "close", None)
|
|
36
|
+
if close is not None:
|
|
37
|
+
close()
|
|
38
|
+
raise RuntimeError("Use aconnect() inside an event loop")
|
|
39
|
+
|
|
40
|
+
@staticmethod
|
|
41
|
+
def _model_from_endpoint(endpoint: str) -> str | None:
|
|
42
|
+
path = urlsplit(endpoint).path
|
|
43
|
+
prefix = "/model/"
|
|
44
|
+
suffix = "/invoke-with-bidirectional-stream"
|
|
45
|
+
if path.startswith(prefix) and path.endswith(suffix):
|
|
46
|
+
return unquote(path[len(prefix) : -len(suffix)])
|
|
47
|
+
return None
|
|
48
|
+
|
|
49
|
+
@staticmethod
|
|
50
|
+
def _profile_credentials(profile: str) -> dict[str, str]:
|
|
51
|
+
from botocore.session import Session
|
|
52
|
+
|
|
53
|
+
credentials = Session(profile=profile).get_credentials()
|
|
54
|
+
if credentials is None:
|
|
55
|
+
raise ValueError(f"AWS profile {profile!r} has no credentials")
|
|
56
|
+
frozen = credentials.get_frozen_credentials()
|
|
57
|
+
result = {
|
|
58
|
+
"aws_access_key_id": frozen.access_key,
|
|
59
|
+
"aws_secret_access_key": frozen.secret_key,
|
|
60
|
+
}
|
|
61
|
+
if frozen.token:
|
|
62
|
+
result["aws_session_token"] = frozen.token
|
|
63
|
+
return result
|
|
64
|
+
|
|
65
|
+
@staticmethod
|
|
66
|
+
def _default_sdk():
|
|
67
|
+
try:
|
|
68
|
+
from aws_sdk_bedrock_runtime.config import Config
|
|
69
|
+
from aws_sdk_bedrock_runtime.models import (
|
|
70
|
+
InvokeModelWithBidirectionalStreamOperationInput,
|
|
71
|
+
)
|
|
72
|
+
except ImportError as exc: # pragma: no cover - Python < 3.12
|
|
73
|
+
raise RuntimeError(
|
|
74
|
+
"AWS Bedrock bidirectional transport requires Python 3.12 "
|
|
75
|
+
"and aws-sdk-bedrock-runtime"
|
|
76
|
+
) from exc
|
|
77
|
+
try:
|
|
78
|
+
from aws_sdk_bedrock_runtime.client import (
|
|
79
|
+
AsyncBedrockRuntimeClient,
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
client_type = AsyncBedrockRuntimeClient
|
|
83
|
+
except ImportError:
|
|
84
|
+
with warnings.catch_warnings():
|
|
85
|
+
warnings.simplefilter("ignore", DeprecationWarning)
|
|
86
|
+
from aws_sdk_bedrock_runtime.client import (
|
|
87
|
+
BedrockRuntimeClient,
|
|
88
|
+
)
|
|
89
|
+
|
|
90
|
+
client_type = BedrockRuntimeClient
|
|
91
|
+
return (
|
|
92
|
+
client_type,
|
|
93
|
+
Config,
|
|
94
|
+
(InvokeModelWithBidirectionalStreamOperationInput),
|
|
95
|
+
)
|
|
96
|
+
|
|
97
|
+
def connect(
|
|
98
|
+
self,
|
|
99
|
+
endpoint: str,
|
|
100
|
+
*,
|
|
101
|
+
headers: Mapping[str, str] | None = None,
|
|
102
|
+
**kwargs: Any,
|
|
103
|
+
) -> AwsEventStreamConnection:
|
|
104
|
+
return self._run_sync(
|
|
105
|
+
self.aconnect(endpoint, headers=headers, **kwargs)
|
|
106
|
+
)
|
|
107
|
+
|
|
108
|
+
async def aconnect(
|
|
109
|
+
self,
|
|
110
|
+
endpoint: str,
|
|
111
|
+
*,
|
|
112
|
+
headers: Mapping[str, str] | None = None,
|
|
113
|
+
**kwargs: Any,
|
|
114
|
+
) -> AwsEventStreamConnection:
|
|
115
|
+
if headers:
|
|
116
|
+
raise ValueError(
|
|
117
|
+
"Custom headers are unsupported; the AWS SDK owns "
|
|
118
|
+
"SigV4 headers"
|
|
119
|
+
)
|
|
120
|
+
|
|
121
|
+
model_id = kwargs.pop("model_id", None) or self._model_from_endpoint(
|
|
122
|
+
endpoint
|
|
123
|
+
)
|
|
124
|
+
if not model_id:
|
|
125
|
+
raise ValueError("model_id is required for a Bedrock event stream")
|
|
126
|
+
|
|
127
|
+
region = str(kwargs.pop("region", None) or self.region)
|
|
128
|
+
profile = kwargs.pop("profile", None)
|
|
129
|
+
credential_options = {
|
|
130
|
+
key: kwargs.pop(key)
|
|
131
|
+
for key in (
|
|
132
|
+
"aws_access_key_id",
|
|
133
|
+
"aws_secret_access_key",
|
|
134
|
+
"aws_session_token",
|
|
135
|
+
)
|
|
136
|
+
if kwargs.get(key) is not None
|
|
137
|
+
}
|
|
138
|
+
if ("aws_access_key_id" in credential_options) != (
|
|
139
|
+
"aws_secret_access_key" in credential_options
|
|
140
|
+
):
|
|
141
|
+
raise ValueError(
|
|
142
|
+
"aws_access_key_id and aws_secret_access_key must be "
|
|
143
|
+
"provided together"
|
|
144
|
+
)
|
|
145
|
+
if profile and not credential_options:
|
|
146
|
+
credential_options = await asyncio.to_thread(
|
|
147
|
+
self._profile_credentials, str(profile)
|
|
148
|
+
)
|
|
149
|
+
if kwargs:
|
|
150
|
+
names = ", ".join(sorted(kwargs))
|
|
151
|
+
raise TypeError(f"Unsupported AWS event stream options: {names}")
|
|
152
|
+
|
|
153
|
+
client_type, config_type, input_type = self._default_sdk()
|
|
154
|
+
split = urlsplit(endpoint)
|
|
155
|
+
if split.scheme not in {"http", "https"} or not split.netloc:
|
|
156
|
+
raise ValueError("endpoint must be an absolute HTTP or HTTPS URL")
|
|
157
|
+
endpoint_uri = f"{split.scheme}://{split.netloc}"
|
|
158
|
+
config = config_type(
|
|
159
|
+
region=region,
|
|
160
|
+
endpoint_uri=endpoint_uri,
|
|
161
|
+
**credential_options,
|
|
162
|
+
)
|
|
163
|
+
factory = self.sdk_client_factory or client_type
|
|
164
|
+
client = factory(config)
|
|
165
|
+
stream = client.invoke_model_with_bidirectional_stream(
|
|
166
|
+
input_type(model_id=str(model_id))
|
|
167
|
+
)
|
|
168
|
+
return AwsEventStreamConnection.from_stream(stream)
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
from .AwsEventStreamError import AwsEventStreamError
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class AwsEventStreamCodec:
|
|
10
|
+
"""Translate Python payloads to and from Bedrock Smithy events."""
|
|
11
|
+
|
|
12
|
+
@staticmethod
|
|
13
|
+
def _models():
|
|
14
|
+
try:
|
|
15
|
+
from aws_sdk_bedrock_runtime import models
|
|
16
|
+
except ImportError as exc: # pragma: no cover - Python < 3.12
|
|
17
|
+
raise RuntimeError(
|
|
18
|
+
"AWS Bedrock bidirectional transport requires Python 3.12 "
|
|
19
|
+
"and aws-sdk-bedrock-runtime"
|
|
20
|
+
) from exc
|
|
21
|
+
return models
|
|
22
|
+
|
|
23
|
+
def encode(self, event: Any) -> Any:
|
|
24
|
+
models = self._models()
|
|
25
|
+
payload = (
|
|
26
|
+
event
|
|
27
|
+
if isinstance(event, bytes)
|
|
28
|
+
else json.dumps(event, separators=(",", ":")).encode("utf-8")
|
|
29
|
+
)
|
|
30
|
+
return models.InvokeModelWithBidirectionalStreamInputChunk(
|
|
31
|
+
models.BidirectionalInputPayloadPart(bytes_=payload)
|
|
32
|
+
)
|
|
33
|
+
|
|
34
|
+
def decode(self, event: Any) -> Any:
|
|
35
|
+
if event.__class__.__name__ == (
|
|
36
|
+
"InvokeModelWithBidirectionalStreamOutputChunk"
|
|
37
|
+
):
|
|
38
|
+
payload = event.value.bytes_ or b""
|
|
39
|
+
return json.loads(payload.decode("utf-8"))
|
|
40
|
+
|
|
41
|
+
value = getattr(event, "value", None)
|
|
42
|
+
message = getattr(value, "message", None)
|
|
43
|
+
if message is None:
|
|
44
|
+
message = f"Unexpected AWS event stream event: {event!r}"
|
|
45
|
+
raise AwsEventStreamError(str(message))
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
from collections.abc import AsyncIterator, Iterator
|
|
4
|
+
from typing import Any
|
|
5
|
+
|
|
6
|
+
from pydantic import PrivateAttr
|
|
7
|
+
from swarmauri_base.transports.bidirectional.BidirectionalStreamConnectionBase import ( # noqa: E501
|
|
8
|
+
BidirectionalStreamConnectionBase,
|
|
9
|
+
)
|
|
10
|
+
from swarmauri_base.transports.bidirectional.BidirectionalStreamState import (
|
|
11
|
+
BidirectionalStreamState,
|
|
12
|
+
)
|
|
13
|
+
|
|
14
|
+
from .AwsEventStreamCodec import AwsEventStreamCodec
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class AwsEventStreamConnection(BidirectionalStreamConnectionBase):
|
|
18
|
+
"""Swarmauri connection adapter for a Smithy duplex event stream."""
|
|
19
|
+
|
|
20
|
+
_stream: Any = PrivateAttr()
|
|
21
|
+
_receiver: Any = PrivateAttr(default=None)
|
|
22
|
+
_stream_closed: bool = PrivateAttr(default=False)
|
|
23
|
+
_codec: AwsEventStreamCodec = PrivateAttr(
|
|
24
|
+
default_factory=AwsEventStreamCodec
|
|
25
|
+
)
|
|
26
|
+
|
|
27
|
+
@classmethod
|
|
28
|
+
def from_stream(cls, stream: Any) -> "AwsEventStreamConnection":
|
|
29
|
+
connection = cls(state=BidirectionalStreamState.OPEN.value)
|
|
30
|
+
connection._stream = stream
|
|
31
|
+
return connection
|
|
32
|
+
|
|
33
|
+
async def _output_receiver(self) -> Any:
|
|
34
|
+
if self._receiver is None:
|
|
35
|
+
_, self._receiver = await self._stream.await_output()
|
|
36
|
+
return self._receiver
|
|
37
|
+
|
|
38
|
+
async def asend(self, event: Any) -> None:
|
|
39
|
+
if self.state != BidirectionalStreamState.OPEN.value:
|
|
40
|
+
raise RuntimeError("AWS event stream is not open")
|
|
41
|
+
await self._stream.input_stream.send(self._codec.encode(event))
|
|
42
|
+
|
|
43
|
+
async def areceive(self) -> Any:
|
|
44
|
+
if self.state != BidirectionalStreamState.OPEN.value:
|
|
45
|
+
raise EOFError("AWS event stream is closed")
|
|
46
|
+
receiver = await self._output_receiver()
|
|
47
|
+
event = await receiver.receive()
|
|
48
|
+
if event is None:
|
|
49
|
+
await self._close_stream()
|
|
50
|
+
raise EOFError("AWS event stream ended")
|
|
51
|
+
return self._codec.decode(event)
|
|
52
|
+
|
|
53
|
+
def iter_events(self) -> Iterator[Any]:
|
|
54
|
+
while self.state == BidirectionalStreamState.OPEN.value:
|
|
55
|
+
try:
|
|
56
|
+
yield self.receive()
|
|
57
|
+
except EOFError:
|
|
58
|
+
return
|
|
59
|
+
|
|
60
|
+
async def aiter_events(self) -> AsyncIterator[Any]:
|
|
61
|
+
while self.state == BidirectionalStreamState.OPEN.value:
|
|
62
|
+
try:
|
|
63
|
+
yield await self.areceive()
|
|
64
|
+
except EOFError:
|
|
65
|
+
return
|
|
66
|
+
|
|
67
|
+
async def _close_stream(self) -> None:
|
|
68
|
+
if self._stream_closed:
|
|
69
|
+
self.state = BidirectionalStreamState.CLOSED.value
|
|
70
|
+
return
|
|
71
|
+
self.state = BidirectionalStreamState.CLOSING.value
|
|
72
|
+
try:
|
|
73
|
+
await self._stream.close()
|
|
74
|
+
finally:
|
|
75
|
+
self._stream_closed = True
|
|
76
|
+
self.state = BidirectionalStreamState.CLOSED.value
|
|
77
|
+
|
|
78
|
+
async def aclose(self) -> None:
|
|
79
|
+
await self._close_stream()
|
swarmauri_transport_aws_eventstream-0.1.0.dev1/swarmauri_transport_aws_eventstream/__init__.py
ADDED
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
from .AwsEventStreamClient import AwsEventStreamClient
|
|
2
|
+
from .AwsEventStreamCodec import AwsEventStreamCodec
|
|
3
|
+
from .AwsEventStreamConnection import AwsEventStreamConnection
|
|
4
|
+
from .AwsEventStreamError import AwsEventStreamError
|
|
5
|
+
|
|
6
|
+
__all__ = [
|
|
7
|
+
"AwsEventStreamClient",
|
|
8
|
+
"AwsEventStreamCodec",
|
|
9
|
+
"AwsEventStreamConnection",
|
|
10
|
+
"AwsEventStreamError",
|
|
11
|
+
]
|