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

embed_client/cli_core.py

@@ -0,0 +1,427 @@
+"""
+Core implementation of the embed-client CLI.
+
+Provides the `VectorizationCLI` class and configuration helpers used by the
+user-facing `embed_client.cli` entry point.
+
+Author: Vasiliy Zdanovskiy
+email: vasilyvz@gmail.com
+"""
+
+from __future__ import annotations
+
+import json
+import sys
+import time
+from typing import Any, Callable, Dict, List, Optional
+
+from embed_client.embedding_client import EmbeddingClient
+from embed_client.client_factory import ClientFactory
+from embed_client.config import ClientConfig
+from embed_client.constants import DEFAULT_POLL_INTERVAL
+from embed_client.embed_file_output_validate import (
+    EmbedFileJobPerfReport,
+    format_perf_report_lines,
+    format_validation_checklist_lines,
+    validate_embed_file_job_result,
+)
+from embed_client.exceptions import EmbedClientError
+from embed_client.json_embed_file_preflight import iter_json_array_block_texts
+from embed_client.response_parsers import extract_embeddings
+
+
+def _embed_file_ws_progress_stderr(msg: Dict[str, Any]) -> None:
+    """Print one WebSocket/status JSON object per line to stderr (embed_file wait)."""
+    print(json.dumps(msg, ensure_ascii=False), file=sys.stderr, flush=True)
+
+
+class VectorizationCLI:
+    """CLI application for text vectorization using embed-client library.
+
+    This class encapsulates the core logic of the `embed-vectorize` command,
+    including connection management, health checks, embedding requests, and
+    queue operations. It relies exclusively on `EmbeddingClient`
+    and related helpers from the embed-client package.
+    """
+
+    def __init__(self) -> None:
+        """Initialize a new VectorizationCLI instance."""
+        self.client: Optional[EmbeddingClient] = None
+
+    async def connect(
+        self,
+        config: Optional[ClientConfig] = None,
+        config_dict: Optional[Dict[str, Any]] = None,
+    ) -> bool:
+        """Connect to the embedding service using provided configuration.
+
+        The client is created via `ClientFactory` or directly from a
+        `ClientConfig` / configuration dictionary. All low-level transport and
+        authentication details are handled by the embed-client library.
+
+        Args:
+            config: Optional `ClientConfig` instance.
+            config_dict: Optional raw configuration dictionary.
+
+        Returns:
+            True if connection was established successfully, False otherwise.
+        """
+        try:
+            if config is not None:
+                self.client = EmbeddingClient.from_config(config)
+            elif config_dict is not None:
+                from embed_client.embedding_client import (
+                    embedding_client_from_config_dict,
+                )
+
+                self.client = embedding_client_from_config_dict(config_dict)
+            else:
+                # Use ClientFactory.from_environment() as fallback
+                self.client = ClientFactory.from_environment()
+
+            await self.client.__aenter__()
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to connect: {exc}", file=sys.stderr)
+            return False
+
+    async def disconnect(self) -> None:
+        """Close the underlying client connection if it exists."""
+        if self.client is not None:
+            await self.client.close()
+
+    async def health_check(self) -> bool:
+        """Perform a health check against the embedding service."""
+        try:
+            assert self.client is not None
+            result = await self.client.health()
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Health check failed: {exc}", file=sys.stderr)
+            return False
+
+    async def get_models(self) -> bool:
+        """Fetch and print available embedding models from the service."""
+        try:
+            assert self.client is not None
+            result = await self.client.models()
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to get models: {exc}", file=sys.stderr)
+            return False
+
+    async def get_timing_stats(self, file_path: Optional[str] = None) -> bool:
+        """Fetch and print timing statistics from the service."""
+        try:
+            assert self.client is not None
+            result = await self.client.timing_stats(file_path=file_path)
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to get timing stats: {exc}", file=sys.stderr)
+            return False
+
+    async def embed_file_local_json(
+        self,
+        json_path: str,
+        *,
+        max_block_bytes: Optional[int] = None,
+        upload_filename: Optional[str] = None,
+        wait_timeout: Optional[float] = None,
+        download_to: Optional[str] = None,
+        skip_download: bool = False,
+        model: Optional[str] = None,
+        dimension: Optional[int] = None,
+        device: Optional[str] = None,
+        error_policy: str = "continue",
+        show_progress: bool = False,
+    ) -> bool:
+        """Local JSON array file → adapter upload → ``embed_file`` (json_array)."""
+        try:
+            assert self.client is not None
+            on_ws: Optional[Callable[[Dict[str, Any]], None]] = None
+            if show_progress:
+                on_ws = _embed_file_ws_progress_stderr
+            ef_kw: Dict[str, Any] = {
+                "max_block_bytes": max_block_bytes,
+                "upload_filename": upload_filename,
+                "wait": True,
+                "download_to": download_to,
+                "skip_download": skip_download,
+                "model": model,
+                "dimension": dimension,
+                "device": device,
+                "error_policy": error_policy,
+                "on_ws_event": on_ws,
+            }
+            if wait_timeout is not None:
+                ef_kw["wait_timeout"] = int(wait_timeout)
+            blocks = iter_json_array_block_texts(json_path)
+            t0 = time.perf_counter()
+            result = await self.client.embed_file_from_local_json_file(
+                json_path,
+                **ef_kw,
+            )
+            elapsed = time.perf_counter() - t0
+            perf = EmbedFileJobPerfReport(
+                total_seconds=elapsed, block_count=len(blocks)
+            )
+            for ln in format_perf_report_lines(perf):
+                print(ln, file=sys.stderr, flush=True)
+            try:
+                vreport = validate_embed_file_job_result(
+                    blocks, result, rpc_dimension=dimension
+                )
+            except EmbedClientError as ve:
+                print(f"❌ embed_file output validation: {ve}", file=sys.stderr)
+                print(json.dumps(result, indent=2))
+                return False
+            for ln in format_validation_checklist_lines(vreport):
+                print(ln, file=sys.stderr, flush=True)
+            if not vreport.ok:
+                print("❌ embed_file output validation failed.", file=sys.stderr)
+                print(json.dumps(result, indent=2))
+                return False
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ embed-file-json failed: {exc}", file=sys.stderr)
+            return False
+
+    async def vectorize_texts(
+        self,
+        texts: List[str],
+        output_format: str = "json",
+        show_full_data: bool = False,
+        timeout: Optional[float] = None,
+        model: Optional[str] = None,
+        dimension: Optional[int] = None,
+        error_policy: str = "continue",
+        device: Optional[str] = None,
+    ) -> Optional[List[List[float]]]:
+        """Vectorize texts using the high-level `embed` helper.
+
+        Notes:
+            - Default ``error_policy=\"continue\"`` so per-item errors are
+              exposed via ``data[\"results\"][i][\"error\"]``.
+            - All input validation is delegated to the server.
+
+        Args:
+            texts: List of input texts to embed.
+            output_format: One of ``\"json\"``, ``\"csv\"``, or ``\"vectors\"``.
+            show_full_data: If True, print full embedding payload instead of
+                only vectors.
+            timeout: Maximum time to wait for vectorization completion in seconds.
+                    If None, return immediately without waiting.
+                    If 0, wait indefinitely (no timeout).
+                    If > 0, wait up to specified seconds.
+            model: Optional model name (server-side).
+            dimension: Optional expected vector dimension (server-side).
+            error_policy: "continue" or "fail_fast" (default: "continue").
+            device: Optional device: "auto", "cpu", "cuda", "cuda:0" (server-side).
+
+        Returns:
+            A list of embedding vectors on success (possibly empty when
+            `show_full_data` is True and no plain vectors are available),
+            or ``None`` on failure.
+        """
+        try:
+            assert self.client is not None
+
+            wt = 60 if timeout is None else int(timeout)
+            data = await self.client.embed(
+                texts,
+                model=model,
+                dimension=dimension,
+                error_policy=error_policy,
+                device=device,
+                wait=True,
+                wait_timeout=wt,
+            )
+
+            if show_full_data:
+                if output_format == "json":
+                    print(json.dumps(data, indent=2))
+                    return data.get("embeddings") or []
+                embeddings = data.get("embeddings") or []
+                if output_format == "csv":
+                    self._print_csv(embeddings)
+                elif output_format == "vectors":
+                    self._print_vectors(embeddings)
+                return embeddings
+
+            embeddings = data.get("embeddings")
+            if embeddings is None:
+                # Fallback to response_parsers for legacy-style responses
+                wrapped: Dict[str, Any] = {"result": {"success": True, "data": data}}
+                embeddings = extract_embeddings(wrapped)
+
+            if output_format == "json":
+                print(json.dumps(embeddings, indent=2))
+            elif output_format == "csv":
+                self._print_csv(embeddings)
+            elif output_format == "vectors":
+                self._print_vectors(embeddings)
+
+            return embeddings
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Vectorization failed: {exc}", file=sys.stderr)
+            return None
+
+    def _print_csv(self, embeddings: List[List[float]]) -> None:
+        """Print embeddings in CSV format."""
+        for idx, embedding in enumerate(embeddings):
+            print(f"text_{idx}," + ",".join(map(str, embedding)))
+
+    def _print_vectors(self, embeddings: List[List[float]]) -> None:
+        """Print embeddings as human-readable vectors."""
+        for idx, embedding in enumerate(embeddings):
+            print(f"Text {idx}: [{', '.join(map(str, embedding))}]")
+
+    async def get_help(self, command: Optional[str] = None) -> bool:
+        """Fetch help information from the embedding service."""
+        try:
+            assert self.client is not None
+            if command:
+                result = await self.client.cmd("help", params={"command": command})
+            else:
+                result = await self.client.cmd("help")
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Help request failed: {exc}", file=sys.stderr)
+            return False
+
+    async def get_commands(self) -> bool:
+        """List all available commands from the embedding service."""
+        try:
+            assert self.client is not None
+            result = await self.client.list_commands()
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Commands request failed: {exc}", file=sys.stderr)
+            return False
+
+    async def queue_list(
+        self, status: Optional[str] = None, limit: Optional[int] = None
+    ) -> bool:
+        """List queued commands."""
+        try:
+            assert self.client is not None
+            result = await self.client.list_queued_commands(status=status, limit=limit)
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to list queue: {exc}", file=sys.stderr)
+            return False
+
+    async def queue_status(self, job_id: str) -> bool:
+        """Get the status of a specific job."""
+        try:
+            assert self.client is not None
+            result = await self.client.job_status(job_id)
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to get job status: {exc}", file=sys.stderr)
+            return False
+
+    async def queue_cancel(self, job_id: str) -> bool:
+        """Cancel a queued or running job."""
+        try:
+            assert self.client is not None
+            result = await self.client.cancel_command(job_id)
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to cancel job: {exc}", file=sys.stderr)
+            return False
+
+    async def queue_wait(
+        self,
+        job_id: str,
+        timeout: float = 60.0,
+        poll_interval: float = DEFAULT_POLL_INTERVAL,
+    ) -> bool:
+        """Wait for completion of a specific job."""
+        try:
+            assert self.client is not None
+            result = await self.client.wait_for_job(
+                job_id, timeout=timeout, poll_interval=poll_interval
+            )
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to wait for job: {exc}", file=sys.stderr)
+            return False
+
+    async def queue_logs(self, job_id: str) -> bool:
+        """Fetch logs for a specific job."""
+        try:
+            assert self.client is not None
+            result = await self.client.get_job_logs(job_id)
+            print(json.dumps(result, indent=2))
+            return True
+        except Exception as exc:  # noqa: BLE001
+            print(f"❌ Failed to get job logs: {exc}", file=sys.stderr)
+            return False
+
+
+def create_config_from_args(args: Any) -> Dict[str, Any]:
+    """Create a configuration dictionary from parsed CLI arguments.
+
+    This helper mirrors the shape of the JSON configuration files used by
+    `ClientConfig` and `ClientFactory`, allowing the CLI to construct a
+    `ClientConfig` instance without requiring callers to manage JSON files.
+
+    The ``--insecure`` (``--no-check-hostname``) flag maps to
+    ``ssl.check_hostname = False``, which the library propagates as
+    ``check_hostname=False`` to both the HTTPS transport and the WebSocket
+    (wss) SSL context. Default (no flag) keeps strict TLS verification.
+    """
+    insecure = bool(getattr(args, "insecure", False))
+    config: Dict[str, Any] = {
+        "server": {"host": args.host, "port": args.port},
+        "auth": {"method": "none"},
+        "ssl": {"enabled": False, "check_hostname": not insecure},
+        "client": {"timeout": getattr(args, "timeout", 30.0)},
+    }
+
+    # Add authentication if specified
+    if getattr(args, "api_key", None):
+        config["auth"]["method"] = "api_key"
+        config["auth"]["api_keys"] = {"user": args.api_key}
+        if getattr(args, "api_key_header", None):
+            config["auth"]["api_key_header"] = args.api_key_header
+    elif getattr(args, "jwt_secret", None):
+        config["auth"]["method"] = "jwt"
+        config["auth"]["jwt"] = {
+            "secret": args.jwt_secret,
+            "username": getattr(args, "jwt_username", ""),
+            "password": getattr(args, "jwt_password", ""),
+        }
+    elif getattr(args, "basic_username", None):
+        config["auth"]["method"] = "basic"
+        config["auth"]["basic"] = {
+            "username": args.basic_username,
+            "password": getattr(args, "basic_password", ""),
+        }
+
+    # Add SSL if specified
+    if (
+        getattr(args, "ssl", False)
+        or getattr(args, "cert_file", None)
+        or getattr(args, "key_file", None)
+    ):
+        config["ssl"]["enabled"] = True
+        if getattr(args, "cert_file", None):
+            config["ssl"]["cert_file"] = args.cert_file
+        if getattr(args, "key_file", None):
+            config["ssl"]["key_file"] = args.key_file
+        if getattr(args, "ca_cert_file", None):
+            config["ssl"]["ca_cert_file"] = args.ca_cert_file
+
+    return config