hostel-protocol-python 0.1.0__tar.gz → 0.2.0__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.
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/Makefile +6 -3
- hostel_protocol_python-0.2.0/PKG-INFO +12 -0
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/pyproject.toml +13 -4
- hostel_protocol_python-0.2.0/src/hostel/client/__init__.py +5 -0
- hostel_protocol_python-0.2.0/src/hostel/client/client.py +302 -0
- {hostel_protocol_python-0.1.0/src/hostel_protocol → hostel_protocol_python-0.2.0/src/hostel/protocol}/__init__.py +7 -3
- {hostel_protocol_python-0.1.0/src/hostel_protocol → hostel_protocol_python-0.2.0/src/hostel/protocol}/converter.py +1 -1
- hostel_protocol_python-0.2.0/src/hostel/transport/__init__.py +11 -0
- hostel_protocol_python-0.2.0/src/hostel/transport/transport.py +45 -0
- hostel_protocol_python-0.2.0/src/hostel/transport/zeromq.py +119 -0
- hostel_protocol_python-0.2.0/tests/__init__.py +0 -0
- hostel_protocol_python-0.2.0/tests/test_client.py +512 -0
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/tests/test_converter.py +2 -2
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/tests/test_models.py +1 -1
- hostel_protocol_python-0.2.0/tests/test_transport.py +289 -0
- hostel_protocol_python-0.1.0/PKG-INFO +0 -16
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/.github/workflows/ci.yml +0 -0
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/.github/workflows/publish.yml +0 -0
- {hostel_protocol_python-0.1.0 → hostel_protocol_python-0.2.0}/.gitignore +0 -0
- {hostel_protocol_python-0.1.0/src/hostel_protocol → hostel_protocol_python-0.2.0/src/hostel/protocol}/models.py +0 -0
- /hostel_protocol_python-0.1.0/tests/__init__.py → /hostel_protocol_python-0.2.0/src/hostel/py.typed +0 -0
|
@@ -1,7 +1,7 @@
|
|
|
1
|
-
.PHONY: install format lint typecheck test
|
|
1
|
+
.PHONY: install format lint typecheck test build
|
|
2
2
|
|
|
3
3
|
install:
|
|
4
|
-
uv sync --all-
|
|
4
|
+
uv sync --all-groups
|
|
5
5
|
|
|
6
6
|
format:
|
|
7
7
|
uv run black src/ tests/
|
|
@@ -12,7 +12,10 @@ lint:
|
|
|
12
12
|
uv run black --check src/ tests/
|
|
13
13
|
|
|
14
14
|
typecheck:
|
|
15
|
-
uv run mypy
|
|
15
|
+
uv run mypy --package hostel.protocol --package hostel.transport --package hostel.client
|
|
16
16
|
|
|
17
17
|
test:
|
|
18
18
|
uv run pytest
|
|
19
|
+
|
|
20
|
+
build:
|
|
21
|
+
uv build
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: hostel-protocol-python
|
|
3
|
+
Version: 0.2.0
|
|
4
|
+
Summary: Pydantic models, transport and client for the Hostel protocol
|
|
5
|
+
Requires-Python: >=3.12
|
|
6
|
+
Requires-Dist: hostel-protocol>=0.1.0
|
|
7
|
+
Requires-Dist: protobuf>=7.0
|
|
8
|
+
Requires-Dist: pydantic>=2.5.0
|
|
9
|
+
Provides-Extra: client
|
|
10
|
+
Requires-Dist: pyzmq>=26.0; extra == 'client'
|
|
11
|
+
Provides-Extra: transport
|
|
12
|
+
Requires-Dist: pyzmq>=26.0; extra == 'transport'
|
|
@@ -4,8 +4,8 @@ build-backend = "hatchling.build"
|
|
|
4
4
|
|
|
5
5
|
[project]
|
|
6
6
|
name = "hostel-protocol-python"
|
|
7
|
-
version = "0.
|
|
8
|
-
description = "Pydantic
|
|
7
|
+
version = "0.2.0"
|
|
8
|
+
description = "Pydantic models, transport and client for the Hostel protocol"
|
|
9
9
|
requires-python = ">=3.12"
|
|
10
10
|
dependencies = [
|
|
11
11
|
"hostel-protocol>=0.1.0",
|
|
@@ -14,10 +14,15 @@ dependencies = [
|
|
|
14
14
|
]
|
|
15
15
|
|
|
16
16
|
[project.optional-dependencies]
|
|
17
|
+
transport = ["pyzmq>=26.0"]
|
|
18
|
+
client = ["pyzmq>=26.0"]
|
|
19
|
+
|
|
20
|
+
[dependency-groups]
|
|
17
21
|
dev = [
|
|
18
22
|
"pytest",
|
|
19
23
|
"pytest-cov",
|
|
20
24
|
"pytest-asyncio",
|
|
25
|
+
"pyzmq>=26.0",
|
|
21
26
|
"black",
|
|
22
27
|
"ruff",
|
|
23
28
|
"mypy",
|
|
@@ -25,7 +30,7 @@ dev = [
|
|
|
25
30
|
]
|
|
26
31
|
|
|
27
32
|
[tool.hatch.build.targets.wheel]
|
|
28
|
-
packages = ["src/
|
|
33
|
+
packages = ["src/hostel"]
|
|
29
34
|
|
|
30
35
|
[tool.black]
|
|
31
36
|
line-length = 120
|
|
@@ -45,9 +50,13 @@ plugins = ["pydantic.mypy"]
|
|
|
45
50
|
module = "hostel.protocol.v1.*"
|
|
46
51
|
ignore_missing_imports = true
|
|
47
52
|
|
|
53
|
+
[[tool.mypy.overrides]]
|
|
54
|
+
module = "zmq.*"
|
|
55
|
+
ignore_missing_imports = true
|
|
56
|
+
|
|
48
57
|
[tool.pytest.ini_options]
|
|
49
58
|
testpaths = ["tests"]
|
|
50
|
-
addopts = "--cov=
|
|
59
|
+
addopts = "--cov=hostel --cov-report=term-missing"
|
|
51
60
|
|
|
52
61
|
[tool.uv]
|
|
53
62
|
constraint-dependencies = ["protobuf>=7.0"]
|
|
@@ -0,0 +1,302 @@
|
|
|
1
|
+
"""Async client for the Hostel protocol."""
|
|
2
|
+
|
|
3
|
+
import asyncio
|
|
4
|
+
import logging
|
|
5
|
+
import uuid
|
|
6
|
+
from collections.abc import AsyncGenerator
|
|
7
|
+
from typing import Any
|
|
8
|
+
|
|
9
|
+
from hostel.protocol.models import (
|
|
10
|
+
ChatMessage,
|
|
11
|
+
ChatRequest,
|
|
12
|
+
ChatResponse,
|
|
13
|
+
CreateComponentRequest,
|
|
14
|
+
CreateTaskRequest,
|
|
15
|
+
DeleteComponentRequest,
|
|
16
|
+
DeleteTaskRequest,
|
|
17
|
+
GetComponentRequest,
|
|
18
|
+
GetTaskRequest,
|
|
19
|
+
HostelMessage,
|
|
20
|
+
ListAgentsRequest,
|
|
21
|
+
ListComponentsRequest,
|
|
22
|
+
ListTasksRequest,
|
|
23
|
+
TaskData,
|
|
24
|
+
UpdateComponentRequest,
|
|
25
|
+
UpdateTaskRequest,
|
|
26
|
+
)
|
|
27
|
+
from hostel.transport import DEFAULT_IPC_ENDPOINT
|
|
28
|
+
from hostel.transport.transport import Transport
|
|
29
|
+
from hostel.transport.zeromq import ZmqTransport
|
|
30
|
+
|
|
31
|
+
logger = logging.getLogger(__name__)
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
class HostelClient:
|
|
35
|
+
|
|
36
|
+
def __init__(self, transport: Transport | None = None, endpoint: str | None = None):
|
|
37
|
+
if transport is not None:
|
|
38
|
+
self.transport = transport
|
|
39
|
+
else:
|
|
40
|
+
self.transport = ZmqTransport(endpoint or DEFAULT_IPC_ENDPOINT)
|
|
41
|
+
self._pending: dict[str, asyncio.Queue[HostelMessage]] = {}
|
|
42
|
+
|
|
43
|
+
async def start(self) -> None:
|
|
44
|
+
logger.debug("Starting Hostel Client...")
|
|
45
|
+
await self.transport.connect()
|
|
46
|
+
logger.debug("IPC connected, sending hello...")
|
|
47
|
+
|
|
48
|
+
request_id = uuid.uuid4().hex
|
|
49
|
+
await self.transport.send(HostelMessage(id=request_id, type="system", action="hello", system_payload={}))
|
|
50
|
+
|
|
51
|
+
msg = await self.transport.receive()
|
|
52
|
+
assert msg.action == "ready"
|
|
53
|
+
|
|
54
|
+
logger.debug("Received ready from server, starting receiver...")
|
|
55
|
+
asyncio.create_task(self._receiver())
|
|
56
|
+
logger.info("Hostel Client started.")
|
|
57
|
+
|
|
58
|
+
async def _receiver(self) -> None:
|
|
59
|
+
while True:
|
|
60
|
+
msg = await self.transport.receive()
|
|
61
|
+
request_id = msg.id
|
|
62
|
+
if request_id and request_id in self._pending:
|
|
63
|
+
await self._pending[request_id].put(msg)
|
|
64
|
+
|
|
65
|
+
# ── Agent operations ─────────────────────────────────────────────────
|
|
66
|
+
|
|
67
|
+
async def list_agents(self) -> list[dict[str, Any]]:
|
|
68
|
+
request_id = uuid.uuid4().hex
|
|
69
|
+
queue: asyncio.Queue[HostelMessage] = asyncio.Queue()
|
|
70
|
+
self._pending[request_id] = queue
|
|
71
|
+
|
|
72
|
+
await self.transport.send(
|
|
73
|
+
HostelMessage(
|
|
74
|
+
id=request_id,
|
|
75
|
+
type="request",
|
|
76
|
+
action="list_agents",
|
|
77
|
+
agent_list_request=ListAgentsRequest(),
|
|
78
|
+
)
|
|
79
|
+
)
|
|
80
|
+
msg = await queue.get()
|
|
81
|
+
del self._pending[request_id]
|
|
82
|
+
|
|
83
|
+
assert msg.agent_list_response is not None
|
|
84
|
+
return msg.agent_list_response.agents
|
|
85
|
+
|
|
86
|
+
# ── Chat operations ──────────────────────────────────────────────────
|
|
87
|
+
|
|
88
|
+
async def chat(self, agent_name: str, messages: list[ChatMessage]) -> AsyncGenerator[ChatResponse, None]:
|
|
89
|
+
request_id = uuid.uuid4().hex
|
|
90
|
+
queue: asyncio.Queue[HostelMessage] = asyncio.Queue()
|
|
91
|
+
self._pending[request_id] = queue
|
|
92
|
+
|
|
93
|
+
await self.transport.send(
|
|
94
|
+
HostelMessage(
|
|
95
|
+
id=request_id,
|
|
96
|
+
type="request",
|
|
97
|
+
action="chat",
|
|
98
|
+
chat_request=ChatRequest(agent_name=agent_name, messages=messages),
|
|
99
|
+
)
|
|
100
|
+
)
|
|
101
|
+
|
|
102
|
+
while True:
|
|
103
|
+
msg = await queue.get()
|
|
104
|
+
assert msg.chat_response_chunk is not None
|
|
105
|
+
chunk = msg.chat_response_chunk
|
|
106
|
+
|
|
107
|
+
logger.debug(f"Received chat chunk: {chunk}")
|
|
108
|
+
yield chunk
|
|
109
|
+
|
|
110
|
+
if chunk.role == "agent" and chunk.scope in ("complete", "error"):
|
|
111
|
+
break
|
|
112
|
+
|
|
113
|
+
del self._pending[request_id]
|
|
114
|
+
|
|
115
|
+
# ── Task operations ──────────────────────────────────────────────────
|
|
116
|
+
|
|
117
|
+
async def create_task(
|
|
118
|
+
self,
|
|
119
|
+
agent_name: str,
|
|
120
|
+
prompt: str,
|
|
121
|
+
start_datetime: str,
|
|
122
|
+
webhook_url: str | None = None,
|
|
123
|
+
) -> str:
|
|
124
|
+
request_id = uuid.uuid4().hex
|
|
125
|
+
await self.transport.send(
|
|
126
|
+
HostelMessage(
|
|
127
|
+
id=request_id,
|
|
128
|
+
type="request",
|
|
129
|
+
action="create_task",
|
|
130
|
+
task_create=CreateTaskRequest(
|
|
131
|
+
agent_name=agent_name,
|
|
132
|
+
prompt=prompt,
|
|
133
|
+
start_datetime=start_datetime,
|
|
134
|
+
webhook_url=webhook_url,
|
|
135
|
+
),
|
|
136
|
+
)
|
|
137
|
+
)
|
|
138
|
+
|
|
139
|
+
msg = await self.transport.receive()
|
|
140
|
+
assert msg.task_create_response is not None
|
|
141
|
+
assert msg.task_create_response.task is not None
|
|
142
|
+
task_id = msg.task_create_response.task.id
|
|
143
|
+
logger.debug(f"Created task with ID: {task_id}")
|
|
144
|
+
return task_id
|
|
145
|
+
|
|
146
|
+
async def list_tasks(self) -> list[TaskData]:
|
|
147
|
+
request_id = uuid.uuid4().hex
|
|
148
|
+
await self.transport.send(
|
|
149
|
+
HostelMessage(
|
|
150
|
+
id=request_id,
|
|
151
|
+
type="request",
|
|
152
|
+
action="list_tasks",
|
|
153
|
+
task_list=ListTasksRequest(),
|
|
154
|
+
)
|
|
155
|
+
)
|
|
156
|
+
|
|
157
|
+
msg = await self.transport.receive()
|
|
158
|
+
assert msg.task_list_response is not None
|
|
159
|
+
tasks = msg.task_list_response.tasks
|
|
160
|
+
logger.debug(f"Retrieved {len(tasks)} tasks")
|
|
161
|
+
return tasks
|
|
162
|
+
|
|
163
|
+
async def get_task(self, task_id: str) -> TaskData | None:
|
|
164
|
+
request_id = uuid.uuid4().hex
|
|
165
|
+
await self.transport.send(
|
|
166
|
+
HostelMessage(
|
|
167
|
+
id=request_id,
|
|
168
|
+
type="request",
|
|
169
|
+
action="get_task",
|
|
170
|
+
task_get=GetTaskRequest(task_id=task_id),
|
|
171
|
+
)
|
|
172
|
+
)
|
|
173
|
+
|
|
174
|
+
msg = await self.transport.receive()
|
|
175
|
+
assert msg.task_get_response is not None
|
|
176
|
+
return msg.task_get_response.task
|
|
177
|
+
|
|
178
|
+
async def update_task(
|
|
179
|
+
self,
|
|
180
|
+
task_id: str,
|
|
181
|
+
agent_name: str | None = None,
|
|
182
|
+
prompt: str | None = None,
|
|
183
|
+
start_datetime: str | None = None,
|
|
184
|
+
webhook_url: str | None = None,
|
|
185
|
+
status: str | None = None,
|
|
186
|
+
) -> bool:
|
|
187
|
+
request_id = uuid.uuid4().hex
|
|
188
|
+
await self.transport.send(
|
|
189
|
+
HostelMessage(
|
|
190
|
+
id=request_id,
|
|
191
|
+
type="request",
|
|
192
|
+
action="update_task",
|
|
193
|
+
task_update=UpdateTaskRequest(
|
|
194
|
+
task_id=task_id,
|
|
195
|
+
agent_name=agent_name,
|
|
196
|
+
prompt=prompt,
|
|
197
|
+
start_datetime=start_datetime,
|
|
198
|
+
webhook_url=webhook_url,
|
|
199
|
+
status=status,
|
|
200
|
+
),
|
|
201
|
+
)
|
|
202
|
+
)
|
|
203
|
+
|
|
204
|
+
msg = await self.transport.receive()
|
|
205
|
+
assert msg.task_update_response is not None
|
|
206
|
+
success = msg.task_update_response.task is not None
|
|
207
|
+
logger.debug(f"Task {task_id} update {'succeeded' if success else 'failed'}")
|
|
208
|
+
return success
|
|
209
|
+
|
|
210
|
+
async def delete_task(self, task_id: str) -> bool:
|
|
211
|
+
request_id = uuid.uuid4().hex
|
|
212
|
+
await self.transport.send(
|
|
213
|
+
HostelMessage(
|
|
214
|
+
id=request_id,
|
|
215
|
+
type="request",
|
|
216
|
+
action="delete_task",
|
|
217
|
+
task_delete=DeleteTaskRequest(task_id=task_id),
|
|
218
|
+
)
|
|
219
|
+
)
|
|
220
|
+
|
|
221
|
+
msg = await self.transport.receive()
|
|
222
|
+
assert msg.task_delete_response is not None
|
|
223
|
+
success = msg.task_delete_response.success
|
|
224
|
+
logger.debug(f"Task {task_id} delete {'succeeded' if success else 'failed'}")
|
|
225
|
+
return success
|
|
226
|
+
|
|
227
|
+
# ── Component operations ─────────────────────────────────────────────
|
|
228
|
+
|
|
229
|
+
async def create_component(self, component_type: str, data: dict[str, Any]) -> bool:
|
|
230
|
+
request_id = uuid.uuid4().hex
|
|
231
|
+
await self.transport.send(
|
|
232
|
+
HostelMessage(
|
|
233
|
+
id=request_id,
|
|
234
|
+
type="request",
|
|
235
|
+
action="create_component",
|
|
236
|
+
component_create=CreateComponentRequest(component_type=component_type, data=data),
|
|
237
|
+
)
|
|
238
|
+
)
|
|
239
|
+
|
|
240
|
+
msg = await self.transport.receive()
|
|
241
|
+
assert msg.component_create_response is not None
|
|
242
|
+
return msg.component_create_response.success
|
|
243
|
+
|
|
244
|
+
async def get_component(self, component_type: str, name: str) -> dict[str, Any] | None:
|
|
245
|
+
request_id = uuid.uuid4().hex
|
|
246
|
+
await self.transport.send(
|
|
247
|
+
HostelMessage(
|
|
248
|
+
id=request_id,
|
|
249
|
+
type="request",
|
|
250
|
+
action="get_component",
|
|
251
|
+
component_get=GetComponentRequest(component_type=component_type, name=name),
|
|
252
|
+
)
|
|
253
|
+
)
|
|
254
|
+
|
|
255
|
+
msg = await self.transport.receive()
|
|
256
|
+
assert msg.component_get_response is not None
|
|
257
|
+
return msg.component_get_response.data
|
|
258
|
+
|
|
259
|
+
async def list_components(self, component_type: str) -> list[dict[str, Any]]:
|
|
260
|
+
request_id = uuid.uuid4().hex
|
|
261
|
+
await self.transport.send(
|
|
262
|
+
HostelMessage(
|
|
263
|
+
id=request_id,
|
|
264
|
+
type="request",
|
|
265
|
+
action="list_components",
|
|
266
|
+
component_list=ListComponentsRequest(component_type=component_type),
|
|
267
|
+
)
|
|
268
|
+
)
|
|
269
|
+
|
|
270
|
+
msg = await self.transport.receive()
|
|
271
|
+
assert msg.component_list_response is not None
|
|
272
|
+
return msg.component_list_response.components
|
|
273
|
+
|
|
274
|
+
async def update_component(self, component_type: str, name: str, data: dict[str, Any]) -> bool:
|
|
275
|
+
request_id = uuid.uuid4().hex
|
|
276
|
+
await self.transport.send(
|
|
277
|
+
HostelMessage(
|
|
278
|
+
id=request_id,
|
|
279
|
+
type="request",
|
|
280
|
+
action="update_component",
|
|
281
|
+
component_update=UpdateComponentRequest(component_type=component_type, name=name, data=data),
|
|
282
|
+
)
|
|
283
|
+
)
|
|
284
|
+
|
|
285
|
+
msg = await self.transport.receive()
|
|
286
|
+
assert msg.component_update_response is not None
|
|
287
|
+
return msg.component_update_response.success
|
|
288
|
+
|
|
289
|
+
async def delete_component(self, component_type: str, name: str) -> bool:
|
|
290
|
+
request_id = uuid.uuid4().hex
|
|
291
|
+
await self.transport.send(
|
|
292
|
+
HostelMessage(
|
|
293
|
+
id=request_id,
|
|
294
|
+
type="request",
|
|
295
|
+
action="delete_component",
|
|
296
|
+
component_delete=DeleteComponentRequest(component_type=component_type, name=name),
|
|
297
|
+
)
|
|
298
|
+
)
|
|
299
|
+
|
|
300
|
+
msg = await self.transport.receive()
|
|
301
|
+
assert msg.component_delete_response is not None
|
|
302
|
+
return msg.component_delete_response.success
|
|
@@ -1,7 +1,11 @@
|
|
|
1
|
-
"""
|
|
1
|
+
"""hostel.protocol — Pydantic convenience layer for the Hostel protocol."""
|
|
2
2
|
|
|
3
|
-
from
|
|
4
|
-
|
|
3
|
+
from pkgutil import extend_path
|
|
4
|
+
|
|
5
|
+
__path__ = extend_path(__path__, __name__)
|
|
6
|
+
|
|
7
|
+
from hostel.protocol.converter import proto_to_pydantic, pydantic_to_proto
|
|
8
|
+
from hostel.protocol.models import (
|
|
5
9
|
AIChatMessage,
|
|
6
10
|
ChatMessage,
|
|
7
11
|
ChatRequest,
|
|
@@ -16,7 +16,7 @@ from hostel.protocol.v1 import (
|
|
|
16
16
|
)
|
|
17
17
|
from pydantic import BaseModel
|
|
18
18
|
|
|
19
|
-
from
|
|
19
|
+
from hostel.protocol import models
|
|
20
20
|
|
|
21
21
|
# ---------------------------------------------------------------------------
|
|
22
22
|
# Registry: Pydantic model class <-> Protobuf message class
|
|
@@ -0,0 +1,11 @@
|
|
|
1
|
+
"""hostel.transport — Transport abstractions and implementations."""
|
|
2
|
+
|
|
3
|
+
from hostel.transport.transport import ClientId, ServerTransport, Transport
|
|
4
|
+
|
|
5
|
+
__all__ = [
|
|
6
|
+
"ClientId",
|
|
7
|
+
"Transport",
|
|
8
|
+
"ServerTransport",
|
|
9
|
+
]
|
|
10
|
+
|
|
11
|
+
DEFAULT_IPC_ENDPOINT = "ipc:///tmp/hostel.ipc"
|
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
"""Abstract base classes for Hostel transports."""
|
|
2
|
+
|
|
3
|
+
from abc import ABC, abstractmethod
|
|
4
|
+
|
|
5
|
+
from hostel.protocol.models import HostelMessage
|
|
6
|
+
|
|
7
|
+
ClientId = bytes
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
class Transport(ABC):
|
|
11
|
+
|
|
12
|
+
@abstractmethod
|
|
13
|
+
async def connect(self) -> None:
|
|
14
|
+
pass
|
|
15
|
+
|
|
16
|
+
@abstractmethod
|
|
17
|
+
async def send(self, message: HostelMessage) -> None:
|
|
18
|
+
pass
|
|
19
|
+
|
|
20
|
+
@abstractmethod
|
|
21
|
+
async def receive(self) -> HostelMessage:
|
|
22
|
+
pass
|
|
23
|
+
|
|
24
|
+
@abstractmethod
|
|
25
|
+
async def close(self) -> None:
|
|
26
|
+
pass
|
|
27
|
+
|
|
28
|
+
|
|
29
|
+
class ServerTransport(ABC):
|
|
30
|
+
|
|
31
|
+
@abstractmethod
|
|
32
|
+
async def connect(self) -> None:
|
|
33
|
+
pass
|
|
34
|
+
|
|
35
|
+
@abstractmethod
|
|
36
|
+
async def send(self, client_id: ClientId, message: HostelMessage) -> None:
|
|
37
|
+
pass
|
|
38
|
+
|
|
39
|
+
@abstractmethod
|
|
40
|
+
async def receive(self) -> tuple[ClientId, HostelMessage]:
|
|
41
|
+
pass
|
|
42
|
+
|
|
43
|
+
@abstractmethod
|
|
44
|
+
async def close(self) -> None:
|
|
45
|
+
pass
|
|
@@ -0,0 +1,119 @@
|
|
|
1
|
+
"""ZeroMQ transport implementations for Hostel."""
|
|
2
|
+
|
|
3
|
+
import logging
|
|
4
|
+
import uuid
|
|
5
|
+
|
|
6
|
+
import zmq
|
|
7
|
+
import zmq.asyncio
|
|
8
|
+
from hostel.protocol.v1 import message_pb2
|
|
9
|
+
|
|
10
|
+
from hostel.protocol.converter import proto_to_pydantic, pydantic_to_proto
|
|
11
|
+
from hostel.protocol.models import HostelMessage
|
|
12
|
+
from hostel.transport.transport import ClientId, ServerTransport, Transport
|
|
13
|
+
|
|
14
|
+
logger = logging.getLogger(__name__)
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class ZmqTransport(Transport):
|
|
18
|
+
"""ZeroMQ DEALER transport for client-side IPC."""
|
|
19
|
+
|
|
20
|
+
def __init__(
|
|
21
|
+
self,
|
|
22
|
+
endpoint: str,
|
|
23
|
+
*,
|
|
24
|
+
identity: str | None = None,
|
|
25
|
+
linger_ms: int = 0,
|
|
26
|
+
):
|
|
27
|
+
self.endpoint = endpoint
|
|
28
|
+
self.identity = identity or uuid.uuid4().hex
|
|
29
|
+
self.linger_ms = linger_ms
|
|
30
|
+
|
|
31
|
+
self._ctx: zmq.asyncio.Context | None = None
|
|
32
|
+
self._socket: zmq.asyncio.Socket | None = None
|
|
33
|
+
|
|
34
|
+
async def connect(self) -> None:
|
|
35
|
+
if self._socket is not None:
|
|
36
|
+
return
|
|
37
|
+
|
|
38
|
+
self._ctx = zmq.asyncio.Context.instance()
|
|
39
|
+
self._socket = self._ctx.socket(zmq.DEALER)
|
|
40
|
+
self._socket.setsockopt(zmq.IDENTITY, self.identity.encode())
|
|
41
|
+
self._socket.setsockopt(zmq.LINGER, self.linger_ms)
|
|
42
|
+
self._socket.connect(self.endpoint)
|
|
43
|
+
|
|
44
|
+
async def send(self, message: HostelMessage) -> None:
|
|
45
|
+
if self._socket is None:
|
|
46
|
+
raise RuntimeError("Transport not connected")
|
|
47
|
+
|
|
48
|
+
proto = pydantic_to_proto(message)
|
|
49
|
+
raw = proto.SerializeToString()
|
|
50
|
+
await self._socket.send_multipart([b"", raw])
|
|
51
|
+
|
|
52
|
+
async def receive(self) -> HostelMessage:
|
|
53
|
+
if self._socket is None:
|
|
54
|
+
raise RuntimeError("Transport not connected")
|
|
55
|
+
|
|
56
|
+
frames = await self._socket.recv_multipart()
|
|
57
|
+
if len(frames) != 2 or frames[0] != b"":
|
|
58
|
+
raise ValueError(f"Invalid DEALER frame: {frames}")
|
|
59
|
+
|
|
60
|
+
_, raw = frames
|
|
61
|
+
proto = message_pb2.HostelMessage()
|
|
62
|
+
proto.ParseFromString(raw)
|
|
63
|
+
return proto_to_pydantic(proto) # type: ignore[return-value]
|
|
64
|
+
|
|
65
|
+
async def close(self) -> None:
|
|
66
|
+
if self._socket is not None:
|
|
67
|
+
self._socket.close()
|
|
68
|
+
self._socket = None
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
class ZmqRouterTransport(ServerTransport):
|
|
72
|
+
"""ZeroMQ ROUTER transport for the service side."""
|
|
73
|
+
|
|
74
|
+
def __init__(
|
|
75
|
+
self,
|
|
76
|
+
endpoint: str,
|
|
77
|
+
*,
|
|
78
|
+
linger_ms: int = 0,
|
|
79
|
+
):
|
|
80
|
+
self.endpoint = endpoint
|
|
81
|
+
self.linger_ms = linger_ms
|
|
82
|
+
|
|
83
|
+
self._ctx: zmq.asyncio.Context | None = None
|
|
84
|
+
self._socket: zmq.asyncio.Socket | None = None
|
|
85
|
+
|
|
86
|
+
async def connect(self) -> None:
|
|
87
|
+
if self._socket is not None:
|
|
88
|
+
return
|
|
89
|
+
|
|
90
|
+
self._ctx = zmq.asyncio.Context.instance()
|
|
91
|
+
self._socket = self._ctx.socket(zmq.ROUTER)
|
|
92
|
+
self._socket.setsockopt(zmq.LINGER, self.linger_ms)
|
|
93
|
+
self._socket.bind(self.endpoint)
|
|
94
|
+
|
|
95
|
+
async def receive(self) -> tuple[ClientId, HostelMessage]:
|
|
96
|
+
if self._socket is None:
|
|
97
|
+
raise RuntimeError("Transport not bound")
|
|
98
|
+
|
|
99
|
+
frames = await self._socket.recv_multipart()
|
|
100
|
+
if len(frames) != 3 or frames[1] != b"":
|
|
101
|
+
raise ValueError(f"Invalid ROUTER frame: {frames}")
|
|
102
|
+
|
|
103
|
+
client_id, _, raw = frames
|
|
104
|
+
proto = message_pb2.HostelMessage()
|
|
105
|
+
proto.ParseFromString(raw)
|
|
106
|
+
return client_id, proto_to_pydantic(proto) # type: ignore[return-value]
|
|
107
|
+
|
|
108
|
+
async def send(self, client_id: ClientId, message: HostelMessage) -> None:
|
|
109
|
+
if self._socket is None:
|
|
110
|
+
raise RuntimeError("Transport not bound")
|
|
111
|
+
|
|
112
|
+
proto = pydantic_to_proto(message)
|
|
113
|
+
raw = proto.SerializeToString()
|
|
114
|
+
await self._socket.send_multipart([client_id, b"", raw])
|
|
115
|
+
|
|
116
|
+
async def close(self) -> None:
|
|
117
|
+
if self._socket is not None:
|
|
118
|
+
self._socket.close()
|
|
119
|
+
self._socket = None
|
|
File without changes
|