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.
Files changed (48) hide show
  1. embed_client/__init__.py +36 -0
  2. embed_client/__main__.py +15 -0
  3. embed_client/_version.py +3 -0
  4. embed_client/adapter_config_factory.py +213 -0
  5. embed_client/adapter_config_normalizer.py +135 -0
  6. embed_client/adapter_generator_bridge.py +166 -0
  7. embed_client/auth.py +189 -0
  8. embed_client/cli.py +483 -0
  9. embed_client/cli_core.py +427 -0
  10. embed_client/client_factory.py +515 -0
  11. embed_client/client_factory_examples.py +375 -0
  12. embed_client/client_factory_helpers.py +59 -0
  13. embed_client/config.py +408 -0
  14. embed_client/config_defaults.py +77 -0
  15. embed_client/config_examples.py +228 -0
  16. embed_client/config_generator.py +652 -0
  17. embed_client/config_generator_cli.py +151 -0
  18. embed_client/config_presets.py +245 -0
  19. embed_client/config_validator.py +246 -0
  20. embed_client/config_validator_cli.py +141 -0
  21. embed_client/constants.py +52 -0
  22. embed_client/demo_cli.py +78 -0
  23. embed_client/demo_security_cli.py +206 -0
  24. embed_client/embed_file_output_validate.py +348 -0
  25. embed_client/embed_file_status_store.py +80 -0
  26. embed_client/embedding_client.py +1929 -0
  27. embed_client/embedding_client_helpers.py +105 -0
  28. embed_client/example_async_usage.py +357 -0
  29. embed_client/example_async_usage_demo.py +282 -0
  30. embed_client/example_async_usage_ru.py +340 -0
  31. embed_client/example_async_usage_ru_demo.py +313 -0
  32. embed_client/exceptions.py +122 -0
  33. embed_client/identifier_types.py +56 -0
  34. embed_client/json_embed_file_preflight.py +128 -0
  35. embed_client/local_json_embed_file.py +73 -0
  36. embed_client/md_embed_file_prep.py +197 -0
  37. embed_client/response_normalizer.py +160 -0
  38. embed_client/response_parsers.py +219 -0
  39. embed_client/security_cli.py +200 -0
  40. embed_client/security_cli_core.py +227 -0
  41. embed_client/ssl_examples.py +278 -0
  42. embed_client/ssl_manager.py +141 -0
  43. embed_client_dev-4.0.78.dist-info/METADATA +155 -0
  44. embed_client_dev-4.0.78.dist-info/RECORD +48 -0
  45. embed_client_dev-4.0.78.dist-info/WHEEL +5 -0
  46. embed_client_dev-4.0.78.dist-info/entry_points.txt +4 -0
  47. embed_client_dev-4.0.78.dist-info/licenses/LICENSE +21 -0
  48. 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