embed-client-dev 4.0.78__py3-none-any.whl
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.
- embed_client/__init__.py +36 -0
- embed_client/__main__.py +15 -0
- embed_client/_version.py +3 -0
- embed_client/adapter_config_factory.py +213 -0
- embed_client/adapter_config_normalizer.py +135 -0
- embed_client/adapter_generator_bridge.py +166 -0
- embed_client/auth.py +189 -0
- embed_client/cli.py +483 -0
- embed_client/cli_core.py +427 -0
- embed_client/client_factory.py +515 -0
- embed_client/client_factory_examples.py +375 -0
- embed_client/client_factory_helpers.py +59 -0
- embed_client/config.py +408 -0
- embed_client/config_defaults.py +77 -0
- embed_client/config_examples.py +228 -0
- embed_client/config_generator.py +652 -0
- embed_client/config_generator_cli.py +151 -0
- embed_client/config_presets.py +245 -0
- embed_client/config_validator.py +246 -0
- embed_client/config_validator_cli.py +141 -0
- embed_client/constants.py +52 -0
- embed_client/demo_cli.py +78 -0
- embed_client/demo_security_cli.py +206 -0
- embed_client/embed_file_output_validate.py +348 -0
- embed_client/embed_file_status_store.py +80 -0
- embed_client/embedding_client.py +1929 -0
- embed_client/embedding_client_helpers.py +105 -0
- embed_client/example_async_usage.py +357 -0
- embed_client/example_async_usage_demo.py +282 -0
- embed_client/example_async_usage_ru.py +340 -0
- embed_client/example_async_usage_ru_demo.py +313 -0
- embed_client/exceptions.py +122 -0
- embed_client/identifier_types.py +56 -0
- embed_client/json_embed_file_preflight.py +128 -0
- embed_client/local_json_embed_file.py +73 -0
- embed_client/md_embed_file_prep.py +197 -0
- embed_client/response_normalizer.py +160 -0
- embed_client/response_parsers.py +219 -0
- embed_client/security_cli.py +200 -0
- embed_client/security_cli_core.py +227 -0
- embed_client/ssl_examples.py +278 -0
- embed_client/ssl_manager.py +141 -0
- embed_client_dev-4.0.78.dist-info/METADATA +155 -0
- embed_client_dev-4.0.78.dist-info/RECORD +48 -0
- embed_client_dev-4.0.78.dist-info/WHEEL +5 -0
- embed_client_dev-4.0.78.dist-info/entry_points.txt +4 -0
- embed_client_dev-4.0.78.dist-info/licenses/LICENSE +21 -0
- embed_client_dev-4.0.78.dist-info/top_level.txt +1 -0
|
@@ -0,0 +1,1929 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Embedding service client for interacting with embedding service API.
|
|
3
|
+
|
|
4
|
+
Author: Vasiliy Zdanovskiy
|
|
5
|
+
email: vasilyvz@gmail.com
|
|
6
|
+
"""
|
|
7
|
+
|
|
8
|
+
import asyncio
|
|
9
|
+
import copy
|
|
10
|
+
import hashlib
|
|
11
|
+
import inspect
|
|
12
|
+
import logging
|
|
13
|
+
import os
|
|
14
|
+
import tempfile
|
|
15
|
+
from pathlib import Path
|
|
16
|
+
from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Union
|
|
17
|
+
|
|
18
|
+
if TYPE_CHECKING:
|
|
19
|
+
from embed_client.config import ClientConfig
|
|
20
|
+
|
|
21
|
+
from embed_client.embed_file_status_store import EmbedFileJobStatusStore
|
|
22
|
+
|
|
23
|
+
from mcp_proxy_adapter.client.jsonrpc_client import JsonRpcClient
|
|
24
|
+
from mcp_proxy_adapter.transfer.models import TransferReceipt
|
|
25
|
+
from embed_client.embedding_client_helpers import (
|
|
26
|
+
list_commands_via_client,
|
|
27
|
+
wait_for_job_poll,
|
|
28
|
+
)
|
|
29
|
+
from embed_client.constants import EMBED_COMMAND_SCHEMA_DEFAULT_ERROR_POLICY
|
|
30
|
+
from embed_client.exceptions import EmbedAPIError, EmbedClientError, EmbedError
|
|
31
|
+
|
|
32
|
+
logger = logging.getLogger(__name__)
|
|
33
|
+
|
|
34
|
+
try: # mcp-proxy-adapter >= 8.10 raises this when /ws closes mid-wait
|
|
35
|
+
from mcp_proxy_adapter.client.jsonrpc_client.exceptions import (
|
|
36
|
+
ClientWebSocketClosedError as _AdapterWsClosedError,
|
|
37
|
+
)
|
|
38
|
+
except ImportError: # pragma: no cover - older adapter builds
|
|
39
|
+
_AdapterWsClosedError = None # type: ignore[assignment,misc]
|
|
40
|
+
|
|
41
|
+
# Terminal queue states reported by queue_get_job_status / embed_job_status.
|
|
42
|
+
_POLL_COMPLETED_STATES = ("completed", "success")
|
|
43
|
+
_POLL_FAILED_STATES = ("failed", "error")
|
|
44
|
+
_POLL_STOPPED_STATES = ("stopped", "cancelled", "canceled")
|
|
45
|
+
# Marks an input_limit_update cell that the caller did not name (distinct from null).
|
|
46
|
+
_NOT_GIVEN: Any = object()
|
|
47
|
+
|
|
48
|
+
|
|
49
|
+
def _is_ws_closed_before_terminal(exc: BaseException) -> bool:
|
|
50
|
+
"""Whether *exc* means the /ws stream closed before a terminal job event.
|
|
51
|
+
|
|
52
|
+
The job itself is unaffected by a dropped stream: the server keeps
|
|
53
|
+
executing it and ``queue_get_job_status`` still reports its terminal
|
|
54
|
+
state. A heartbeat miss while the server is busy closes the stream
|
|
55
|
+
without any job having failed.
|
|
56
|
+
"""
|
|
57
|
+
if _AdapterWsClosedError is not None and isinstance(exc, _AdapterWsClosedError):
|
|
58
|
+
return True
|
|
59
|
+
return "websocket closed" in str(exc).lower()
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _client_has_async_method(client: Any, name: str) -> bool:
|
|
63
|
+
fn = getattr(client, name, None)
|
|
64
|
+
return fn is not None and inspect.iscoroutinefunction(fn)
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
def _uses_adapter_job_message_registry(client: Any) -> bool:
|
|
68
|
+
"""True if *client* exposes async ``get_job_messages`` (adapter job registry)."""
|
|
69
|
+
return _client_has_async_method(client, "get_job_messages")
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
async def _maybe_await(value: Any) -> None:
|
|
73
|
+
if inspect.isawaitable(value):
|
|
74
|
+
await value
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def _sha256_hex_file(path: str) -> str:
|
|
78
|
+
"""SHA-256 hex digest of file at ``path`` (streaming read)."""
|
|
79
|
+
h = hashlib.sha256()
|
|
80
|
+
with open(path, "rb") as f:
|
|
81
|
+
for chunk in iter(lambda: f.read(1024 * 1024), b""):
|
|
82
|
+
h.update(chunk)
|
|
83
|
+
return h.hexdigest()
|
|
84
|
+
|
|
85
|
+
|
|
86
|
+
def _normalize_hex_checksum(value: Any) -> str:
|
|
87
|
+
if not isinstance(value, str):
|
|
88
|
+
return ""
|
|
89
|
+
return value.strip().lower()
|
|
90
|
+
|
|
91
|
+
|
|
92
|
+
def _merge_jsonrpc_result_data_layers(raw: Dict[str, Any]) -> Dict[str, Any]:
|
|
93
|
+
"""
|
|
94
|
+
Merge ``data`` dicts found along nested JSON-RPC-style ``result`` wrappers.
|
|
95
|
+
|
|
96
|
+
Walks *raw* depth-first: at each dict node, if ``data`` is a dict it is
|
|
97
|
+
collected; then ``result`` is visited if it is a dict. Shallow ``data``
|
|
98
|
+
layers are merged first; deeper layers overwrite on key conflicts so full
|
|
99
|
+
command payloads (e.g. ``input_checksum_value`` under inner
|
|
100
|
+
``result.result.data``) win over thin outer ``data`` envelopes.
|
|
101
|
+
"""
|
|
102
|
+
layers: List[Dict[str, Any]] = []
|
|
103
|
+
|
|
104
|
+
def visit(node: Any) -> None:
|
|
105
|
+
if not isinstance(node, dict):
|
|
106
|
+
return
|
|
107
|
+
d = node.get("data")
|
|
108
|
+
if isinstance(d, dict):
|
|
109
|
+
layers.append(d)
|
|
110
|
+
nxt = node.get("result")
|
|
111
|
+
if isinstance(nxt, dict):
|
|
112
|
+
visit(nxt)
|
|
113
|
+
|
|
114
|
+
visit(raw)
|
|
115
|
+
if not layers:
|
|
116
|
+
return raw
|
|
117
|
+
merged: Dict[str, Any] = {}
|
|
118
|
+
for d in layers:
|
|
119
|
+
merged = {**merged, **d}
|
|
120
|
+
return merged
|
|
121
|
+
|
|
122
|
+
|
|
123
|
+
def _embed_file_queue_result_error_message(data: Dict[str, Any]) -> Optional[str]:
|
|
124
|
+
"""
|
|
125
|
+
If a queued ``embed_file`` job completed with ``success: false`` in the
|
|
126
|
+
normalized WebSocket payload, return a human-readable error message.
|
|
127
|
+
"""
|
|
128
|
+
inner = data.get("result")
|
|
129
|
+
if not isinstance(inner, dict):
|
|
130
|
+
return None
|
|
131
|
+
if inner.get("success") is True:
|
|
132
|
+
return None
|
|
133
|
+
err = inner.get("error")
|
|
134
|
+
if inner.get("success") is not False and err is None:
|
|
135
|
+
return None
|
|
136
|
+
if isinstance(err, dict):
|
|
137
|
+
msg = err.get("message")
|
|
138
|
+
if isinstance(msg, str) and msg.strip():
|
|
139
|
+
return msg.strip()
|
|
140
|
+
return str(err) if err else "embed_file job failed"
|
|
141
|
+
if err is not None:
|
|
142
|
+
return str(err)
|
|
143
|
+
return "embed_file job failed"
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
#: Adapter terminal events that carry a failure. ``delivery_failed`` is what the
|
|
147
|
+
#: adapter pushes when a queued or execute_async command fails or times out
|
|
148
|
+
#: (its ``error`` field holds the real reason, e.g. "Command timed out after 300.0s").
|
|
149
|
+
_FAILED_TERMINAL_EVENTS = frozenset({"job_failed", "delivery_failed"})
|
|
150
|
+
#: Adapter terminal events for a job that was cancelled or stopped before completion.
|
|
151
|
+
_CANCELLED_TERMINAL_EVENTS = frozenset({"delivery_cancelled", "delivery_stopped"})
|
|
152
|
+
|
|
153
|
+
|
|
154
|
+
def _normalize_terminal_job_payload(payload: Dict[str, Any]) -> Dict[str, Any]:
|
|
155
|
+
"""
|
|
156
|
+
Normalize adapter /ws terminal payload to embed result ``data`` shape.
|
|
157
|
+
|
|
158
|
+
Failure events (``job_failed``, ``delivery_failed``) and cancellation events
|
|
159
|
+
(``delivery_cancelled``, ``delivery_stopped``) raise :class:`EmbedClientError`
|
|
160
|
+
carrying the event's ``error`` text. Used by :meth:`wait_for_job` (adapter path)
|
|
161
|
+
and embed_file WS completion.
|
|
162
|
+
"""
|
|
163
|
+
ev = payload.get("event") or payload.get("type")
|
|
164
|
+
if ev in _FAILED_TERMINAL_EVENTS:
|
|
165
|
+
err = payload.get("error", "Unknown error")
|
|
166
|
+
raise EmbedClientError(f"Job failed: {err}", error=payload.get("error"))
|
|
167
|
+
if ev in _CANCELLED_TERMINAL_EVENTS:
|
|
168
|
+
err = payload.get("error", "Unknown error")
|
|
169
|
+
raise EmbedClientError(f"Job cancelled: {err}", error=payload.get("error"))
|
|
170
|
+
if ev in ("job_completed", "job_completed_success", "delivery_completed"):
|
|
171
|
+
raw = payload.get("result") or {}
|
|
172
|
+
error = _find_failed_command_result(raw)
|
|
173
|
+
if error is not None:
|
|
174
|
+
raise EmbedClientError(
|
|
175
|
+
f"Job command failed: {error}",
|
|
176
|
+
error=_find_failed_command_error_object(raw),
|
|
177
|
+
)
|
|
178
|
+
return _merge_jsonrpc_result_data_layers(raw)
|
|
179
|
+
raise EmbedClientError(
|
|
180
|
+
f"Unexpected WebSocket event: {ev} (payload keys: {list(payload.keys())})"
|
|
181
|
+
)
|
|
182
|
+
|
|
183
|
+
|
|
184
|
+
def _missing_from_worker_registry(exc: Exception) -> bool:
|
|
185
|
+
"""Whether a queued command failed because the worker has no such command.
|
|
186
|
+
|
|
187
|
+
The adapter runs a queued command in a child process whose registry is
|
|
188
|
+
rebuilt from the project's own modules, so the adapter's own builtins —
|
|
189
|
+
help, health, echo — are absent there even though they answer over
|
|
190
|
+
JSON-RPC. Those commands are not broken and must not be reported as such.
|
|
191
|
+
"""
|
|
192
|
+
details = getattr(exc, "details", None)
|
|
193
|
+
if not isinstance(details, dict):
|
|
194
|
+
return False
|
|
195
|
+
status = details.get("result_status")
|
|
196
|
+
if not isinstance(status, dict):
|
|
197
|
+
return False
|
|
198
|
+
return "not found" in str(status.get("description") or "").lower()
|
|
199
|
+
|
|
200
|
+
|
|
201
|
+
def _find_failed_command_error_object(value: Any) -> Any:
|
|
202
|
+
"""Return the raw error object of the first nested failed command envelope.
|
|
203
|
+
|
|
204
|
+
Sibling of :func:`_find_failed_command_result` (same traversal) that returns
|
|
205
|
+
the error exactly as received (dict, str or ``None``) instead of a message.
|
|
206
|
+
"""
|
|
207
|
+
if not isinstance(value, dict):
|
|
208
|
+
return None
|
|
209
|
+
if value.get("success") is False:
|
|
210
|
+
return value.get("error")
|
|
211
|
+
for key in ("result", "data"):
|
|
212
|
+
if _find_failed_command_result(value.get(key)) is not None:
|
|
213
|
+
return _find_failed_command_error_object(value.get(key))
|
|
214
|
+
return None
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
def _find_failed_command_result(value: Any) -> Optional[str]:
|
|
218
|
+
"""Return the first nested command error from a completed queue envelope.
|
|
219
|
+
|
|
220
|
+
A queue job whose command failed still arrives as ``job_completed``: the
|
|
221
|
+
refusal sits inside the envelope with ``success: False``. Without this the
|
|
222
|
+
caller would read a failed command as a successful job.
|
|
223
|
+
"""
|
|
224
|
+
if not isinstance(value, dict):
|
|
225
|
+
return None
|
|
226
|
+
if value.get("success") is False:
|
|
227
|
+
error = value.get("error")
|
|
228
|
+
if isinstance(error, dict):
|
|
229
|
+
return str(error.get("message") or error)
|
|
230
|
+
return str(error or "Unknown command error")
|
|
231
|
+
for key in ("result", "data"):
|
|
232
|
+
nested = _find_failed_command_result(value.get(key))
|
|
233
|
+
if nested is not None:
|
|
234
|
+
return nested
|
|
235
|
+
return None
|
|
236
|
+
|
|
237
|
+
|
|
238
|
+
class EmbeddingClient:
|
|
239
|
+
"""
|
|
240
|
+
Client for embedding service API.
|
|
241
|
+
|
|
242
|
+
Provides high-level interface for:
|
|
243
|
+
- Generating embeddings (via queue)
|
|
244
|
+
- Checking job status
|
|
245
|
+
- Listing available models
|
|
246
|
+
- Health checks
|
|
247
|
+
|
|
248
|
+
This client always uses WebSocket (/ws) for job result delivery when
|
|
249
|
+
wait=True (server exposes /ws; job_push is enabled when queue is on).
|
|
250
|
+
No HTTP polling for completion.
|
|
251
|
+
|
|
252
|
+
**embed_file status store (optional):** pass ``enable_embed_file_status_store=True``
|
|
253
|
+
to accumulate per-``job_id`` submit payload and WebSocket frames for
|
|
254
|
+
:meth:`embed_file`. Query with :meth:`get_embed_file_status_snapshot`,
|
|
255
|
+
:meth:`list_embed_file_status_messages`, :meth:`alist_embed_file_status_messages`,
|
|
256
|
+
:meth:`list_embed_file_status_job_ids`, and clear with
|
|
257
|
+
:meth:`clear_embed_file_status`. With mcp-proxy-adapter ``JobTrackingMixin``
|
|
258
|
+
(>= 6.10), WS completion uses :meth:`JsonRpcClient.wait_for_job` and
|
|
259
|
+
:meth:`clear_embed_file_status` also clears the adapter registry; from async
|
|
260
|
+
code prefer :meth:`alist_embed_file_status_messages` and
|
|
261
|
+
:meth:`aget_embed_file_status_snapshot` to read adapter history.
|
|
262
|
+
Default is off; behavior is unchanged when disabled.
|
|
263
|
+
|
|
264
|
+
With **mcp-proxy-adapter >= 8.10**, a synchronous command that hits the server
|
|
265
|
+
sync cap is **re-queued** instead of failing the HTTP/JSON-RPC call; the JSON
|
|
266
|
+
payload may include ``queued_after_timeout`` (see :meth:`embed` with
|
|
267
|
+
``wait=False``).
|
|
268
|
+
|
|
269
|
+
Example:
|
|
270
|
+
client = EmbeddingClient(
|
|
271
|
+
host="127.0.0.1", port=8080, enable_embed_file_status_store=True
|
|
272
|
+
)
|
|
273
|
+
await client.embed_file("/srv/data.txt", wait=True)
|
|
274
|
+
snap = await client.aget_embed_file_status_snapshot("job-uuid")
|
|
275
|
+
"""
|
|
276
|
+
|
|
277
|
+
def __init__(
|
|
278
|
+
self,
|
|
279
|
+
protocol: str = "http",
|
|
280
|
+
host: str = "127.0.0.1",
|
|
281
|
+
port: int = 8080,
|
|
282
|
+
cert: Optional[str] = None,
|
|
283
|
+
key: Optional[str] = None,
|
|
284
|
+
ca: Optional[str] = None,
|
|
285
|
+
check_hostname: bool = False,
|
|
286
|
+
token: Optional[str] = None,
|
|
287
|
+
token_header: Optional[str] = None,
|
|
288
|
+
timeout: float = 120.0,
|
|
289
|
+
*,
|
|
290
|
+
enable_embed_file_status_store: bool = False,
|
|
291
|
+
verify_server: bool = True,
|
|
292
|
+
):
|
|
293
|
+
"""
|
|
294
|
+
Initialize embedding client.
|
|
295
|
+
|
|
296
|
+
Args:
|
|
297
|
+
protocol: Transport protocol (http, https, mtls)
|
|
298
|
+
host: Server hostname
|
|
299
|
+
port: Server port
|
|
300
|
+
cert: Path to client certificate file (for mTLS)
|
|
301
|
+
key: Path to client private key file (for mTLS)
|
|
302
|
+
ca: Path to CA certificate file (for mTLS)
|
|
303
|
+
check_hostname: Whether to verify hostname in SSL connections
|
|
304
|
+
token: Authentication token
|
|
305
|
+
token_header: Header name for authentication token
|
|
306
|
+
timeout: Request/timeout in seconds for JSON-RPC and wait operations
|
|
307
|
+
enable_embed_file_status_store: If True, keep an in-memory record of
|
|
308
|
+
each ``embed_file`` job (submit + WS messages) keyed by ``job_id``.
|
|
309
|
+
When False (default), no extra storage and query methods return
|
|
310
|
+
empty/None.
|
|
311
|
+
verify_server: When False, the constructed client does not verify the
|
|
312
|
+
server certificate or its name. Required when dialling a host by an
|
|
313
|
+
address that the certificate does not cover.
|
|
314
|
+
"""
|
|
315
|
+
self.client = JsonRpcClient(
|
|
316
|
+
protocol=protocol,
|
|
317
|
+
host=host,
|
|
318
|
+
port=port,
|
|
319
|
+
cert=cert,
|
|
320
|
+
key=key,
|
|
321
|
+
ca=ca,
|
|
322
|
+
check_hostname=check_hostname,
|
|
323
|
+
token=token,
|
|
324
|
+
token_header=token_header or ("X-API-Key" if token else None),
|
|
325
|
+
timeout=timeout,
|
|
326
|
+
)
|
|
327
|
+
if not verify_server:
|
|
328
|
+
self.client.verify = False
|
|
329
|
+
self.protocol = protocol
|
|
330
|
+
self.host = host
|
|
331
|
+
self.port = port
|
|
332
|
+
self._cert = cert
|
|
333
|
+
self._key = key
|
|
334
|
+
self._ca = ca
|
|
335
|
+
self._token = token
|
|
336
|
+
self._token_header = token_header or "X-API-Key"
|
|
337
|
+
self._timeout = float(timeout)
|
|
338
|
+
self._check_hostname = check_hostname
|
|
339
|
+
self._verify_server = verify_server
|
|
340
|
+
self._auth_method: Optional[str] = None
|
|
341
|
+
self._ssl_settings: Optional[Dict[str, Any]] = None
|
|
342
|
+
self._embed_file_status_store: Optional[EmbedFileJobStatusStore] = (
|
|
343
|
+
EmbedFileJobStatusStore() if enable_embed_file_status_store else None
|
|
344
|
+
)
|
|
345
|
+
|
|
346
|
+
@property
|
|
347
|
+
def base_url(self) -> str:
|
|
348
|
+
"""Base URL without path; mTLS is shown as ``https`` for compatibility."""
|
|
349
|
+
scheme = "https" if self.protocol in ("https", "mtls") else "http"
|
|
350
|
+
return f"{scheme}://{self.host}"
|
|
351
|
+
|
|
352
|
+
@property
|
|
353
|
+
def timeout(self) -> float:
|
|
354
|
+
"""Configured client timeout (seconds)."""
|
|
355
|
+
return self._timeout
|
|
356
|
+
|
|
357
|
+
def get_auth_method(self) -> Optional[str]:
|
|
358
|
+
"""Auth method from config (``none``, ``api_key``, ``jwt``, …), if known."""
|
|
359
|
+
return self._auth_method
|
|
360
|
+
|
|
361
|
+
def is_authenticated(self) -> bool:
|
|
362
|
+
"""Whether the client is configured for non-empty authentication."""
|
|
363
|
+
m = self._auth_method
|
|
364
|
+
if m in (None, "none"):
|
|
365
|
+
return False
|
|
366
|
+
if m == "api_key":
|
|
367
|
+
return bool(self._token)
|
|
368
|
+
if m in ("jwt", "basic", "certificate"):
|
|
369
|
+
return True
|
|
370
|
+
return bool(self._token)
|
|
371
|
+
|
|
372
|
+
def is_ssl_enabled(self) -> bool:
|
|
373
|
+
"""True if TLS is enabled in config or transport uses HTTPS/mTLS."""
|
|
374
|
+
if isinstance(self._ssl_settings, dict) and "enabled" in self._ssl_settings:
|
|
375
|
+
return bool(self._ssl_settings["enabled"])
|
|
376
|
+
return self.protocol in ("https", "mtls")
|
|
377
|
+
|
|
378
|
+
def is_mtls_enabled(self) -> bool:
|
|
379
|
+
"""True if mutual TLS is in use."""
|
|
380
|
+
return self.protocol == "mtls"
|
|
381
|
+
|
|
382
|
+
def get_ssl_config(self) -> Dict[str, Any]:
|
|
383
|
+
"""SSL subsection from client config, or a minimal derived dict."""
|
|
384
|
+
if isinstance(self._ssl_settings, dict) and self._ssl_settings:
|
|
385
|
+
return dict(self._ssl_settings)
|
|
386
|
+
return {
|
|
387
|
+
"enabled": self.is_ssl_enabled(),
|
|
388
|
+
"verify_mode": "CERT_REQUIRED",
|
|
389
|
+
"check_hostname": self._check_hostname,
|
|
390
|
+
}
|
|
391
|
+
|
|
392
|
+
def get_auth_headers(self) -> Dict[str, str]:
|
|
393
|
+
"""Return headers for API-key style auth (empty if no token)."""
|
|
394
|
+
if not self._token:
|
|
395
|
+
return {}
|
|
396
|
+
name = self._token_header or "X-API-Key"
|
|
397
|
+
return {str(name): str(self._token)}
|
|
398
|
+
|
|
399
|
+
@classmethod
|
|
400
|
+
def from_base_url_port(
|
|
401
|
+
cls,
|
|
402
|
+
base_url: str,
|
|
403
|
+
port: Optional[int] = None,
|
|
404
|
+
*,
|
|
405
|
+
timeout: float = 30.0,
|
|
406
|
+
) -> "EmbeddingClient":
|
|
407
|
+
"""Construct client from ``base_url`` and optional ``port`` (tests / scripts)."""
|
|
408
|
+
from urllib.parse import urlparse
|
|
409
|
+
|
|
410
|
+
parsed = urlparse(base_url)
|
|
411
|
+
scheme = (parsed.scheme or "http").lower()
|
|
412
|
+
if scheme not in ("http", "https"):
|
|
413
|
+
scheme = "http"
|
|
414
|
+
host = parsed.hostname or "localhost"
|
|
415
|
+
use_port = int(
|
|
416
|
+
parsed.port
|
|
417
|
+
if parsed.port is not None
|
|
418
|
+
else (port if port is not None else 8001)
|
|
419
|
+
)
|
|
420
|
+
return cls(
|
|
421
|
+
protocol="https" if scheme == "https" else "http",
|
|
422
|
+
host=host,
|
|
423
|
+
port=use_port,
|
|
424
|
+
timeout=float(timeout),
|
|
425
|
+
)
|
|
426
|
+
|
|
427
|
+
@classmethod
|
|
428
|
+
def from_config_dict(cls, config_dict: Dict[str, Any]) -> "EmbeddingClient":
|
|
429
|
+
"""Same as :func:`embedding_client_from_config_dict`."""
|
|
430
|
+
return embedding_client_from_config_dict(config_dict)
|
|
431
|
+
|
|
432
|
+
@classmethod
|
|
433
|
+
def from_config_file(cls, config_path: Union[str, Path]) -> "EmbeddingClient":
|
|
434
|
+
"""Load config from path and build a client."""
|
|
435
|
+
from embed_client.config import ClientConfig
|
|
436
|
+
|
|
437
|
+
cc = ClientConfig()
|
|
438
|
+
cc.load_from_file(str(config_path))
|
|
439
|
+
return cls.from_config(cc)
|
|
440
|
+
|
|
441
|
+
@classmethod
|
|
442
|
+
def with_auth(
|
|
443
|
+
cls,
|
|
444
|
+
base_url: str,
|
|
445
|
+
port: int,
|
|
446
|
+
auth_method: str,
|
|
447
|
+
**kwargs: Any,
|
|
448
|
+
) -> "EmbeddingClient":
|
|
449
|
+
"""Build a client from URL/port and auth method (legacy factory helper)."""
|
|
450
|
+
config_dict: Dict[str, Any] = {
|
|
451
|
+
"server": {"host": base_url, "port": port},
|
|
452
|
+
"client": {"timeout": kwargs.get("timeout", 30.0)},
|
|
453
|
+
"auth": {"method": auth_method},
|
|
454
|
+
}
|
|
455
|
+
if auth_method == "api_key":
|
|
456
|
+
if "api_keys" in kwargs:
|
|
457
|
+
config_dict["auth"]["api_keys"] = kwargs["api_keys"]
|
|
458
|
+
elif "api_key" in kwargs:
|
|
459
|
+
config_dict["auth"]["api_keys"] = {"user": kwargs["api_key"]}
|
|
460
|
+
else:
|
|
461
|
+
raise ValueError(
|
|
462
|
+
"api_keys or api_key parameter required for api_key authentication"
|
|
463
|
+
)
|
|
464
|
+
elif auth_method == "jwt":
|
|
465
|
+
required_params = ["secret", "username", "password"]
|
|
466
|
+
for param in required_params:
|
|
467
|
+
if param not in kwargs:
|
|
468
|
+
raise ValueError(
|
|
469
|
+
f"{param} parameter required for jwt authentication"
|
|
470
|
+
)
|
|
471
|
+
config_dict["auth"]["jwt"] = {
|
|
472
|
+
"secret": kwargs["secret"],
|
|
473
|
+
"username": kwargs["username"],
|
|
474
|
+
"password": kwargs["password"],
|
|
475
|
+
"expiry_hours": kwargs.get("expiry_hours", 24),
|
|
476
|
+
}
|
|
477
|
+
elif auth_method == "basic":
|
|
478
|
+
required_params = ["username", "password"]
|
|
479
|
+
for param in required_params:
|
|
480
|
+
if param not in kwargs:
|
|
481
|
+
raise ValueError(
|
|
482
|
+
f"{param} parameter required for basic authentication"
|
|
483
|
+
)
|
|
484
|
+
config_dict["auth"]["basic"] = {
|
|
485
|
+
"username": kwargs["username"],
|
|
486
|
+
"password": kwargs["password"],
|
|
487
|
+
}
|
|
488
|
+
elif auth_method == "certificate":
|
|
489
|
+
required_params = ["cert_file", "key_file"]
|
|
490
|
+
for param in required_params:
|
|
491
|
+
if param not in kwargs:
|
|
492
|
+
raise ValueError(
|
|
493
|
+
f"{param} parameter required for certificate authentication"
|
|
494
|
+
)
|
|
495
|
+
config_dict["auth"]["certificate"] = {
|
|
496
|
+
"cert_file": kwargs["cert_file"],
|
|
497
|
+
"key_file": kwargs["key_file"],
|
|
498
|
+
}
|
|
499
|
+
else:
|
|
500
|
+
raise ValueError(f"Unsupported authentication method: {auth_method}")
|
|
501
|
+
|
|
502
|
+
ssl_enabled = kwargs.get("ssl_enabled")
|
|
503
|
+
if ssl_enabled is None:
|
|
504
|
+
ssl_enabled = base_url.startswith("https://")
|
|
505
|
+
|
|
506
|
+
if ssl_enabled or any(
|
|
507
|
+
key in kwargs
|
|
508
|
+
for key in ["ca_cert_file", "cert_file", "key_file", "ssl_enabled"]
|
|
509
|
+
):
|
|
510
|
+
ssl_config: Dict[str, Any] = {
|
|
511
|
+
"enabled": ssl_enabled,
|
|
512
|
+
"verify_mode": kwargs.get("verify_mode", "CERT_REQUIRED"),
|
|
513
|
+
"check_hostname": kwargs.get("check_hostname", True),
|
|
514
|
+
"check_expiry": kwargs.get("check_expiry", True),
|
|
515
|
+
}
|
|
516
|
+
if "ca_cert_file" in kwargs:
|
|
517
|
+
ssl_config["ca_cert_file"] = kwargs["ca_cert_file"]
|
|
518
|
+
if "cert_file" in kwargs:
|
|
519
|
+
ssl_config["cert_file"] = kwargs["cert_file"]
|
|
520
|
+
if "key_file" in kwargs:
|
|
521
|
+
ssl_config["key_file"] = kwargs["key_file"]
|
|
522
|
+
config_dict["ssl"] = ssl_config
|
|
523
|
+
|
|
524
|
+
return embedding_client_from_config_dict(config_dict)
|
|
525
|
+
|
|
526
|
+
def get_embed_file_status_snapshot(self, job_id: str) -> Optional[Dict[str, Any]]:
|
|
527
|
+
"""
|
|
528
|
+
Return a copy of accumulated state for an ``embed_file`` job, or ``None``.
|
|
529
|
+
|
|
530
|
+
Only populated when the client was constructed with
|
|
531
|
+
``enable_embed_file_status_store=True`` and the job was recorded (submit
|
|
532
|
+
and/or WebSocket messages). Keys typically include ``job_id``, ``submit``,
|
|
533
|
+
``messages``, and ``last_event`` after WS traffic.
|
|
534
|
+
|
|
535
|
+
Args:
|
|
536
|
+
job_id: Queue job id returned by the server for ``embed_file``.
|
|
537
|
+
|
|
538
|
+
Returns:
|
|
539
|
+
Deep copy of internal record, or ``None`` if unknown or store disabled.
|
|
540
|
+
|
|
541
|
+
When the adapter exposes ``get_job_messages`` and this method is called
|
|
542
|
+
from synchronous code (no running event loop), ``messages`` / ``last_event``
|
|
543
|
+
are refreshed from the adapter registry and merged with the local submit
|
|
544
|
+
snapshot. Under an active asyncio loop, use :meth:`aget_embed_file_status_snapshot`
|
|
545
|
+
so messages can be read from the adapter registry (local store does not
|
|
546
|
+
duplicate WS frames when the adapter tracks them).
|
|
547
|
+
"""
|
|
548
|
+
if self._embed_file_status_store is None:
|
|
549
|
+
return None
|
|
550
|
+
local = self._embed_file_status_store.snapshot(job_id)
|
|
551
|
+
if local is None:
|
|
552
|
+
return None
|
|
553
|
+
if not _uses_adapter_job_message_registry(self.client):
|
|
554
|
+
return local
|
|
555
|
+
try:
|
|
556
|
+
asyncio.get_running_loop()
|
|
557
|
+
except RuntimeError:
|
|
558
|
+
msgs = asyncio.run(self.client.get_job_messages(job_id))
|
|
559
|
+
merged: Dict[str, Any] = dict(local)
|
|
560
|
+
merged["messages"] = copy.deepcopy(msgs)
|
|
561
|
+
if msgs:
|
|
562
|
+
merged["last_event"] = copy.deepcopy(msgs[-1])
|
|
563
|
+
return merged
|
|
564
|
+
return local
|
|
565
|
+
|
|
566
|
+
async def aget_embed_file_status_snapshot(
|
|
567
|
+
self, job_id: str
|
|
568
|
+
) -> Optional[Dict[str, Any]]:
|
|
569
|
+
"""
|
|
570
|
+
Async snapshot of embed_file job state: local ``submit`` plus adapter messages.
|
|
571
|
+
|
|
572
|
+
When ``get_job_messages`` exists on the underlying client, ``messages`` and
|
|
573
|
+
``last_event`` come from the adapter job registry; otherwise from the local
|
|
574
|
+
store. Returns ``None`` if the status store is disabled or ``job_id`` is unknown.
|
|
575
|
+
"""
|
|
576
|
+
if self._embed_file_status_store is None:
|
|
577
|
+
return None
|
|
578
|
+
local = self._embed_file_status_store.snapshot(job_id)
|
|
579
|
+
if local is None:
|
|
580
|
+
return None
|
|
581
|
+
if not _uses_adapter_job_message_registry(self.client):
|
|
582
|
+
return local
|
|
583
|
+
msgs = list(await self.client.get_job_messages(job_id))
|
|
584
|
+
merged = dict(local)
|
|
585
|
+
merged["messages"] = copy.deepcopy(msgs)
|
|
586
|
+
if msgs:
|
|
587
|
+
merged["last_event"] = copy.deepcopy(msgs[-1])
|
|
588
|
+
return merged
|
|
589
|
+
|
|
590
|
+
def list_embed_file_status_messages(self, job_id: str) -> List[Dict[str, Any]]:
|
|
591
|
+
"""
|
|
592
|
+
List WebSocket JSON payloads accumulated for ``job_id`` (order preserved).
|
|
593
|
+
|
|
594
|
+
Returns an empty list if the store is disabled or ``job_id`` is unknown.
|
|
595
|
+
|
|
596
|
+
When ``JsonRpcClient.get_job_messages`` exists and there is no running
|
|
597
|
+
event loop, returns the adapter's bounded history. With an active event
|
|
598
|
+
loop and adapter-side tracking, WebSocket frames are not copied into the
|
|
599
|
+
local store — use :meth:`alist_embed_file_status_messages` (or await
|
|
600
|
+
:meth:`aget_embed_file_status_snapshot`) to read from the adapter registry.
|
|
601
|
+
Otherwise returns the in-memory store filled during ``embed_file``.
|
|
602
|
+
"""
|
|
603
|
+
if self._embed_file_status_store is None:
|
|
604
|
+
return []
|
|
605
|
+
if not _uses_adapter_job_message_registry(self.client):
|
|
606
|
+
return self._embed_file_status_store.messages(job_id)
|
|
607
|
+
try:
|
|
608
|
+
asyncio.get_running_loop()
|
|
609
|
+
except RuntimeError:
|
|
610
|
+
return list(asyncio.run(self.client.get_job_messages(job_id)))
|
|
611
|
+
return self._embed_file_status_store.messages(job_id)
|
|
612
|
+
|
|
613
|
+
async def alist_embed_file_status_messages(
|
|
614
|
+
self, job_id: str
|
|
615
|
+
) -> List[Dict[str, Any]]:
|
|
616
|
+
"""
|
|
617
|
+
Async view of WS/job messages for ``job_id`` (adapter registry when available).
|
|
618
|
+
|
|
619
|
+
Empty list if the embed_file status store is disabled.
|
|
620
|
+
"""
|
|
621
|
+
if self._embed_file_status_store is None:
|
|
622
|
+
return []
|
|
623
|
+
if _uses_adapter_job_message_registry(self.client):
|
|
624
|
+
return list(await self.client.get_job_messages(job_id))
|
|
625
|
+
return self._embed_file_status_store.messages(job_id)
|
|
626
|
+
|
|
627
|
+
def clear_embed_file_status(self, job_id: str) -> None:
|
|
628
|
+
"""Remove stored submit/messages for ``job_id``. No-op if store disabled."""
|
|
629
|
+
if self._embed_file_status_store is None:
|
|
630
|
+
return
|
|
631
|
+
self._embed_file_status_store.clear(job_id)
|
|
632
|
+
if not _client_has_async_method(self.client, "clear_job_tracking"):
|
|
633
|
+
return
|
|
634
|
+
try:
|
|
635
|
+
loop = asyncio.get_running_loop()
|
|
636
|
+
except RuntimeError:
|
|
637
|
+
asyncio.run(self.client.clear_job_tracking(job_id))
|
|
638
|
+
else:
|
|
639
|
+
loop.create_task(self.client.clear_job_tracking(job_id))
|
|
640
|
+
|
|
641
|
+
def list_embed_file_status_job_ids(self) -> List[str]:
|
|
642
|
+
"""
|
|
643
|
+
Return ``job_id`` values currently present in the embed_file status store.
|
|
644
|
+
|
|
645
|
+
Empty list if the store is disabled or nothing was recorded yet.
|
|
646
|
+
"""
|
|
647
|
+
if self._embed_file_status_store is None:
|
|
648
|
+
return []
|
|
649
|
+
return self._embed_file_status_store.job_ids()
|
|
650
|
+
|
|
651
|
+
async def health(self) -> Dict[str, Any]:
|
|
652
|
+
"""
|
|
653
|
+
Check server health.
|
|
654
|
+
|
|
655
|
+
Returns:
|
|
656
|
+
Health status information
|
|
657
|
+
"""
|
|
658
|
+
return await self.client.health()
|
|
659
|
+
|
|
660
|
+
async def embed(
|
|
661
|
+
self,
|
|
662
|
+
texts: List[str],
|
|
663
|
+
model: Optional[str] = None,
|
|
664
|
+
dimension: Optional[int] = None,
|
|
665
|
+
job_id: Optional[str] = None,
|
|
666
|
+
error_policy: Optional[str] = None,
|
|
667
|
+
device: Optional[str] = None,
|
|
668
|
+
aggregation_method: Optional[str] = None,
|
|
669
|
+
parts: Optional[List[Optional[List[str]]]] = None,
|
|
670
|
+
wait: bool = False,
|
|
671
|
+
wait_timeout: int = 60,
|
|
672
|
+
poll_interval: float = 1.0,
|
|
673
|
+
use_push: bool = True,
|
|
674
|
+
provider: Optional[str] = None,
|
|
675
|
+
) -> Dict[str, Any]:
|
|
676
|
+
"""
|
|
677
|
+
Generate embeddings for texts (via queue).
|
|
678
|
+
|
|
679
|
+
Parameters match server embed command schema. When wait=True, completion
|
|
680
|
+
is always received via WebSocket (/ws) push (use_push is ignored; server
|
|
681
|
+
always uses WS for result delivery).
|
|
682
|
+
|
|
683
|
+
Args:
|
|
684
|
+
texts: List of texts to create embeddings for
|
|
685
|
+
model: Optional model name
|
|
686
|
+
dimension: Optional expected vector dimension
|
|
687
|
+
job_id: Optional custom job ID (server may ignore and generate own)
|
|
688
|
+
error_policy: 'fail_fast' or 'continue'; default from server
|
|
689
|
+
device: 'auto', 'cpu', 'cuda', 'cuda:0'; default from server
|
|
690
|
+
aggregation_method: How to combine part vectors when a text is
|
|
691
|
+
longer than the model can read. For the local model, when
|
|
692
|
+
neither the call nor the server config sets it, the default is
|
|
693
|
+
``weighted_mean_by_tokens_normalized`` (``method_source``
|
|
694
|
+
``constant``); a config value overrides that default (an
|
|
695
|
+
explicit null in the config disables aggregation, so a text past
|
|
696
|
+
the limit is refused rather than truncated) and a call value
|
|
697
|
+
overrides both. For a nonlocal provider a text past its input limit is
|
|
698
|
+
split and aggregated server-side, with the default
|
|
699
|
+
``weighted_mean_by_tokens_normalized`` when this is omitted.
|
|
700
|
+
parts: One entry per text: a list of parts to combine, or ``None``
|
|
701
|
+
to let the server derive the parts from the text itself.
|
|
702
|
+
wait: If True, wait for job completion and return results (via /ws).
|
|
703
|
+
wait_timeout: Max time to wait (seconds); 0 = no timeout.
|
|
704
|
+
poll_interval: Unused (kept for API compatibility).
|
|
705
|
+
use_push: Unused; WS is always used when wait=True.
|
|
706
|
+
provider: Optional provider id of the full model id (provider +
|
|
707
|
+
model). ``"local"`` is the server's own model; any other id is
|
|
708
|
+
a live macsc vectorization provider (for example ``"openai"``).
|
|
709
|
+
Sent only when not ``None`` (never injected by the client).
|
|
710
|
+
Resolution on the server (default first): provider given ->
|
|
711
|
+
must be known (else UNKNOWN_PROVIDER); model omitted -> the
|
|
712
|
+
default model, which must belong to that provider (else
|
|
713
|
+
UNKNOWN_MODEL). Provider omitted and model omitted -> the
|
|
714
|
+
default pair. Provider omitted and model given -> the default
|
|
715
|
+
provider is tried first (if it lists the model it is used,
|
|
716
|
+
macsc is not consulted for a local default); otherwise the
|
|
717
|
+
other providers are searched: exactly one offering the model
|
|
718
|
+
is used, several -> PROVIDER_REQUIRED, none -> UNKNOWN_MODEL.
|
|
719
|
+
``dimension`` omitted -> the entry's default dimension; given
|
|
720
|
+
-> must satisfy its supported dimensions (else
|
|
721
|
+
UNSUPPORTED_DIMENSION). Errors arrive with ``error.data.code``
|
|
722
|
+
and ``error.data.kind`` unchanged: -32602 UNKNOWN_PROVIDER,
|
|
723
|
+
UNKNOWN_MODEL, UNSUPPORTED_DIMENSION, PROVIDER_REQUIRED;
|
|
724
|
+
-32003 EXTERNAL_UNAVAILABLE, EXTERNAL_UPSTREAM_FAILURE,
|
|
725
|
+
EXTERNAL_BAD_RESPONSE, EXTERNAL_INVALID_REQUEST; -32001
|
|
726
|
+
MODEL_LOADING (retry); -32004 OWN_MODEL_DISABLED; -32005
|
|
727
|
+
GPU_DISABLED; -32603 MODEL_LOAD_FAILED, NO_GPU_AVAILABLE,
|
|
728
|
+
-32603 TOKEN_MEASUREMENT_FAILED (retry; the input length could
|
|
729
|
+
not be measured, nothing was embedded).
|
|
730
|
+
``aggregation_method`` and ``parts`` are accepted for any
|
|
731
|
+
effective provider. A text longer than a nonlocal provider's input
|
|
732
|
+
limit is split and aggregated server-side: with no
|
|
733
|
+
``aggregation_method`` the default is
|
|
734
|
+
``weighted_mean_by_tokens_normalized`` (reported as
|
|
735
|
+
``method_source`` ``provider_default``); it is never refused for
|
|
736
|
+
lack of a method. ``device`` is accepted only with the effective
|
|
737
|
+
provider ``"local"``. The client does not validate these
|
|
738
|
+
parameters; it forwards them as given.
|
|
739
|
+
|
|
740
|
+
Returns:
|
|
741
|
+
If wait=False: ``{"job_id", "status", "message"}``; when the server hit
|
|
742
|
+
its synchronous execution cap and fell back to the queue, responses may
|
|
743
|
+
include ``queued_after_timeout: true`` (mcp-proxy-adapter >= 8.10).
|
|
744
|
+
If wait=True: {"results": [...], "model", "dimension"} (and optional
|
|
745
|
+
"device", "provider", "backend" and timing fields such as ``timing_total_seconds``,
|
|
746
|
+
``timing_encode_seconds``, ``timing_tokenization_seconds``,
|
|
747
|
+
``timing_median_per_text_encode_seconds``), matching server embed
|
|
748
|
+
result data.
|
|
749
|
+
|
|
750
|
+
Raises:
|
|
751
|
+
EmbedClientError: If embedding request or job fails.
|
|
752
|
+
TimeoutError: If wait_timeout > 0 and job does not complete in time.
|
|
753
|
+
|
|
754
|
+
Example:
|
|
755
|
+
# Submit job and get job_id immediately
|
|
756
|
+
result = await client.embed(["Hello, world!"], wait=False)
|
|
757
|
+
job_id = result["job_id"]
|
|
758
|
+
|
|
759
|
+
# Wait for completion with timeout
|
|
760
|
+
result = await client.embed(
|
|
761
|
+
["Hello, world!"],
|
|
762
|
+
wait=True,
|
|
763
|
+
wait_timeout=120 # Wait up to 2 minutes
|
|
764
|
+
)
|
|
765
|
+
|
|
766
|
+
# With error_policy and device (match server schema)
|
|
767
|
+
result = await client.embed(
|
|
768
|
+
["Hello!"], error_policy="continue", device="cpu", wait=True
|
|
769
|
+
)
|
|
770
|
+
"""
|
|
771
|
+
# Params match server embed schema (texts required; rest optional)
|
|
772
|
+
params: Dict[str, Any] = {"texts": texts}
|
|
773
|
+
if model is not None:
|
|
774
|
+
params["model"] = model
|
|
775
|
+
if dimension is not None:
|
|
776
|
+
params["dimension"] = dimension
|
|
777
|
+
if job_id is not None:
|
|
778
|
+
params["job_id"] = job_id
|
|
779
|
+
if error_policy is not None:
|
|
780
|
+
params["error_policy"] = error_policy
|
|
781
|
+
else:
|
|
782
|
+
params["error_policy"] = EMBED_COMMAND_SCHEMA_DEFAULT_ERROR_POLICY
|
|
783
|
+
if device is not None:
|
|
784
|
+
params["device"] = device
|
|
785
|
+
if aggregation_method is not None:
|
|
786
|
+
params["aggregation_method"] = aggregation_method
|
|
787
|
+
if parts is not None:
|
|
788
|
+
params["parts"] = parts
|
|
789
|
+
if provider is not None:
|
|
790
|
+
params["provider"] = provider
|
|
791
|
+
|
|
792
|
+
# When wait=True use execute_async("embed_execute") for in-process execution
|
|
793
|
+
# (no queue, no worker spawn) to avoid large per-request delay.
|
|
794
|
+
if wait:
|
|
795
|
+
return await self._embed_via_execute_async(params, wait_timeout)
|
|
796
|
+
|
|
797
|
+
response = await self.client.jsonrpc_call("embed", params)
|
|
798
|
+
|
|
799
|
+
# Extract result
|
|
800
|
+
if "error" in response:
|
|
801
|
+
raise EmbedClientError(f"Embedding failed: {response['error']}")
|
|
802
|
+
|
|
803
|
+
result = response.get("result", {})
|
|
804
|
+
if not result.get("success"):
|
|
805
|
+
error = result.get("error", {})
|
|
806
|
+
raise EmbedClientError(
|
|
807
|
+
f"Embedding failed: {error.get('message', 'Unknown error')}"
|
|
808
|
+
)
|
|
809
|
+
|
|
810
|
+
# Extract data - server may return full result (sync) or job_id (async queue)
|
|
811
|
+
data = result.get("data", {})
|
|
812
|
+
out_job_id = data.get("job_id") or result.get("job_id")
|
|
813
|
+
|
|
814
|
+
# Sync response: full result in one response (no job_id)
|
|
815
|
+
if not out_job_id and isinstance(data.get("results"), list):
|
|
816
|
+
return data
|
|
817
|
+
|
|
818
|
+
if not out_job_id:
|
|
819
|
+
raise EmbedClientError("No job_id or results returned from server")
|
|
820
|
+
|
|
821
|
+
qat = data.get("queued_after_timeout")
|
|
822
|
+
if qat is None:
|
|
823
|
+
qat = result.get("queued_after_timeout")
|
|
824
|
+
out: Dict[str, Any] = {
|
|
825
|
+
"job_id": out_job_id,
|
|
826
|
+
"status": data.get("status") or result.get("status", "pending"),
|
|
827
|
+
"message": data.get("message")
|
|
828
|
+
or result.get("message", "Job queued successfully"),
|
|
829
|
+
}
|
|
830
|
+
if qat is not None:
|
|
831
|
+
out["queued_after_timeout"] = bool(qat)
|
|
832
|
+
return out
|
|
833
|
+
|
|
834
|
+
async def wait_for_job(
|
|
835
|
+
self,
|
|
836
|
+
job_id: str,
|
|
837
|
+
timeout: Union[int, float] = 60,
|
|
838
|
+
poll_interval: float = 1.0,
|
|
839
|
+
) -> Dict[str, Any]:
|
|
840
|
+
"""
|
|
841
|
+
Wait for job completion.
|
|
842
|
+
|
|
843
|
+
When the underlying ``JsonRpcClient`` provides ``wait_for_job``
|
|
844
|
+
(mcp-proxy-adapter ``JobTrackingMixin``, >= 6.10), completion is waited
|
|
845
|
+
via WebSocket and the result is normalized to the same shape as the poll
|
|
846
|
+
path. Otherwise polls :meth:`job_status` every *poll_interval* seconds.
|
|
847
|
+
|
|
848
|
+
Args:
|
|
849
|
+
job_id: Job ID to wait for
|
|
850
|
+
timeout: Max wait (seconds); 0 = no timeout (WS path: no receive
|
|
851
|
+
deadline; poll path: unbounded loop)
|
|
852
|
+
poll_interval: Seconds between status checks (HTTP poll fallback only)
|
|
853
|
+
|
|
854
|
+
Returns:
|
|
855
|
+
{"results": [...], "model": str, "dimension": int} (and optional "device")
|
|
856
|
+
|
|
857
|
+
Raises:
|
|
858
|
+
EmbedClientError: If job not found or job failed. A job that runs
|
|
859
|
+
longer than the server's 180 s job limit is terminated by the
|
|
860
|
+
server and fails with "Worker request timed out waiting for a
|
|
861
|
+
response from the worker process" (-32603, WORKER_RESPONSE_TIMEOUT);
|
|
862
|
+
the server stays usable for the next request.
|
|
863
|
+
TimeoutError: If timeout > 0 and job does not complete in time
|
|
864
|
+
"""
|
|
865
|
+
if _client_has_async_method(self.client, "wait_for_job"):
|
|
866
|
+
t_arg = float(timeout) if timeout > 0 else 0.0
|
|
867
|
+
started = asyncio.get_running_loop().time()
|
|
868
|
+
try:
|
|
869
|
+
raw = await self.client.wait_for_job(job_id, timeout=t_arg)
|
|
870
|
+
except asyncio.TimeoutError as e:
|
|
871
|
+
if t_arg > 0:
|
|
872
|
+
msg = f"Job {job_id} did not complete within {t_arg} seconds"
|
|
873
|
+
else:
|
|
874
|
+
msg = f"Job {job_id} WebSocket wait timed out"
|
|
875
|
+
raise TimeoutError(msg) from e
|
|
876
|
+
except Exception as e: # noqa: BLE001 - classified below
|
|
877
|
+
if not _is_ws_closed_before_terminal(e):
|
|
878
|
+
if isinstance(e, RuntimeError):
|
|
879
|
+
raise EmbedClientError(f"WebSocket error: {e}") from e
|
|
880
|
+
raise
|
|
881
|
+
return await self._poll_terminal_after_ws_close(
|
|
882
|
+
job_id, t_arg, started, e, poll_interval=poll_interval
|
|
883
|
+
)
|
|
884
|
+
return _normalize_terminal_job_payload(raw)
|
|
885
|
+
return await wait_for_job_poll(
|
|
886
|
+
self.job_status, job_id, timeout=timeout, poll_interval=poll_interval
|
|
887
|
+
)
|
|
888
|
+
|
|
889
|
+
async def _poll_terminal_after_ws_close(
|
|
890
|
+
self,
|
|
891
|
+
job_id: str,
|
|
892
|
+
timeout: float,
|
|
893
|
+
started: float,
|
|
894
|
+
cause: BaseException,
|
|
895
|
+
poll_interval: float = 1.0,
|
|
896
|
+
) -> Dict[str, Any]:
|
|
897
|
+
"""Finish a wait whose /ws stream closed before the terminal event.
|
|
898
|
+
|
|
899
|
+
Polls :meth:`job_status` (``queue_get_job_status``) for the rest of the
|
|
900
|
+
caller's *timeout* and returns exactly what the WebSocket path would
|
|
901
|
+
have returned: the terminal status is wrapped as the equivalent
|
|
902
|
+
``job_completed`` / ``job_failed`` event and passed through the same
|
|
903
|
+
normalizer, so a failed command inside a completed envelope is still
|
|
904
|
+
reported as a failure. ``timeout`` 0 means no deadline, as on the WS path.
|
|
905
|
+
The terminal events the normalizer accepts are ``job_completed``,
|
|
906
|
+
``job_completed_success``, ``delivery_completed``, ``job_failed``,
|
|
907
|
+
``delivery_failed``, ``delivery_cancelled`` and ``delivery_stopped``.
|
|
908
|
+
"""
|
|
909
|
+
loop = asyncio.get_running_loop()
|
|
910
|
+
deadline = None if timeout <= 0 else started + timeout
|
|
911
|
+
logger.warning(
|
|
912
|
+
"WebSocket closed before a terminal event for job %s (%s); "
|
|
913
|
+
"polling queue_get_job_status for its terminal state",
|
|
914
|
+
job_id,
|
|
915
|
+
cause,
|
|
916
|
+
)
|
|
917
|
+
interval = poll_interval if poll_interval and poll_interval > 0 else 1.0
|
|
918
|
+
while True:
|
|
919
|
+
status = await self.job_status(job_id)
|
|
920
|
+
if isinstance(status, dict):
|
|
921
|
+
if status.get("exists") is False:
|
|
922
|
+
raise EmbedClientError(f"Job {job_id} not found") from cause
|
|
923
|
+
state = str(status.get("status") or "").lower()
|
|
924
|
+
if state in _POLL_COMPLETED_STATES:
|
|
925
|
+
return _normalize_terminal_job_payload(
|
|
926
|
+
{
|
|
927
|
+
"event": "job_completed",
|
|
928
|
+
"job_id": job_id,
|
|
929
|
+
"result": status.get("result") or {},
|
|
930
|
+
}
|
|
931
|
+
)
|
|
932
|
+
if state in _POLL_FAILED_STATES:
|
|
933
|
+
failed_event: Dict[str, Any] = {
|
|
934
|
+
"event": "job_failed",
|
|
935
|
+
"job_id": job_id,
|
|
936
|
+
}
|
|
937
|
+
# Omit "error" when the status has none: the normalizer then
|
|
938
|
+
# words the message "Unknown error" and exc.error stays None.
|
|
939
|
+
failed_error = (
|
|
940
|
+
status.get("error")
|
|
941
|
+
or status.get("description")
|
|
942
|
+
or status.get("message")
|
|
943
|
+
)
|
|
944
|
+
if failed_error:
|
|
945
|
+
failed_event["error"] = failed_error
|
|
946
|
+
return _normalize_terminal_job_payload(failed_event)
|
|
947
|
+
if state in _POLL_STOPPED_STATES:
|
|
948
|
+
raise EmbedClientError(f"Job {job_id} was stopped") from cause
|
|
949
|
+
if deadline is not None and loop.time() >= deadline:
|
|
950
|
+
raise TimeoutError(
|
|
951
|
+
f"Job {job_id} did not complete within {timeout} seconds"
|
|
952
|
+
) from cause
|
|
953
|
+
await asyncio.sleep(interval)
|
|
954
|
+
|
|
955
|
+
async def _embed_via_execute_async(
|
|
956
|
+
self, params: Dict[str, Any], timeout: int = 60
|
|
957
|
+
) -> Dict[str, Any]:
|
|
958
|
+
"""
|
|
959
|
+
Run embed in-process via execute_async (no queue, no worker spawn).
|
|
960
|
+
|
|
961
|
+
Uses adapter execute_async RPC; server runs embed_execute in main process
|
|
962
|
+
and pushes result via WebSocket. Returns embed result shape; raises on
|
|
963
|
+
failure or timeout.
|
|
964
|
+
"""
|
|
965
|
+
timeout_sec = float(timeout) if timeout > 0 else 60.0
|
|
966
|
+
loop = asyncio.get_running_loop()
|
|
967
|
+
fut: asyncio.Future[Dict[str, Any]] = loop.create_future()
|
|
968
|
+
|
|
969
|
+
def on_result(event: Dict[str, Any]) -> None:
|
|
970
|
+
if not fut.done():
|
|
971
|
+
fut.set_result(event)
|
|
972
|
+
|
|
973
|
+
await self.client.execute_async(
|
|
974
|
+
"embed_execute",
|
|
975
|
+
params,
|
|
976
|
+
on_result=on_result,
|
|
977
|
+
timeout=timeout_sec,
|
|
978
|
+
)
|
|
979
|
+
try:
|
|
980
|
+
event = await asyncio.wait_for(fut, timeout=timeout_sec + 5.0)
|
|
981
|
+
except asyncio.TimeoutError as e:
|
|
982
|
+
raise TimeoutError(
|
|
983
|
+
f"Embed (execute_async) did not complete within {timeout_sec}s"
|
|
984
|
+
) from e
|
|
985
|
+
|
|
986
|
+
ev = event.get("event") or event.get("type")
|
|
987
|
+
if ev in _FAILED_TERMINAL_EVENTS:
|
|
988
|
+
err = event.get("error", "Unknown error")
|
|
989
|
+
raise EmbedClientError(f"Embed failed: {err}", error=event.get("error"))
|
|
990
|
+
if ev in _CANCELLED_TERMINAL_EVENTS:
|
|
991
|
+
err = event.get("error", "Unknown error")
|
|
992
|
+
raise EmbedClientError(f"Embed cancelled: {err}", error=event.get("error"))
|
|
993
|
+
# Accept job_completed, job_completed_success, and delivery_completed
|
|
994
|
+
# (delivery_completed is used in newer adapter versions)
|
|
995
|
+
if ev not in ("job_completed", "job_completed_success", "delivery_completed"):
|
|
996
|
+
raise EmbedClientError(
|
|
997
|
+
f"Unexpected WebSocket event: {ev} (keys: {list(event.keys())})"
|
|
998
|
+
)
|
|
999
|
+
|
|
1000
|
+
return _normalize_terminal_job_payload(event)
|
|
1001
|
+
|
|
1002
|
+
async def _wait_for_job_via_adapter_ws(
|
|
1003
|
+
self,
|
|
1004
|
+
job_id: str,
|
|
1005
|
+
timeout: int = 60,
|
|
1006
|
+
*,
|
|
1007
|
+
on_ws_event: Optional[Callable[[Dict[str, Any]], Any]] = None,
|
|
1008
|
+
) -> Dict[str, Any]:
|
|
1009
|
+
"""
|
|
1010
|
+
Wait for job completion via adapter WebSocket (/ws queue_job subscribe).
|
|
1011
|
+
|
|
1012
|
+
Prefers ``JsonRpcClient.wait_for_job`` when present so frames are recorded
|
|
1013
|
+
in the adapter job registry; optional ``on_ws_event`` is chained after
|
|
1014
|
+
adapter recording. Falls back to ``wait_for_job_via_websocket`` with
|
|
1015
|
+
``on_inbound_message`` when ``wait_for_job`` is not available.
|
|
1016
|
+
"""
|
|
1017
|
+
if _client_has_async_method(self.client, "wait_for_job"):
|
|
1018
|
+
t_arg = float(timeout) if timeout > 0 else 0.0
|
|
1019
|
+
started = asyncio.get_running_loop().time()
|
|
1020
|
+
|
|
1021
|
+
async def _on_event(msg: Dict[str, Any]) -> None:
|
|
1022
|
+
if on_ws_event is not None:
|
|
1023
|
+
await _maybe_await(on_ws_event(msg))
|
|
1024
|
+
|
|
1025
|
+
try:
|
|
1026
|
+
payload = await self.client.wait_for_job(
|
|
1027
|
+
job_id,
|
|
1028
|
+
timeout=t_arg,
|
|
1029
|
+
on_event=_on_event,
|
|
1030
|
+
)
|
|
1031
|
+
except asyncio.TimeoutError as e:
|
|
1032
|
+
if t_arg > 0:
|
|
1033
|
+
msg = f"Job {job_id} did not complete within {t_arg} seconds"
|
|
1034
|
+
else:
|
|
1035
|
+
msg = f"Job {job_id} WebSocket wait timed out"
|
|
1036
|
+
raise TimeoutError(msg) from e
|
|
1037
|
+
except Exception as e: # noqa: BLE001 - classified below
|
|
1038
|
+
if not _is_ws_closed_before_terminal(e):
|
|
1039
|
+
if isinstance(e, RuntimeError):
|
|
1040
|
+
raise EmbedClientError(f"WebSocket error: {e}") from e
|
|
1041
|
+
raise
|
|
1042
|
+
return await self._poll_terminal_after_ws_close(
|
|
1043
|
+
job_id, t_arg, started, e
|
|
1044
|
+
)
|
|
1045
|
+
return _normalize_terminal_job_payload(payload)
|
|
1046
|
+
|
|
1047
|
+
timeout_sec = float(timeout) if timeout > 0 else 60.0
|
|
1048
|
+
try:
|
|
1049
|
+
payload = await self.client.wait_for_job_via_websocket(
|
|
1050
|
+
job_id,
|
|
1051
|
+
timeout=timeout_sec,
|
|
1052
|
+
on_inbound_message=on_ws_event,
|
|
1053
|
+
)
|
|
1054
|
+
except asyncio.TimeoutError as e:
|
|
1055
|
+
raise TimeoutError(
|
|
1056
|
+
f"Job {job_id} did not complete within {timeout_sec} seconds"
|
|
1057
|
+
) from e
|
|
1058
|
+
except RuntimeError as e:
|
|
1059
|
+
raise EmbedClientError(f"WebSocket error: {e}") from e
|
|
1060
|
+
return _normalize_terminal_job_payload(payload)
|
|
1061
|
+
|
|
1062
|
+
def open_bidirectional_ws_channel(
|
|
1063
|
+
self,
|
|
1064
|
+
receive_timeout: float = 60.0,
|
|
1065
|
+
heartbeat: float = 30.0,
|
|
1066
|
+
) -> Any:
|
|
1067
|
+
"""
|
|
1068
|
+
Open adapter's bidirectional WebSocket channel to /ws (send + receive).
|
|
1069
|
+
|
|
1070
|
+
Returns the same context manager as the adapter's client: use
|
|
1071
|
+
send_json() to send (e.g. subscribe/unsubscribe), receive_iter() to
|
|
1072
|
+
consume server events (job_completed, job_failed, etc.).
|
|
1073
|
+
|
|
1074
|
+
Returns:
|
|
1075
|
+
Adapter BidirectionalWsChannel (use with async with).
|
|
1076
|
+
"""
|
|
1077
|
+
from mcp_proxy_adapter.client.jsonrpc_client import (
|
|
1078
|
+
open_bidirectional_ws_channel,
|
|
1079
|
+
)
|
|
1080
|
+
|
|
1081
|
+
return open_bidirectional_ws_channel(
|
|
1082
|
+
self.client,
|
|
1083
|
+
receive_timeout=receive_timeout,
|
|
1084
|
+
heartbeat=heartbeat,
|
|
1085
|
+
)
|
|
1086
|
+
|
|
1087
|
+
async def job_status(self, job_id: str) -> Dict[str, Any]:
|
|
1088
|
+
"""
|
|
1089
|
+
Get job status (shape matches server embed_job_status result.data).
|
|
1090
|
+
|
|
1091
|
+
Tries queue_get_job_status first, then embed_job_status. Returns dict
|
|
1092
|
+
with exists, status, done, and when completed: result (with success/data).
|
|
1093
|
+
|
|
1094
|
+
Args:
|
|
1095
|
+
job_id: Job ID to check
|
|
1096
|
+
|
|
1097
|
+
Returns:
|
|
1098
|
+
Server job status data (exists, status, done, result?).
|
|
1099
|
+
|
|
1100
|
+
Raises:
|
|
1101
|
+
EmbedClientError: If status request fails
|
|
1102
|
+
"""
|
|
1103
|
+
# Try built-in queue command first (for commands with use_queue=True)
|
|
1104
|
+
try:
|
|
1105
|
+
response = await self.client.jsonrpc_call(
|
|
1106
|
+
"queue_get_job_status", {"job_id": job_id}
|
|
1107
|
+
)
|
|
1108
|
+
if "error" not in response:
|
|
1109
|
+
result = response.get("result", {})
|
|
1110
|
+
if result.get("success"):
|
|
1111
|
+
return result.get("data", {})
|
|
1112
|
+
except Exception as e:
|
|
1113
|
+
logger.debug(
|
|
1114
|
+
"queue_get_job_status not available (%s), using embed_job_status",
|
|
1115
|
+
e,
|
|
1116
|
+
)
|
|
1117
|
+
|
|
1118
|
+
# Use embed_job_status when queue does not provide queue_get_job_status
|
|
1119
|
+
try:
|
|
1120
|
+
response = await self.client.jsonrpc_call(
|
|
1121
|
+
"embed_job_status", {"job_id": job_id}
|
|
1122
|
+
)
|
|
1123
|
+
|
|
1124
|
+
if "error" in response:
|
|
1125
|
+
raise EmbedClientError(
|
|
1126
|
+
f"Status check failed: {response['error']}",
|
|
1127
|
+
error=response["error"],
|
|
1128
|
+
)
|
|
1129
|
+
|
|
1130
|
+
result = response.get("result", {})
|
|
1131
|
+
if not result.get("success"):
|
|
1132
|
+
error = result.get("error", {})
|
|
1133
|
+
raise EmbedClientError(
|
|
1134
|
+
f"Status check failed: {error.get('message', 'Unknown error')}",
|
|
1135
|
+
error=result.get("error"),
|
|
1136
|
+
)
|
|
1137
|
+
|
|
1138
|
+
return result.get("data", {})
|
|
1139
|
+
except Exception as e:
|
|
1140
|
+
raise EmbedClientError(
|
|
1141
|
+
f"Status check failed: {e}",
|
|
1142
|
+
error=e.error if isinstance(e, EmbedClientError) else None,
|
|
1143
|
+
) from e
|
|
1144
|
+
|
|
1145
|
+
async def upload_file(
|
|
1146
|
+
self,
|
|
1147
|
+
source_path: str,
|
|
1148
|
+
*,
|
|
1149
|
+
filename: Optional[str] = None,
|
|
1150
|
+
) -> TransferReceipt:
|
|
1151
|
+
"""
|
|
1152
|
+
Stage a local file via the adapter HTTP upload (same transport as JSON-RPC).
|
|
1153
|
+
|
|
1154
|
+
Delegates to :meth:`JsonRpcClient.upload_file` on the underlying client.
|
|
1155
|
+
After upload, call :meth:`JsonRpcClient.get_upload_session` with the receipt's
|
|
1156
|
+
``transfer_id`` to obtain ``storage_path``, then pass that path as
|
|
1157
|
+
``input_file`` to :meth:`embed_file` (see :meth:`embed_file_from_local_json_file`).
|
|
1158
|
+
"""
|
|
1159
|
+
return await self.client.upload_file(source_path, filename=filename)
|
|
1160
|
+
|
|
1161
|
+
async def embed_file(
|
|
1162
|
+
self,
|
|
1163
|
+
input_file: str,
|
|
1164
|
+
*,
|
|
1165
|
+
file_format: str = "lines",
|
|
1166
|
+
model: Optional[str] = None,
|
|
1167
|
+
dimension: Optional[int] = None,
|
|
1168
|
+
device: Optional[str] = None,
|
|
1169
|
+
error_policy: Optional[str] = None,
|
|
1170
|
+
aggregation_method: Optional[str] = None,
|
|
1171
|
+
provider: Optional[str] = None,
|
|
1172
|
+
wait: bool = True,
|
|
1173
|
+
wait_timeout: int = 600,
|
|
1174
|
+
download_to: Optional[str] = None,
|
|
1175
|
+
skip_download: bool = False,
|
|
1176
|
+
local_source_path: Optional[str] = None,
|
|
1177
|
+
on_ws_event: Optional[Callable[[Dict[str, Any]], Any]] = None,
|
|
1178
|
+
on_job_id_known: Optional[Callable[[str], Any]] = None,
|
|
1179
|
+
on_download_start: Optional[Callable[[str, str], Any]] = None,
|
|
1180
|
+
on_download_done: Optional[Callable[[str], Any]] = None,
|
|
1181
|
+
) -> Dict[str, Any]:
|
|
1182
|
+
"""
|
|
1183
|
+
Run ``embed_file`` on the server: read texts from ``input_file``, write JSON.
|
|
1184
|
+
|
|
1185
|
+
Matches server command ``embed_file`` (queue job). When ``wait=True``, waits for
|
|
1186
|
+
completion via WebSocket. If the result includes an outbound transfer id, the
|
|
1187
|
+
client downloads via the adapter transfer HTTP API
|
|
1188
|
+
(``JsonRpcClient.download_file``) and sends ``transfer_ack`` so the server can
|
|
1189
|
+
remove the temp output file (when the session used ``delete_source_on_ack``).
|
|
1190
|
+
|
|
1191
|
+
``input_file`` must be a path on the **server** host. After :meth:`upload_file`,
|
|
1192
|
+
resolve the server-visible path with :meth:`JsonRpcClient.get_upload_session`
|
|
1193
|
+
(``storage_path``) and pass it here — :meth:`embed_file_from_local_json_file`
|
|
1194
|
+
does this for local ``json_array`` files.
|
|
1195
|
+
|
|
1196
|
+
Args:
|
|
1197
|
+
input_file: Absolute path to UTF-8 file on the **server** filesystem.
|
|
1198
|
+
file_format: ``lines`` or ``json_array`` (maps to server param ``format``).
|
|
1199
|
+
model, dimension, device, error_policy: Same semantics as :meth:`embed`.
|
|
1200
|
+
aggregation_method: Same semantics as :meth:`embed`: accepted for any
|
|
1201
|
+
provider; for a nonlocal provider an over-limit entry is split
|
|
1202
|
+
and aggregated server-side (default
|
|
1203
|
+
``weighted_mean_by_tokens_normalized``). For the local model it
|
|
1204
|
+
is resolved by priority call -> config -> constant when omitted,
|
|
1205
|
+
and the constant is ``weighted_mean_by_tokens_normalized``;
|
|
1206
|
+
applied per-item, one line/array entry at a time.
|
|
1207
|
+
provider: Same semantics as :meth:`embed` (``"local"`` or a live
|
|
1208
|
+
macsc provider id; sent only when not ``None``; resolution order
|
|
1209
|
+
and error codes as documented there).
|
|
1210
|
+
wait: If True, block until job completes (WebSocket).
|
|
1211
|
+
wait_timeout: Seconds for wait (0 treated as 600 for WS wait).
|
|
1212
|
+
download_to: If set, write the downloaded result JSON to this path. If
|
|
1213
|
+
``None`` and ``skip_download`` is False, uses a temporary local path.
|
|
1214
|
+
skip_download: If True, do not call ``download_file``; ``transfer_ack`` uses
|
|
1215
|
+
``checksum_value`` from the job completion payload (server metadata).
|
|
1216
|
+
local_source_path: If set, after the job completes verify this file's
|
|
1217
|
+
SHA-256 against ``input_checksum_value`` from the server (when
|
|
1218
|
+
present).
|
|
1219
|
+
on_ws_event: Optional callback for each WebSocket JSON payload while waiting
|
|
1220
|
+
for the queue job (including non-terminal events). May be sync or async.
|
|
1221
|
+
on_job_id_known: Called with ``job_id`` after the RPC returns and before the
|
|
1222
|
+
WebSocket wait (only when ``wait=True``).
|
|
1223
|
+
on_download_start: Called with ``(outbound_transfer_id, local_path)``
|
|
1224
|
+
before ``download_file`` (only when a transfer id is present and
|
|
1225
|
+
download runs).
|
|
1226
|
+
on_download_done: Called with ``local_path`` after a successful download
|
|
1227
|
+
and checksum, before ``transfer_ack``.
|
|
1228
|
+
|
|
1229
|
+
If the client was created with ``enable_embed_file_status_store=True``, the
|
|
1230
|
+
same submit payload and WebSocket frames used for this call are recorded
|
|
1231
|
+
under the job's ``job_id`` (see :meth:`get_embed_file_status_snapshot`).
|
|
1232
|
+
|
|
1233
|
+
Returns:
|
|
1234
|
+
With wait=True: merged job ``data`` plus ``local_output_path`` when a file
|
|
1235
|
+
was downloaded, and ``transfer_ack_result`` after a successful ACK.
|
|
1236
|
+
With wait=False: dict with at least ``job_id`` and ``status``.
|
|
1237
|
+
"""
|
|
1238
|
+
if not input_file or not str(input_file).strip():
|
|
1239
|
+
raise EmbedClientError("embed_file: input_file is required")
|
|
1240
|
+
|
|
1241
|
+
params: Dict[str, Any] = {
|
|
1242
|
+
"input_file": str(input_file).strip(),
|
|
1243
|
+
"format": file_format,
|
|
1244
|
+
}
|
|
1245
|
+
if model is not None:
|
|
1246
|
+
params["model"] = model
|
|
1247
|
+
if dimension is not None:
|
|
1248
|
+
params["dimension"] = dimension
|
|
1249
|
+
if device is not None:
|
|
1250
|
+
params["device"] = device
|
|
1251
|
+
if error_policy is not None:
|
|
1252
|
+
params["error_policy"] = error_policy
|
|
1253
|
+
else:
|
|
1254
|
+
params["error_policy"] = EMBED_COMMAND_SCHEMA_DEFAULT_ERROR_POLICY
|
|
1255
|
+
if aggregation_method is not None:
|
|
1256
|
+
params["aggregation_method"] = aggregation_method
|
|
1257
|
+
if provider is not None:
|
|
1258
|
+
params["provider"] = provider
|
|
1259
|
+
|
|
1260
|
+
response = await self.client.jsonrpc_call("embed_file", params)
|
|
1261
|
+
if "error" in response:
|
|
1262
|
+
raise EmbedClientError(f"embed_file failed: {response['error']}")
|
|
1263
|
+
|
|
1264
|
+
result = response.get("result", {})
|
|
1265
|
+
if not result.get("success"):
|
|
1266
|
+
err = result.get("error") or {}
|
|
1267
|
+
msg = (
|
|
1268
|
+
err.get("message", "Unknown error")
|
|
1269
|
+
if isinstance(err, dict)
|
|
1270
|
+
else str(err)
|
|
1271
|
+
)
|
|
1272
|
+
raise EmbedClientError(f"embed_file failed: {msg}")
|
|
1273
|
+
|
|
1274
|
+
nested = result.get("data")
|
|
1275
|
+
nested_dict = nested if isinstance(nested, dict) else {}
|
|
1276
|
+
job_id = nested_dict.get("job_id") or result.get("job_id")
|
|
1277
|
+
job_id_str = str(job_id).strip() if job_id else ""
|
|
1278
|
+
|
|
1279
|
+
if self._embed_file_status_store is not None and job_id_str:
|
|
1280
|
+
submit_snap: Dict[str, Any] = dict(nested_dict) if nested_dict else {}
|
|
1281
|
+
submit_snap.setdefault("job_id", job_id_str)
|
|
1282
|
+
self._embed_file_status_store.record_submit(job_id_str, submit_snap)
|
|
1283
|
+
|
|
1284
|
+
if not wait:
|
|
1285
|
+
out: Dict[str, Any] = dict(nested_dict)
|
|
1286
|
+
if job_id_str:
|
|
1287
|
+
out.setdefault("job_id", job_id_str)
|
|
1288
|
+
return out
|
|
1289
|
+
|
|
1290
|
+
if not job_id_str:
|
|
1291
|
+
if isinstance(nested_dict.get("output_path"), str):
|
|
1292
|
+
return dict(nested_dict)
|
|
1293
|
+
raise EmbedClientError(
|
|
1294
|
+
"embed_file: no job_id and no output_path in response"
|
|
1295
|
+
)
|
|
1296
|
+
|
|
1297
|
+
tw = int(wait_timeout) if wait_timeout > 0 else 600
|
|
1298
|
+
if on_job_id_known is not None:
|
|
1299
|
+
await _maybe_await(on_job_id_known(job_id_str))
|
|
1300
|
+
|
|
1301
|
+
ws_callback: Optional[Callable[[Dict[str, Any]], Any]] = on_ws_event
|
|
1302
|
+
if self._embed_file_status_store is not None and job_id_str:
|
|
1303
|
+
store = self._embed_file_status_store
|
|
1304
|
+
user_cb = on_ws_event
|
|
1305
|
+
|
|
1306
|
+
async def _embed_file_ws_with_store(ev: Dict[str, Any]) -> None:
|
|
1307
|
+
if not _uses_adapter_job_message_registry(self.client):
|
|
1308
|
+
store.record_message(job_id_str, ev)
|
|
1309
|
+
if user_cb is not None:
|
|
1310
|
+
await _maybe_await(user_cb(ev))
|
|
1311
|
+
|
|
1312
|
+
ws_callback = _embed_file_ws_with_store
|
|
1313
|
+
data = await self._wait_for_job_via_adapter_ws(
|
|
1314
|
+
job_id_str, tw, on_ws_event=ws_callback
|
|
1315
|
+
)
|
|
1316
|
+
|
|
1317
|
+
fail_msg = _embed_file_queue_result_error_message(data)
|
|
1318
|
+
if fail_msg:
|
|
1319
|
+
inner_result = data.get("result")
|
|
1320
|
+
raise EmbedClientError(
|
|
1321
|
+
f"embed_file failed: {fail_msg}",
|
|
1322
|
+
error=(
|
|
1323
|
+
inner_result.get("error")
|
|
1324
|
+
if isinstance(inner_result, dict)
|
|
1325
|
+
else None
|
|
1326
|
+
),
|
|
1327
|
+
)
|
|
1328
|
+
|
|
1329
|
+
if local_source_path:
|
|
1330
|
+
expected_in = _normalize_hex_checksum(data.get("input_checksum_value"))
|
|
1331
|
+
if not expected_in:
|
|
1332
|
+
raise EmbedClientError(
|
|
1333
|
+
"embed_file: local_source_path set but job result has no "
|
|
1334
|
+
"input_checksum_value for verification"
|
|
1335
|
+
)
|
|
1336
|
+
src_digest = _sha256_hex_file(local_source_path).lower()
|
|
1337
|
+
if src_digest != expected_in:
|
|
1338
|
+
raise EmbedClientError(
|
|
1339
|
+
"embed_file: local source file SHA256 does not match "
|
|
1340
|
+
f"server input_checksum_value (local={src_digest}, "
|
|
1341
|
+
f"server={expected_in})"
|
|
1342
|
+
)
|
|
1343
|
+
|
|
1344
|
+
tid = str(
|
|
1345
|
+
data.get("outbound_transfer_id") or data.get("transfer_id") or ""
|
|
1346
|
+
).strip()
|
|
1347
|
+
if not tid:
|
|
1348
|
+
return data
|
|
1349
|
+
|
|
1350
|
+
algo = str(data.get("checksum_algorithm") or "sha256").strip().lower()
|
|
1351
|
+
local_path: Optional[str] = None
|
|
1352
|
+
temp_path: Optional[str] = None
|
|
1353
|
+
checksum_value: str
|
|
1354
|
+
|
|
1355
|
+
if skip_download:
|
|
1356
|
+
cv = data.get("checksum_value")
|
|
1357
|
+
if not isinstance(cv, str) or not cv.strip():
|
|
1358
|
+
raise EmbedClientError(
|
|
1359
|
+
"embed_file: skip_download requires checksum_value in job result"
|
|
1360
|
+
)
|
|
1361
|
+
checksum_value = cv.strip().lower()
|
|
1362
|
+
else:
|
|
1363
|
+
if download_to:
|
|
1364
|
+
local_path = download_to
|
|
1365
|
+
else:
|
|
1366
|
+
fd, local_path = tempfile.mkstemp(prefix="embed_file_", suffix=".json")
|
|
1367
|
+
os.close(fd)
|
|
1368
|
+
temp_path = local_path
|
|
1369
|
+
try:
|
|
1370
|
+
if on_download_start is not None:
|
|
1371
|
+
await _maybe_await(on_download_start(tid, local_path))
|
|
1372
|
+
await self.client.download_file(tid, local_path)
|
|
1373
|
+
expected_out = _normalize_hex_checksum(data.get("checksum_value"))
|
|
1374
|
+
if not expected_out:
|
|
1375
|
+
raise EmbedClientError(
|
|
1376
|
+
"embed_file: job result missing checksum_value; "
|
|
1377
|
+
"cannot verify downloaded file integrity"
|
|
1378
|
+
)
|
|
1379
|
+
local_digest = _sha256_hex_file(local_path).lower()
|
|
1380
|
+
if local_digest != expected_out:
|
|
1381
|
+
raise EmbedClientError(
|
|
1382
|
+
"embed_file: downloaded file SHA256 does not match "
|
|
1383
|
+
f"server checksum_value (local={local_digest}, "
|
|
1384
|
+
f"expected={expected_out})"
|
|
1385
|
+
)
|
|
1386
|
+
checksum_value = local_digest
|
|
1387
|
+
if on_download_done is not None:
|
|
1388
|
+
await _maybe_await(on_download_done(local_path))
|
|
1389
|
+
except Exception:
|
|
1390
|
+
if temp_path:
|
|
1391
|
+
try:
|
|
1392
|
+
os.unlink(temp_path)
|
|
1393
|
+
except OSError:
|
|
1394
|
+
pass
|
|
1395
|
+
raise
|
|
1396
|
+
|
|
1397
|
+
ack_params: Dict[str, Any] = {
|
|
1398
|
+
"outbound_transfer_id": tid,
|
|
1399
|
+
"checksum_algorithm": algo,
|
|
1400
|
+
"checksum_value": checksum_value,
|
|
1401
|
+
}
|
|
1402
|
+
if job_id_str:
|
|
1403
|
+
ack_params["job_id"] = job_id_str
|
|
1404
|
+
|
|
1405
|
+
ack_resp = await self.client.jsonrpc_call("transfer_ack", ack_params)
|
|
1406
|
+
if "error" in ack_resp:
|
|
1407
|
+
err = ack_resp["error"]
|
|
1408
|
+
if temp_path:
|
|
1409
|
+
try:
|
|
1410
|
+
os.unlink(temp_path)
|
|
1411
|
+
except OSError:
|
|
1412
|
+
pass
|
|
1413
|
+
raise EmbedClientError(f"transfer_ack failed: {err}")
|
|
1414
|
+
|
|
1415
|
+
out = dict(data)
|
|
1416
|
+
if local_path is not None:
|
|
1417
|
+
out["local_output_path"] = local_path
|
|
1418
|
+
out["transfer_ack_result"] = ack_resp.get("result", ack_resp)
|
|
1419
|
+
return out
|
|
1420
|
+
|
|
1421
|
+
async def embed_file_from_local_json_file(
|
|
1422
|
+
self,
|
|
1423
|
+
local_path: str,
|
|
1424
|
+
*,
|
|
1425
|
+
max_block_bytes: Optional[int] = None,
|
|
1426
|
+
upload_filename: Optional[str] = None,
|
|
1427
|
+
**kwargs: Any,
|
|
1428
|
+
) -> Dict[str, Any]:
|
|
1429
|
+
"""
|
|
1430
|
+
Validate a local UTF-8 JSON ``json_array`` file (array of strings), upload
|
|
1431
|
+
via the adapter, resolve ``storage_path`` via ``get_upload_session``, then run
|
|
1432
|
+
:meth:`embed_file` with that path and ``format=json_array``.
|
|
1433
|
+
|
|
1434
|
+
See :mod:`embed_client.json_embed_file_preflight` for validation rules and
|
|
1435
|
+
``EMBED_CLIENT_MAX_JSON_BLOCK_BYTES``.
|
|
1436
|
+
|
|
1437
|
+
See also :mod:`embed_client.md_embed_file_prep` when strings came from
|
|
1438
|
+
``prepare-embed-file-json`` (paragraphs split on ``\\n\\n`` only).
|
|
1439
|
+
"""
|
|
1440
|
+
from embed_client.local_json_embed_file import (
|
|
1441
|
+
embed_file_from_local_json_file as _embed_local_json,
|
|
1442
|
+
)
|
|
1443
|
+
|
|
1444
|
+
return await _embed_local_json(
|
|
1445
|
+
self,
|
|
1446
|
+
local_path,
|
|
1447
|
+
max_block_bytes=max_block_bytes,
|
|
1448
|
+
upload_filename=upload_filename,
|
|
1449
|
+
**kwargs,
|
|
1450
|
+
)
|
|
1451
|
+
|
|
1452
|
+
async def jsonrpc(
|
|
1453
|
+
self, method: str, params: Optional[Dict[str, Any]] = None
|
|
1454
|
+
) -> Dict[str, Any]:
|
|
1455
|
+
"""
|
|
1456
|
+
Low-level JSON-RPC (same transport and auth as other methods).
|
|
1457
|
+
|
|
1458
|
+
Prefer :meth:`embed`, :meth:`models`, :meth:`health`, etc. when applicable.
|
|
1459
|
+
"""
|
|
1460
|
+
return await self.client.jsonrpc_call(method, params or {})
|
|
1461
|
+
|
|
1462
|
+
async def models(self) -> Dict[str, Any]:
|
|
1463
|
+
"""
|
|
1464
|
+
List available models (shape matches server models command result.data).
|
|
1465
|
+
|
|
1466
|
+
Returns:
|
|
1467
|
+
``{"use_own_model": "none"|"ram"|"gpu", "default": {provider,
|
|
1468
|
+
model}, "external": {configured, service}, "models": [...]}``; each
|
|
1469
|
+
entry has name, provider, dimension, supported_dimensions,
|
|
1470
|
+
max_tokens, backend and (local entries) residency. For the local
|
|
1471
|
+
entry ``max_tokens`` is the real window of the loaded (own) model,
|
|
1472
|
+
read from the model itself (its on-disk metadata before the load),
|
|
1473
|
+
and ``null`` when unknown; it is not a fixed table. Entries are
|
|
1474
|
+
unique by (provider, name). The local entry is listed only when
|
|
1475
|
+
``use_own_model`` is not ``"none"``; macsc entries are the live
|
|
1476
|
+
list of providers with vectorization. macsc configured but down
|
|
1477
|
+
raises -32003 EXTERNAL_UNAVAILABLE (no static fallback). The method
|
|
1478
|
+
takes no parameters.
|
|
1479
|
+
|
|
1480
|
+
Raises:
|
|
1481
|
+
EmbedClientError: If models request fails
|
|
1482
|
+
"""
|
|
1483
|
+
response = await self.client.jsonrpc_call("models", {})
|
|
1484
|
+
|
|
1485
|
+
if "error" in response:
|
|
1486
|
+
raise EmbedClientError(f"Models list failed: {response['error']}")
|
|
1487
|
+
|
|
1488
|
+
result = response.get("result", {})
|
|
1489
|
+
if not result.get("success"):
|
|
1490
|
+
error = result.get("error", {})
|
|
1491
|
+
raise EmbedClientError(
|
|
1492
|
+
f"Models list failed: {error.get('message', 'Unknown error')}"
|
|
1493
|
+
)
|
|
1494
|
+
|
|
1495
|
+
return result.get("data", {})
|
|
1496
|
+
|
|
1497
|
+
async def default_model(self) -> Dict[str, Any]:
|
|
1498
|
+
"""Default pair of the server (provider + model) and its dimension.
|
|
1499
|
+
|
|
1500
|
+
The server takes no parameters. The default pair is the configured
|
|
1501
|
+
default provider and model; the dimension is the configured default
|
|
1502
|
+
dimension, else the entry's default dimension (local: the loaded
|
|
1503
|
+
dimension). Calls that name no provider use this pair.
|
|
1504
|
+
|
|
1505
|
+
Returns:
|
|
1506
|
+
Exactly ``{model, dimension, provider}``.
|
|
1507
|
+
|
|
1508
|
+
Raises:
|
|
1509
|
+
EmbedAPIError: -32001 MODEL_LOADING (local default not loaded yet),
|
|
1510
|
+
-32003 EXTERNAL_* when macsc fails or does not list the
|
|
1511
|
+
configured default model.
|
|
1512
|
+
"""
|
|
1513
|
+
return await self._service_command("default_model", {})
|
|
1514
|
+
|
|
1515
|
+
async def _service_command(
|
|
1516
|
+
self, command: str, params: Dict[str, Any]
|
|
1517
|
+
) -> Dict[str, Any]:
|
|
1518
|
+
"""Run a synchronous service command and return its ``data``.
|
|
1519
|
+
|
|
1520
|
+
A refusal (error response or ``success: false``) is raised as
|
|
1521
|
+
:class:`EmbedAPIError` carrying the server's own error object, whose
|
|
1522
|
+
``code``/``data.code`` name the refusal (for example MODEL_LOADING).
|
|
1523
|
+
"""
|
|
1524
|
+
response = await self.client.jsonrpc_call(command, params)
|
|
1525
|
+
if isinstance(response, dict) and response.get("error") is not None:
|
|
1526
|
+
raise EmbedAPIError(response["error"])
|
|
1527
|
+
result = response.get("result") if isinstance(response, dict) else None
|
|
1528
|
+
if not isinstance(result, dict):
|
|
1529
|
+
raise EmbedAPIError(f"Command {command!r} returned {response!r}")
|
|
1530
|
+
if result.get("success") is False:
|
|
1531
|
+
raise EmbedAPIError(result.get("error") or result)
|
|
1532
|
+
data = result.get("data")
|
|
1533
|
+
return data if isinstance(data, dict) else result
|
|
1534
|
+
|
|
1535
|
+
async def model_offload(self, model: Optional[str] = None) -> Dict[str, Any]:
|
|
1536
|
+
"""Move the resident local model VRAM -> RAM and pin it there.
|
|
1537
|
+
|
|
1538
|
+
The hold lasts until ``model_load`` (or a model selection): neither the
|
|
1539
|
+
GPU watcher nor a local request moves it back; local requests are
|
|
1540
|
+
served on the CPU from the RAM image meanwhile. Idempotent.
|
|
1541
|
+
use_own_model ``gpu``: VRAM -> RAM, the GPU is released. ``ram``: an
|
|
1542
|
+
idempotent success (residency ``ram``, device ``cpu``, ``changed``
|
|
1543
|
+
false). ``none``: refused with -32004 OWN_MODEL_DISABLED. The command
|
|
1544
|
+
takes no ``provider``.
|
|
1545
|
+
|
|
1546
|
+
Args:
|
|
1547
|
+
model: Local model to offload; defaults to the resident model.
|
|
1548
|
+
|
|
1549
|
+
Returns:
|
|
1550
|
+
``{model, dimension, residency: "ram", previous_residency, device:
|
|
1551
|
+
"cpu", changed, vram_released_mb}``.
|
|
1552
|
+
|
|
1553
|
+
Raises:
|
|
1554
|
+
EmbedAPIError: -32602 invalid parameter or MODEL_NOT_RESIDENT,
|
|
1555
|
+
-32001 MODEL_LOADING, -32004 OWN_MODEL_DISABLED.
|
|
1556
|
+
"""
|
|
1557
|
+
params: Dict[str, Any] = {}
|
|
1558
|
+
if model is not None:
|
|
1559
|
+
params["model"] = model
|
|
1560
|
+
return await self._service_command("model_offload", params)
|
|
1561
|
+
|
|
1562
|
+
async def model_load(
|
|
1563
|
+
self,
|
|
1564
|
+
model: Optional[str] = None,
|
|
1565
|
+
device: Optional[str] = None,
|
|
1566
|
+
wait: Optional[bool] = None,
|
|
1567
|
+
) -> Dict[str, Any]:
|
|
1568
|
+
"""Move a local model RAM -> VRAM (or load it from disk).
|
|
1569
|
+
|
|
1570
|
+
use_own_model ``gpu``: the model goes onto the GPU. ``ram``: refused
|
|
1571
|
+
with -32005 GPU_DISABLED (any arguments). ``none``: refused with -32004
|
|
1572
|
+
OWN_MODEL_DISABLED. The command takes no ``provider``.
|
|
1573
|
+
|
|
1574
|
+
Args:
|
|
1575
|
+
model: Local model to load; defaults to the server's ``embedding.model_name``.
|
|
1576
|
+
device: ``"auto"`` (server default), ``"cuda"`` or ``"cuda:N"``.
|
|
1577
|
+
wait: ``False`` (server default) returns at once with residency
|
|
1578
|
+
``"loading"`` while the load continues in the background;
|
|
1579
|
+
``True`` blocks until the model is in VRAM or the load failed.
|
|
1580
|
+
|
|
1581
|
+
Returns:
|
|
1582
|
+
``{model, dimension, residency: "loading"|"vram",
|
|
1583
|
+
previous_residency, device, changed, async}``.
|
|
1584
|
+
|
|
1585
|
+
Raises:
|
|
1586
|
+
EmbedAPIError: -32602 invalid parameter, -32603 NO_GPU_AVAILABLE or
|
|
1587
|
+
MODEL_LOAD_FAILED, -32004 OWN_MODEL_DISABLED, -32005 GPU_DISABLED.
|
|
1588
|
+
"""
|
|
1589
|
+
params: Dict[str, Any] = {}
|
|
1590
|
+
if model is not None:
|
|
1591
|
+
params["model"] = model
|
|
1592
|
+
if device is not None:
|
|
1593
|
+
params["device"] = device
|
|
1594
|
+
if wait is not None:
|
|
1595
|
+
params["wait"] = wait
|
|
1596
|
+
return await self._service_command("model_load", params)
|
|
1597
|
+
|
|
1598
|
+
async def input_limit_add(
|
|
1599
|
+
self,
|
|
1600
|
+
provider: str,
|
|
1601
|
+
model: str,
|
|
1602
|
+
max_input_tokens: int,
|
|
1603
|
+
tokenizer: str,
|
|
1604
|
+
max_inputs_per_request: Optional[int] = None,
|
|
1605
|
+
max_total_tokens_per_request: Optional[int] = None,
|
|
1606
|
+
source: Optional[str] = None,
|
|
1607
|
+
) -> Dict[str, Any]:
|
|
1608
|
+
"""Add a manual input-limit row for (provider, model).
|
|
1609
|
+
|
|
1610
|
+
Cells not given are omitted from the request (server default: null for
|
|
1611
|
+
the three nullable cells). The client does not validate the values.
|
|
1612
|
+
|
|
1613
|
+
Args:
|
|
1614
|
+
provider: Provider id (``"local"`` or a macsc provider id).
|
|
1615
|
+
model: Model id under that provider.
|
|
1616
|
+
max_input_tokens: Positive integer limit; ``"model"`` is the only
|
|
1617
|
+
value the local row accepts, so a number on ``local`` is refused.
|
|
1618
|
+
tokenizer: Tokenizer name (enum ``cl100k_base``).
|
|
1619
|
+
max_inputs_per_request: Positive integer or ``None``.
|
|
1620
|
+
max_total_tokens_per_request: Positive integer or ``None``.
|
|
1621
|
+
source: Free-text provenance of the limit, or ``None``.
|
|
1622
|
+
|
|
1623
|
+
Returns:
|
|
1624
|
+
``{provider, model, row: {max_input_tokens, tokenizer,
|
|
1625
|
+
max_inputs_per_request, max_total_tokens_per_request, source,
|
|
1626
|
+
origin: "manual"}}``.
|
|
1627
|
+
|
|
1628
|
+
Raises:
|
|
1629
|
+
EmbedAPIError: -32602 INPUT_LIMIT_EXISTS (row already present),
|
|
1630
|
+
-32602 LOCAL_LIMIT_FROM_MODEL (numeric limit on ``local``),
|
|
1631
|
+
-32602 validation (unknown, missing, type or value; expected/got),
|
|
1632
|
+
-32603 INPUT_LIMIT_PERSIST_FAILED (table and file unchanged).
|
|
1633
|
+
"""
|
|
1634
|
+
params: Dict[str, Any] = {
|
|
1635
|
+
"provider": provider,
|
|
1636
|
+
"model": model,
|
|
1637
|
+
"max_input_tokens": max_input_tokens,
|
|
1638
|
+
"tokenizer": tokenizer,
|
|
1639
|
+
}
|
|
1640
|
+
if max_inputs_per_request is not None:
|
|
1641
|
+
params["max_inputs_per_request"] = max_inputs_per_request
|
|
1642
|
+
if max_total_tokens_per_request is not None:
|
|
1643
|
+
params["max_total_tokens_per_request"] = max_total_tokens_per_request
|
|
1644
|
+
if source is not None:
|
|
1645
|
+
params["source"] = source
|
|
1646
|
+
return await self._service_command("input_limit_add", params)
|
|
1647
|
+
|
|
1648
|
+
async def input_limit_delete(self, provider: str, model: str) -> Dict[str, Any]:
|
|
1649
|
+
"""Delete the input-limit row for (provider, model).
|
|
1650
|
+
|
|
1651
|
+
A row registered by a module is remembered as deleted (not re-added at
|
|
1652
|
+
startup); a manual row with no module row is removed.
|
|
1653
|
+
|
|
1654
|
+
Args:
|
|
1655
|
+
provider: Provider id.
|
|
1656
|
+
model: Model id under that provider.
|
|
1657
|
+
|
|
1658
|
+
Returns:
|
|
1659
|
+
``{provider, model, deleted: true, origin_was}``.
|
|
1660
|
+
|
|
1661
|
+
Raises:
|
|
1662
|
+
EmbedAPIError: -32602 INPUT_LIMIT_NOT_FOUND (no such row), -32602
|
|
1663
|
+
validation, -32603 INPUT_LIMIT_PERSIST_FAILED (unchanged).
|
|
1664
|
+
"""
|
|
1665
|
+
return await self._service_command(
|
|
1666
|
+
"input_limit_delete", {"provider": provider, "model": model}
|
|
1667
|
+
)
|
|
1668
|
+
|
|
1669
|
+
async def input_limit_update(
|
|
1670
|
+
self,
|
|
1671
|
+
provider: str,
|
|
1672
|
+
model: str,
|
|
1673
|
+
*,
|
|
1674
|
+
max_input_tokens: Any = _NOT_GIVEN,
|
|
1675
|
+
tokenizer: Any = _NOT_GIVEN,
|
|
1676
|
+
max_inputs_per_request: Any = _NOT_GIVEN,
|
|
1677
|
+
max_total_tokens_per_request: Any = _NOT_GIVEN,
|
|
1678
|
+
source: Any = _NOT_GIVEN,
|
|
1679
|
+
) -> Dict[str, Any]:
|
|
1680
|
+
"""Change only the named cells of the input-limit row for (provider, model).
|
|
1681
|
+
|
|
1682
|
+
Only the cells passed explicitly are sent; every other cell of the row
|
|
1683
|
+
stays unchanged. Passing ``None`` sends null, which the server accepts
|
|
1684
|
+
only for ``max_inputs_per_request``, ``max_total_tokens_per_request`` and
|
|
1685
|
+
``source``. Updating a module row makes it a manual row (manual override).
|
|
1686
|
+
Calling with no cells sends none, and the server refuses it with
|
|
1687
|
+
INPUT_LIMIT_EMPTY_UPDATE.
|
|
1688
|
+
|
|
1689
|
+
Args:
|
|
1690
|
+
provider: Provider id.
|
|
1691
|
+
model: Model id under that provider.
|
|
1692
|
+
max_input_tokens: Positive integer; refused on ``local``
|
|
1693
|
+
(LOCAL_LIMIT_FROM_MODEL).
|
|
1694
|
+
tokenizer: Tokenizer name (enum ``cl100k_base``).
|
|
1695
|
+
max_inputs_per_request: Positive integer or ``None``.
|
|
1696
|
+
max_total_tokens_per_request: Positive integer or ``None``.
|
|
1697
|
+
source: Free-text provenance or ``None``.
|
|
1698
|
+
|
|
1699
|
+
Returns:
|
|
1700
|
+
``{provider, model, row: {all cells}, changed: [names of the cells
|
|
1701
|
+
whose value changed], origin: "manual"}``.
|
|
1702
|
+
|
|
1703
|
+
Raises:
|
|
1704
|
+
EmbedAPIError: -32602 INPUT_LIMIT_NOT_FOUND (no such row),
|
|
1705
|
+
-32602 INPUT_LIMIT_EMPTY_UPDATE (no cell named),
|
|
1706
|
+
-32602 LOCAL_LIMIT_FROM_MODEL, -32602 validation (unknown cell,
|
|
1707
|
+
type or value; expected/got), -32603 INPUT_LIMIT_PERSIST_FAILED
|
|
1708
|
+
(unchanged).
|
|
1709
|
+
"""
|
|
1710
|
+
named = {
|
|
1711
|
+
"max_input_tokens": max_input_tokens,
|
|
1712
|
+
"tokenizer": tokenizer,
|
|
1713
|
+
"max_inputs_per_request": max_inputs_per_request,
|
|
1714
|
+
"max_total_tokens_per_request": max_total_tokens_per_request,
|
|
1715
|
+
"source": source,
|
|
1716
|
+
}
|
|
1717
|
+
params: Dict[str, Any] = {"provider": provider, "model": model}
|
|
1718
|
+
for name, value in named.items():
|
|
1719
|
+
if value is not _NOT_GIVEN:
|
|
1720
|
+
params[name] = value
|
|
1721
|
+
return await self._service_command("input_limit_update", params)
|
|
1722
|
+
|
|
1723
|
+
async def input_limit_list(self, provider: Optional[str] = None) -> Dict[str, Any]:
|
|
1724
|
+
"""List input-limit rows, sorted by (provider, model), unique by key.
|
|
1725
|
+
|
|
1726
|
+
Reads no model and makes no provider call. Deletion marks are not shown.
|
|
1727
|
+
|
|
1728
|
+
Args:
|
|
1729
|
+
provider: Optional provider filter; sent only when not ``None``.
|
|
1730
|
+
|
|
1731
|
+
Returns:
|
|
1732
|
+
``{rows: [{provider, model, max_input_tokens, tokenizer,
|
|
1733
|
+
max_inputs_per_request, max_total_tokens_per_request, source,
|
|
1734
|
+
origin: "module" | "manual"}]}``; a provider with no rows gives
|
|
1735
|
+
``rows: []``.
|
|
1736
|
+
|
|
1737
|
+
Raises:
|
|
1738
|
+
EmbedAPIError: -32602 validation (unknown parameter, value).
|
|
1739
|
+
"""
|
|
1740
|
+
params: Dict[str, Any] = {}
|
|
1741
|
+
if provider is not None:
|
|
1742
|
+
params["provider"] = provider
|
|
1743
|
+
return await self._service_command("input_limit_list", params)
|
|
1744
|
+
|
|
1745
|
+
async def list_commands(self) -> Dict[str, Any]:
|
|
1746
|
+
"""
|
|
1747
|
+
List available commands.
|
|
1748
|
+
|
|
1749
|
+
Returns:
|
|
1750
|
+
Available commands information
|
|
1751
|
+
"""
|
|
1752
|
+
return await list_commands_via_client(self.client.jsonrpc_call)
|
|
1753
|
+
|
|
1754
|
+
async def cmd(
|
|
1755
|
+
self,
|
|
1756
|
+
command: str,
|
|
1757
|
+
params: Optional[Dict[str, Any]] = None,
|
|
1758
|
+
*,
|
|
1759
|
+
auto_poll: bool = True,
|
|
1760
|
+
) -> Dict[str, Any]:
|
|
1761
|
+
"""Execute a command via the adapter unified path (CLI and legacy callers).
|
|
1762
|
+
|
|
1763
|
+
Callers of this package depend on its own error type, not on whichever
|
|
1764
|
+
exception the transport happens to raise, so an adapter failure and a
|
|
1765
|
+
command that answered with a refusal both arrive as
|
|
1766
|
+
:class:`EmbedAPIError` carrying the server's own error.
|
|
1767
|
+
"""
|
|
1768
|
+
try:
|
|
1769
|
+
result = await self.client.execute_command_unified(
|
|
1770
|
+
command=command,
|
|
1771
|
+
params=params or {},
|
|
1772
|
+
auto_poll=auto_poll,
|
|
1773
|
+
)
|
|
1774
|
+
except EmbedError:
|
|
1775
|
+
raise
|
|
1776
|
+
except Exception as exc:
|
|
1777
|
+
if _missing_from_worker_registry(exc):
|
|
1778
|
+
try:
|
|
1779
|
+
return await self._cmd_over_jsonrpc(command, params or {})
|
|
1780
|
+
except EmbedError:
|
|
1781
|
+
raise
|
|
1782
|
+
except Exception as direct_exc:
|
|
1783
|
+
raise EmbedAPIError(str(direct_exc)) from direct_exc
|
|
1784
|
+
raise EmbedAPIError(str(exc)) from exc
|
|
1785
|
+
failure = _find_failed_command_result(result)
|
|
1786
|
+
if failure is not None:
|
|
1787
|
+
raise EmbedAPIError(failure)
|
|
1788
|
+
return result
|
|
1789
|
+
|
|
1790
|
+
async def _cmd_over_jsonrpc(
|
|
1791
|
+
self, command: str, params: Dict[str, Any]
|
|
1792
|
+
) -> Dict[str, Any]:
|
|
1793
|
+
"""Run *command* over plain JSON-RPC and return its payload."""
|
|
1794
|
+
response = await self.client.jsonrpc_call(command, params)
|
|
1795
|
+
if isinstance(response, dict) and response.get("error") is not None:
|
|
1796
|
+
raise EmbedAPIError(response["error"])
|
|
1797
|
+
payload = response.get("result") if isinstance(response, dict) else None
|
|
1798
|
+
if not isinstance(payload, dict):
|
|
1799
|
+
raise EmbedAPIError(f"Command {command!r} returned {response!r}")
|
|
1800
|
+
failure = _find_failed_command_result(payload)
|
|
1801
|
+
if failure is not None:
|
|
1802
|
+
raise EmbedAPIError(failure)
|
|
1803
|
+
return payload
|
|
1804
|
+
|
|
1805
|
+
async def list_queued_commands(
|
|
1806
|
+
self,
|
|
1807
|
+
status: Optional[str] = None,
|
|
1808
|
+
limit: Optional[int] = None,
|
|
1809
|
+
) -> Dict[str, Any]:
|
|
1810
|
+
"""List queued jobs (adapter ``queue_list_jobs``).
|
|
1811
|
+
|
|
1812
|
+
The command takes ``status_filter`` and ``limit``; the transport helper
|
|
1813
|
+
of the same name sends ``status``/``job_type`` instead and the server
|
|
1814
|
+
rejects the call, so the request is made directly.
|
|
1815
|
+
"""
|
|
1816
|
+
params: Dict[str, Any] = {}
|
|
1817
|
+
if status:
|
|
1818
|
+
params["status_filter"] = status
|
|
1819
|
+
if limit:
|
|
1820
|
+
params["limit"] = int(limit)
|
|
1821
|
+
response = await self.client.jsonrpc_call("queue_list_jobs", params)
|
|
1822
|
+
if isinstance(response, dict) and response.get("error") is not None:
|
|
1823
|
+
raise EmbedAPIError(response["error"])
|
|
1824
|
+
result = response.get("result") if isinstance(response, dict) else None
|
|
1825
|
+
if not isinstance(result, dict):
|
|
1826
|
+
raise EmbedAPIError(f"queue_list_jobs returned {response!r}")
|
|
1827
|
+
failure = _find_failed_command_result(result)
|
|
1828
|
+
if failure is not None:
|
|
1829
|
+
raise EmbedAPIError(failure)
|
|
1830
|
+
return result
|
|
1831
|
+
|
|
1832
|
+
async def cancel_command(self, job_id: str) -> Dict[str, Any]:
|
|
1833
|
+
"""Stop and delete a queued job."""
|
|
1834
|
+
await self.client.queue_stop_job(job_id)
|
|
1835
|
+
return await self.client.queue_delete_job(job_id)
|
|
1836
|
+
|
|
1837
|
+
async def get_job_logs(self, job_id: str) -> Dict[str, Any]:
|
|
1838
|
+
"""Fetch logs for a queued job (adapter ``queue_get_job_logs``)."""
|
|
1839
|
+
return await self.client.queue_get_job_logs(job_id)
|
|
1840
|
+
|
|
1841
|
+
async def timing_stats(self, file_path: Optional[str] = None) -> Dict[str, Any]:
|
|
1842
|
+
"""Fetch timing statistics RPC (when enabled on server)."""
|
|
1843
|
+
params: Dict[str, Any] = {}
|
|
1844
|
+
if file_path is not None and str(file_path).strip():
|
|
1845
|
+
params["file_path"] = str(file_path).strip()
|
|
1846
|
+
response = await self.client.jsonrpc_call("timing_stats", params)
|
|
1847
|
+
if "error" in response:
|
|
1848
|
+
raise EmbedClientError(f"timing_stats failed: {response['error']}")
|
|
1849
|
+
result = response.get("result", {})
|
|
1850
|
+
if isinstance(result, dict) and result.get("success") is False:
|
|
1851
|
+
err = result.get("error", {})
|
|
1852
|
+
msg = (
|
|
1853
|
+
err.get("message", "Unknown error")
|
|
1854
|
+
if isinstance(err, dict)
|
|
1855
|
+
else str(err)
|
|
1856
|
+
)
|
|
1857
|
+
raise EmbedClientError(f"timing_stats failed: {msg}")
|
|
1858
|
+
return result.get("data", result)
|
|
1859
|
+
|
|
1860
|
+
async def __aenter__(self) -> "EmbeddingClient":
|
|
1861
|
+
"""Async context manager entry (standalone package convenience)."""
|
|
1862
|
+
return self
|
|
1863
|
+
|
|
1864
|
+
async def __aexit__(self, exc_type: Any, exc: Any, tb: Any) -> None:
|
|
1865
|
+
await self.close()
|
|
1866
|
+
|
|
1867
|
+
async def close(self) -> None:
|
|
1868
|
+
"""Release the adapter HTTP client (httpx) and related resources."""
|
|
1869
|
+
await self.client.close()
|
|
1870
|
+
|
|
1871
|
+
async def fetch_openapi_schema(self) -> Dict[str, Any]:
|
|
1872
|
+
"""
|
|
1873
|
+
Load ``GET /openapi.json`` using the same TLS and auth headers as JSON-RPC.
|
|
1874
|
+
|
|
1875
|
+
Delegates to ``JsonRpcClient.get_openapi_schema`` on the underlying client.
|
|
1876
|
+
"""
|
|
1877
|
+
return await self.client.get_openapi_schema()
|
|
1878
|
+
|
|
1879
|
+
async def http_get(self, path: str) -> tuple[int, str]:
|
|
1880
|
+
"""
|
|
1881
|
+
``GET`` a path under the service base URL (e.g. ``/``, ``/docs``).
|
|
1882
|
+
|
|
1883
|
+
Uses the shared async httpx client from ``JsonRpcClient.get_async_http_client``
|
|
1884
|
+
(same ``base_url`` and ``headers`` as JSON-RPC, including API key or
|
|
1885
|
+
``Authorization``).
|
|
1886
|
+
|
|
1887
|
+
Returns:
|
|
1888
|
+
``(status_code, response_text)``.
|
|
1889
|
+
"""
|
|
1890
|
+
client = await self.client.get_async_http_client()
|
|
1891
|
+
p = path if path.startswith("/") else f"/{path}"
|
|
1892
|
+
response = await client.get(
|
|
1893
|
+
f"{self.client.base_url}{p}",
|
|
1894
|
+
headers=self.client.headers,
|
|
1895
|
+
)
|
|
1896
|
+
return response.status_code, response.text
|
|
1897
|
+
|
|
1898
|
+
def __repr__(self) -> str:
|
|
1899
|
+
"""String representation of client."""
|
|
1900
|
+
return f"EmbeddingClient({self.protocol}://{self.host}:{self.port})"
|
|
1901
|
+
|
|
1902
|
+
@classmethod
|
|
1903
|
+
def from_config(cls, config: "ClientConfig") -> "EmbeddingClient":
|
|
1904
|
+
"""Build client from a loaded :class:`~embed_client.config.ClientConfig`."""
|
|
1905
|
+
return embedding_client_from_config_dict(config.get_all())
|
|
1906
|
+
|
|
1907
|
+
|
|
1908
|
+
def embedding_client_from_config_dict(config_dict: Dict[str, Any]) -> "EmbeddingClient":
|
|
1909
|
+
"""Build :class:`EmbeddingClient` from embed-client ``config_dict`` (adapter-shaped)."""
|
|
1910
|
+
from embed_client.adapter_config_factory import AdapterConfigFactory
|
|
1911
|
+
|
|
1912
|
+
p = AdapterConfigFactory.from_config_dict(config_dict)
|
|
1913
|
+
auth = config_dict.get("auth") or {}
|
|
1914
|
+
ssl = config_dict.get("ssl")
|
|
1915
|
+
client = EmbeddingClient(
|
|
1916
|
+
protocol=p["protocol"],
|
|
1917
|
+
host=p["host"],
|
|
1918
|
+
port=int(p["port"]),
|
|
1919
|
+
cert=p.get("cert"),
|
|
1920
|
+
key=p.get("key"),
|
|
1921
|
+
ca=p.get("ca"),
|
|
1922
|
+
check_hostname=bool(p.get("check_hostname", False)),
|
|
1923
|
+
token=p.get("token"),
|
|
1924
|
+
token_header=p.get("token_header"),
|
|
1925
|
+
timeout=float(p.get("timeout", 120.0)),
|
|
1926
|
+
)
|
|
1927
|
+
client._auth_method = auth.get("method")
|
|
1928
|
+
client._ssl_settings = dict(ssl) if isinstance(ssl, dict) else None
|
|
1929
|
+
return client
|