embed-client-dev 4.0.78__py3-none-any.whl

embed_client/embed_file_output_validate.py

@@ -0,0 +1,348 @@
+"""
+Validate ``embed_file`` downloaded JSON against local ``json_array`` input.
+
+Aligned with ``../embed`` incremental writer shape:
+``{"success": true, "data": {"model", "dimension", "device?", "results": [...]}}``.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+from embed_client.exceptions import EmbedClientError
+
+
+def _is_meaningful_error(err: Any) -> bool:
+    if err is None:
+        return False
+    if isinstance(err, str):
+        return bool(err.strip())
+    if isinstance(err, dict):
+        return bool(str(err.get("message", "")).strip()) or bool(err)
+    return True
+
+
+def load_embed_file_output_json_file(path: str) -> Dict[str, Any]:
+    """Load UTF-8 JSON object from a local file (downloaded embed_file output)."""
+    ap = os.path.abspath(path)
+    if not os.path.isfile(ap):
+        raise EmbedClientError(f"embed_file output file not found: {ap}")
+    if not os.access(ap, os.R_OK):
+        raise EmbedClientError(f"embed_file output file not readable: {ap}")
+    try:
+        with open(ap, encoding="utf-8-sig") as f:
+            raw = f.read()
+    except OSError as e:
+        raise EmbedClientError(f"cannot read embed_file output: {ap}: {e}") from e
+    try:
+        obj: Any = json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise EmbedClientError(
+            f"invalid JSON in embed_file output file {ap}: {e}"
+        ) from e
+    if not isinstance(obj, dict):
+        raise EmbedClientError(
+            f"embed_file output JSON root must be an object, got {type(obj).__name__}"
+        )
+    return obj
+
+
+def _inner_data_dict(payload: Dict[str, Any]) -> Optional[Dict[str, Any]]:
+    d = payload.get("data")
+    return d if isinstance(d, dict) else None
+
+
+def parse_embed_file_output_records(
+    payload: Dict[str, Any],
+) -> Tuple[List[Dict[str, Any]], Optional[int], Optional[str]]:
+    """
+    Extract ``results`` rows and ``dimension`` from a downloaded or inline payload.
+
+    Supports:
+    - Canonical: ``success`` + ``data.results`` + ``data.dimension``
+    - ``data.results`` without top-level ``success``
+    - Top-level ``results`` only (tests / simplified files)
+    """
+    inner = _inner_data_dict(payload)
+    if inner is not None and isinstance(inner.get("results"), list):
+        dim = inner.get("dimension")
+        dim_i = int(dim) if isinstance(dim, int) else None
+        rows = inner["results"]
+        if not all(isinstance(r, dict) for r in rows):
+            raise EmbedClientError("embed_file data.results must be a list of objects")
+        return rows, dim_i, "data.results"
+
+    if isinstance(payload.get("results"), list):
+        dim = payload.get("dimension")
+        dim_i = int(dim) if isinstance(dim, int) else None
+        rows = payload["results"]
+        if not all(isinstance(r, dict) for r in rows):
+            raise EmbedClientError("embed_file results must be a list of objects")
+        return rows, dim_i, "results"
+
+    raise EmbedClientError(
+        "unrecognized embed_file output shape: expected data.results or top-level results"
+    )
+
+
+def resolve_embed_file_output_payload(
+    embed_result: Dict[str, Any],
+) -> Tuple[Dict[str, Any], str]:
+    """
+    Load JSON payload from ``local_output_path`` or parse inline structure in *embed_result*.
+
+    Returns ``(payload_dict, source_description)`` for checklist reporting.
+    """
+    path = embed_result.get("local_output_path")
+    if isinstance(path, str) and path.strip():
+        ap = os.path.abspath(path.strip())
+        if os.path.isfile(ap):
+            pl = load_embed_file_output_json_file(ap)
+            return pl, f"local_output_path file ({ap})"
+
+    # Inline: full document nested under merged job data (unusual but supported)
+    for key in ("embed_file_output", "output_json"):
+        v = embed_result.get(key)
+        if isinstance(v, dict) and (
+            _inner_data_dict(v) is not None or isinstance(v.get("results"), list)
+        ):
+            return v, f"inline key {key!r}"
+
+    # Try canonical nested under result (adapter echo)
+    res = embed_result.get("result")
+    if isinstance(res, dict):
+        inner = res.get("data") if isinstance(res.get("data"), dict) else res
+        if isinstance(inner, dict) and (
+            isinstance(inner.get("results"), list)
+            or (
+                isinstance(inner.get("data"), dict)
+                and isinstance(inner["data"].get("results"), list)
+            )
+        ):
+            pl = res if "success" in res else {"success": True, "data": inner}
+            return pl, "embed_result.result"
+
+    raise EmbedClientError(
+        "cannot validate embed_file output: no readable local_output_path and no "
+        "inline results payload in the job result"
+    )
+
+
+@dataclass
+class EmbedFileOutputValidationReport:
+    """Structured validation outcome + human-readable checklist lines."""
+
+    ok: bool
+    checklist_lines: List[str] = field(default_factory=list)
+    errors: List[str] = field(default_factory=list)
+
+
+def _expected_dimension(
+    *,
+    payload_dim: Optional[int],
+    rpc_dimension: Optional[int],
+) -> Optional[int]:
+    if rpc_dimension is not None:
+        return int(rpc_dimension)
+    return payload_dim
+
+
+def validate_embed_file_output_against_blocks(
+    input_block_texts: List[str],
+    payload: Dict[str, Any],
+    *,
+    rpc_dimension: Optional[int] = None,
+) -> EmbedFileOutputValidationReport:
+    """
+    Check record count, non-empty embeddings (or per-row error), dimension, and
+    ``body`` order vs *input_block_texts* when bodies are present.
+    """
+    lines: List[str] = []
+    errors: List[str] = []
+
+    if payload.get("success") is False:
+        errors.append("payload.success is false in embed_file output document")
+        lines.append("FAIL: payload.success is false.")
+        return EmbedFileOutputValidationReport(
+            ok=False, checklist_lines=lines, errors=errors
+        )
+
+    rows, payload_dim, path_hint = parse_embed_file_output_records(payload)
+    lines.append(f"Parsed output records from {path_hint}.")
+    dim = _expected_dimension(payload_dim=payload_dim, rpc_dimension=rpc_dimension)
+    if dim is not None:
+        lines.append(f"Expected embedding dimension: {dim} (RPC or payload).")
+    else:
+        lines.append(
+            "Expected embedding dimension: not specified (skipping length checks)."
+        )
+
+    n_in = len(input_block_texts)
+    n_out = len(rows)
+    if n_out != n_in:
+        msg = f"record count mismatch: output has {n_out} rows, input has {n_in} blocks"
+        errors.append(msg)
+        lines.append(f"FAIL: {msg}")
+    else:
+        lines.append(f"Record count matches input blocks ({n_in}).")
+
+    ordering_warn: Optional[str] = None
+    for i, row in enumerate(rows):
+        emb = row.get("embedding")
+        err = row.get("error")
+        body = row.get("body")
+        if isinstance(emb, list):
+            if len(emb) == 0:
+                errors.append(f"row {i}: empty embedding list")
+            elif dim is not None and len(emb) != dim:
+                errors.append(
+                    f"row {i}: embedding length {len(emb)} != expected dimension {dim}"
+                )
+        elif not _is_meaningful_error(err):
+            errors.append(
+                f"row {i}: missing embedding and no meaningful error "
+                f"(got embedding={type(emb).__name__!r})"
+            )
+        if isinstance(body, str) and i < n_in:
+            if body != input_block_texts[i]:
+                ordering_warn = (
+                    f"row {i}: body text does not match input block at same index "
+                    "(possible reorder or transform)"
+                )
+
+    if ordering_warn:
+        errors.append(ordering_warn)
+        lines.append(f"ORDERING: {ordering_warn}")
+    else:
+        comparable = sum(
+            1
+            for i, row in enumerate(rows)
+            if i < n_in and isinstance(row.get("body"), str)
+        )
+        if comparable == n_in and n_in > 0:
+            lines.append("Body text matches input block order (index-wise).")
+        elif n_in > 0:
+            lines.append(
+                "Body fields absent or not comparable; skipped strict order check."
+            )
+
+    ok = len(errors) == 0
+    if ok:
+        lines.append("All embedding / error-field checks passed for parsed rows.")
+    else:
+        lines.append(f"Validation failed with {len(errors)} issue(s).")
+
+    return EmbedFileOutputValidationReport(ok=ok, checklist_lines=lines, errors=errors)
+
+
+def validate_embed_file_job_result(
+    input_block_texts: List[str],
+    embed_result: Dict[str, Any],
+    *,
+    rpc_dimension: Optional[int] = None,
+) -> EmbedFileOutputValidationReport:
+    """
+    Resolve payload from *embed_result*, then run
+    :func:`validate_embed_file_output_against_blocks`.
+    """
+    checklist: List[str] = []
+    payload, src = resolve_embed_file_output_payload(embed_result)
+    checklist.append(f"Artifact source: {src}.")
+    inner = validate_embed_file_output_against_blocks(
+        input_block_texts,
+        payload,
+        rpc_dimension=rpc_dimension,
+    )
+    checklist.extend(inner.checklist_lines)
+    return EmbedFileOutputValidationReport(
+        ok=inner.ok,
+        checklist_lines=checklist,
+        errors=list(inner.errors),
+    )
+
+
+@dataclass
+class EmbedFileJobPerfReport:
+    """Wall-clock metrics for a single embed_file_from_local_json_file run."""
+
+    total_seconds: float
+    block_count: int
+    note: str = (
+        "Single perf_counter interval covers upload, queue wait, download, and ACK; "
+        "per-block timings are not available from this client."
+    )
+
+    @property
+    def mean_seconds_per_block(self) -> float:
+        if self.block_count <= 0:
+            return 0.0
+        return self.total_seconds / float(self.block_count)
+
+    @property
+    def median_seconds_per_block(self) -> float:
+        # One interval → median time per block equals mean (total/N).
+        return self.mean_seconds_per_block
+
+    @property
+    def blocks_per_second(self) -> float:
+        if self.total_seconds <= 0:
+            return 0.0
+        return float(self.block_count) / self.total_seconds
+
+    @property
+    def mean_blocks_per_second(self) -> float:
+        return self.blocks_per_second
+
+    @property
+    def median_blocks_per_second(self) -> float:
+        return self.blocks_per_second
+
+
+def format_perf_report_lines(report: EmbedFileJobPerfReport) -> List[str]:
+    """Return printable lines for stderr or stdout."""
+    lines = [
+        "--- embed_file job performance ---",
+        report.note,
+        f"Total wall time: {report.total_seconds:.6f} s",
+        f"Input blocks (N): {report.block_count}",
+        f"Throughput: {report.blocks_per_second:.6f} blocks/s",
+        f"Mean time per block: {report.mean_seconds_per_block:.6f} s/block "
+        "(total_seconds / N; only one wall-clock interval measured).",
+        f"Median time per block: {report.median_seconds_per_block:.6f} s/block "
+        "(equals mean when per-block timings are unavailable).",
+        f"Mean speed: {report.mean_blocks_per_second:.6f} blocks/s",
+        f"Median speed: {report.median_blocks_per_second:.6f} blocks/s",
+    ]
+    return lines
+
+
+def format_validation_checklist_lines(
+    report: EmbedFileOutputValidationReport,
+) -> List[str]:
+    """Prefix lines for clear CLI / example output."""
+    out = ["--- embed_file output validation ---"]
+    out.extend(report.checklist_lines)
+    if not report.ok:
+        out.append("Issues:")
+        for e in report.errors:
+            out.append(f"  - {e}")
+    return out
+
+
+__all__ = [
+    "EmbedFileJobPerfReport",
+    "EmbedFileOutputValidationReport",
+    "format_perf_report_lines",
+    "format_validation_checklist_lines",
+    "load_embed_file_output_json_file",
+    "parse_embed_file_output_records",
+    "resolve_embed_file_output_payload",
+    "validate_embed_file_job_result",
+    "validate_embed_file_output_against_blocks",
+]

embed_client/embed_file_status_store.py

@@ -0,0 +1,80 @@
+"""
+In-memory embed_file job status store (per EmbeddingClient instance).
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+from __future__ import annotations
+
+import copy
+import threading
+import time
+from typing import Any, Dict, List, Optional
+
+
+class EmbedFileJobStatusStore:
+    """
+    Thread-safe accumulation of embed_file RPC submit payload and WebSocket frames.
+
+    Keys are ``job_id`` strings. Each entry may include ``submit``, ``messages``,
+    ``last_event``, and timestamps. Payloads are deep-copied on write.
+    """
+
+    def __init__(self) -> None:
+        self._lock = threading.Lock()
+        self._jobs: Dict[str, Dict[str, Any]] = {}
+
+    def record_submit(self, job_id: str, submit_payload: Dict[str, Any]) -> None:
+        """Record the initial ``embed_file`` RPC success payload for ``job_id``."""
+        with self._lock:
+            entry = self._jobs.setdefault(
+                job_id,
+                {
+                    "job_id": job_id,
+                    "messages": [],
+                },
+            )
+            entry["submit"] = copy.deepcopy(submit_payload)
+            entry["submit_at"] = time.time()
+
+    def record_message(self, job_id: str, event: Dict[str, Any]) -> None:
+        """Append one WebSocket JSON payload (queue job subscription) for ``job_id``."""
+        with self._lock:
+            entry = self._jobs.setdefault(
+                job_id,
+                {
+                    "job_id": job_id,
+                    "messages": [],
+                },
+            )
+            entry["messages"].append(copy.deepcopy(event))
+            entry["last_event"] = copy.deepcopy(event)
+            entry["last_message_at"] = time.time()
+
+    def snapshot(self, job_id: str) -> Optional[Dict[str, Any]]:
+        """Return a deep copy of stored state for ``job_id``, or ``None`` if unknown."""
+        with self._lock:
+            if job_id not in self._jobs:
+                return None
+            return copy.deepcopy(self._jobs[job_id])
+
+    def messages(self, job_id: str) -> List[Dict[str, Any]]:
+        """Return a deep copy of accumulated WS messages for ``job_id``."""
+        with self._lock:
+            if job_id not in self._jobs:
+                return []
+            return copy.deepcopy(self._jobs[job_id].get("messages", []))
+
+    def clear(self, job_id: str) -> None:
+        """Remove stored data for ``job_id``."""
+        with self._lock:
+            self._jobs.pop(job_id, None)
+
+    def job_ids(self) -> List[str]:
+        """List all ``job_id`` keys currently present."""
+        with self._lock:
+            return list(self._jobs.keys())
+
+
+__all__ = ["EmbedFileJobStatusStore"]