tos-vectors-embed-cli 0.1.0__py3-none-any.whl

tos_vectors/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by Beijing Volcano Engine Technology Co., Ltd. on 2026-02-12
+#
+# Original file was released under Apache License 2.0, with the full license text
+# available at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# This modified file is released under the same license.
+"""TOS Vectors Embed CLI package."""

tos_vectors/__version__.py

@@ -0,0 +1,14 @@
+# Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by Beijing Volcano Engine Technology Co., Ltd. on 2026-02-12
+#
+# Original file was released under Apache License 2.0, with the full license text
+# available at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# This modified file is released under the same license.
+"""Version information for TOS Vectors CLI."""
+
+__version__ = "0.1.0"
+__version_info__ = tuple(int(i) for i in __version__.split('.'))

tos_vectors/cli.py

@@ -0,0 +1,62 @@
+#!/usr/bin/env python3
+# Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by Beijing Volcano Engine Technology Co., Ltd. on 2026-02-12
+#
+# Original file was released under Apache License 2.0, with the full license text
+# available at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# This modified file is released under the same license.
+"""Main CLI entry point for TOS Vectors."""
+
+import click
+from rich.console import Console
+from rich.traceback import install
+from tos_vectors.commands.embed_put import embed_put
+from tos_vectors.commands.embed_query import embed_query
+
+# Install rich traceback handler
+install(show_locals=True)
+console = Console()
+
+
+@click.group()
+@click.version_option(version="0.1.0")
+@click.option('--account-id', help='Volcengine account id to use')
+@click.option('--region', default='cn-beijing', help='TOS region name')
+@click.option('--debug', is_flag=True, help='Enable debug mode with detailed logging')
+@click.pass_context
+def cli(ctx, account_id, region, debug):
+    """TOS Vectors Embed CLI - Standalone tool for vector embedding operations with TOS and Ark."""
+    ctx.ensure_object(dict)
+    ctx.obj['account_id'] = account_id
+    ctx.obj['region'] = region
+    ctx.obj['console'] = console
+    ctx.obj['debug'] = debug
+
+    if debug:
+        console.print("[dim] Debug mode enabled[/dim]")
+        console.print(f"[dim] Volcengine Account ID: {account_id}[/dim]")
+        console.print(f"[dim] TOS Region: {region}[/dim]")
+
+
+# Register commands as subcommands
+cli.add_command(embed_put, name='put')
+cli.add_command(embed_query, name='query')
+
+
+def main():
+    """Main entry point."""
+    try:
+        cli()
+    except KeyboardInterrupt:
+        console.print("\n[yellow]Operation cancelled by user[/yellow]")
+    except Exception as e:
+        console.print(f"[red]Error: {str(e)}[/red]")
+        raise
+
+
+if __name__ == '__main__':
+    main()

tos_vectors/commands/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by Beijing Volcano Engine Technology Co., Ltd. on 2026-02-12
+#
+# Original file was released under Apache License 2.0, with the full license text
+# available at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# This modified file is released under the same license.
+"""Commands for TOS Vectors CLI."""

tos_vectors/commands/embed_put.py

@@ -0,0 +1,389 @@
+# Copyright (c) 2024 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+# Copyright (c) 2026 Beijing Volcano Engine Technology Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+#
+# This file has been modified by Beijing Volcano Engine Technology Co., Ltd. on 2026-02-12
+#
+# Original file was released under Apache License 2.0, with the full license text
+# available at http://www.apache.org/licenses/LICENSE-2.0.
+#
+# This modified file is released under the same license.
+"""Command implementation for embedding and storing vectors."""
+
+import json
+import click
+from rich.progress import Progress, SpinnerColumn, TextColumn
+
+from tos_vectors.core.services import ArkService, TOSVectorService
+from tos_vectors.core.unified_processor import UnifiedProcessor
+from tos_vectors.utils.config import get_region
+from tos_vectors.utils.models import (
+    get_model_info,
+    validate_user_parameters,
+    prepare_processing_input,
+    determine_content_type
+)
+from tos_vectors.core.streaming_batch_orchestrator import (
+    StreamingBatchOrchestrator
+)
+
+
+def _create_progress_context(console):
+    """Create progress context for operations."""
+    return Progress(
+        SpinnerColumn(),
+        TextColumn("[progress.description]{task.description}"),
+        console=console,
+        transient=True
+    )
+
+
+def _validate_inputs(text_value, text, image, video, model, key, filename_as_key):
+    """Validate input parameters."""
+    inputs_provided = sum(bool(x) for x in [
+        text_value, text, image, video
+    ])
+
+    if inputs_provided == 0:
+        raise click.ClickException(
+            "At least one input must be provided: --text-value, --text, "
+            "--image, or --video"
+        )
+
+    # Check mutual exclusivity of key parameters
+    if key and filename_as_key:
+        raise click.ClickException(
+            "Cannot use both --key and --filename-as-key. Choose one."
+        )
+
+    # --filename-as-key not allowed with --text-value
+    if filename_as_key and text_value:
+        raise click.ClickException(
+            "--filename-as-key is not supported with --text-value "
+            "(no file or object to extract name from)"
+        )
+
+    # Special case: Allow multimodal input for supported models
+    is_multimodal_input = model.supports_multimodal_input() and sum(
+        1 for inp in [text_value, text, image, video]
+        if inp is not None
+    ) >= 2
+
+    if inputs_provided > 1 and not is_multimodal_input:
+        raise click.ClickException(
+            "Only one input type can be specified at a time, "
+            "except for multimodal input with supported models"
+        )
+
+    return is_multimodal_input
+
+
+@click.command()
+@click.option('--vector-bucket-name', required=True, help='TOS vector bucket name')
+@click.option('--index-name', required=True, help='Vector index name')
+@click.option('--model-id', required=True, help='Ark embedding model ID')
+@click.option('--text-value', help='Direct text input to embed')
+@click.option('--text', help='Text file path (local file or TOS URI)')
+@click.option('--image', help='Image file path (local file or TOS URI)')
+@click.option('--video', help='Video file path (local file)')
+@click.option('--key', help='Custom vector key (auto-generated UUID if not provided)')
+@click.option('--key-prefix', help='Prefix to prepend to all vector keys')
+@click.option('--filename-as-key', is_flag=True, help='Use filename as vector key')
+@click.option('--metadata', help='Additional metadata associated with the vector; provided as JSON string')
+@click.option('--ark-inference-params', help='JSON string with model-specific parameters matching Ark API format')
+@click.option('--max-workers', default=4, type=int, help='Maximum parallel workers for batch processing (default: 4)')
+@click.option('--batch-size', type=click.IntRange(1, 500), default=500, help='Vectors per TOS Vector put_vectors call (1-500, default: 500)')
+@click.option('--output', type=click.Choice(['json', 'table']), default='json', help='Output format')
+@click.option('--region', help='TOS region name (effective in TOS path mode)')
+@click.pass_context
+def embed_put(ctx, vector_bucket_name, index_name, model_id, text_value, text, image, video, ark_inference_params, key, key_prefix, filename_as_key, metadata, max_workers, batch_size, output, region):
+    """Unified embed and store vectors command."""
+
+    console = ctx.obj['console']
+    account_id = ctx.obj.get('account_id')
+    debug = ctx.obj.get('debug', False)
+    region = get_region(region or ctx.obj.get('region'))
+
+    # Load model properties once at start
+    model = get_model_info(model_id)
+    if not model:
+        raise click.ClickException(f"Unsupported model: {model_id}")
+
+    # Parse parameters
+    user_ark_params = {}
+    if ark_inference_params:
+        try:
+            user_ark_params = json.loads(ark_inference_params)
+        except json.JSONDecodeError:
+            raise click.ClickException(
+                "Invalid JSON in --ark-inference-params parameter"
+            )
+
+    metadata_dict = {}
+    if metadata:
+        try:
+            metadata_dict = json.loads(metadata)
+        except json.JSONDecodeError:
+            raise click.ClickException("Invalid JSON in --metadata parameter")
+
+    # Early validation of user parameters before any processing
+    if user_ark_params:
+        try:
+            content_type = determine_content_type(
+                text_value,
+                text,
+                image,
+                video
+            )
+            system_keys = model.get_system_keys(content_type)
+            # Dummy values for validation
+            system_payload = {k: None for k in system_keys}
+
+            # Validate using utility function
+            validate_user_parameters(system_payload, user_ark_params)
+        except ValueError as e:
+            raise click.ClickException(str(e))
+
+    # Validate inputs
+    is_multimodal = _validate_inputs(
+        text_value,
+        text,
+        image,
+        video,
+        model,
+        key,
+        filename_as_key
+    )
+
+    try:
+        # Initialize services
+        ark_service = ArkService(region=region, debug=debug, console=console)
+        tos_vector_service = TOSVectorService(
+            region=region, account_id=account_id,
+            debug=debug, console=console
+        )
+
+        # Create unified processor
+        processor = UnifiedProcessor(ark_service, tos_vector_service, region=region)
+
+        # Fetch index dimensions once at the top level
+        try:
+            index_info = tos_vector_service.get_index(
+                vector_bucket_name, index_name
+            )
+            index_dimensions = index_info.get("dimension")
+            if not index_dimensions:
+                raise click.ClickException(
+                    f"Could not determine dimensions for index {index_name}"
+                )
+        except Exception as e:
+            raise click.ClickException(
+                f"Failed to get index information: {str(e)}"
+            )
+
+        # Prepare processing input
+        processing_input = prepare_processing_input(
+            text_value,
+            text,
+            image,
+            video,
+            is_multimodal,
+            metadata_dict,
+            key,
+            filename_as_key,
+            key_prefix
+        )
+
+        # Check for wildcard patterns (streaming batch processing)
+        eligible = (
+            processing_input.content_type in ["text", "image", "video"]
+            and "file_path" in processing_input.data
+        )
+
+        if eligible:
+            file_path = processing_input.data["file_path"]
+            if '*' in file_path or '?' in file_path:
+                _process_streaming_batch(
+                    file_path,
+                    processing_input.content_type,
+                    vector_bucket_name,
+                    index_name,
+                    model,
+                    metadata_dict,
+                    user_ark_params,
+                    processor,
+                    console,
+                    output,
+                    max_workers,
+                    batch_size,
+                    index_dimensions,
+                    processing_input.filename_as_key,
+                    processing_input.key_prefix
+                )
+                return
+
+        # Process input to generate embeddings
+        with _create_progress_context(console) as progress:
+            progress.add_task(
+                f"[cyan]Embedding {processing_input.content_type}..."
+            )
+            result = processor.process(
+                model=model,
+                processing_input=processing_input,
+                user_ark_params=user_ark_params,
+                vector_bucket_name=vector_bucket_name,
+                index_name=index_name,
+                precomputed_dimensions=index_dimensions
+            )
+
+            # Store vectors with batch_size handling
+            progress.add_task(
+                f"Storing {len(result.vectors)} vector(s)...",
+                total=None
+            )
+
+            # Handle batch_size for single file processing too
+            vector_count = len(result.vectors)
+
+            if vector_count <= batch_size:
+                stored_keys = processor.store_vectors(
+                    result.vectors,
+                    vector_bucket_name,
+                    index_name
+                )
+            else:
+                stored_keys = []
+                for i in range(0, vector_count, batch_size):
+                    chunk = result.vectors[i:i + batch_size]
+                    chunk_keys = processor.store_vectors(
+                        chunk,
+                        vector_bucket_name,
+                        index_name
+                    )
+                    stored_keys.extend(chunk_keys)
+
+        # Prepare output
+        if result.result_type == "multiclip":
+            output_result = {
+                'type': 'multiclip',
+                'bucket': vector_bucket_name,
+                'index': index_name,
+                'model': model.model_id,
+                'contentType': processing_input.content_type,
+                'totalVectors': len(stored_keys),
+                'keys': stored_keys
+            }
+            if result.job_id:
+                output_result['jobId'] = result.job_id
+        else:
+            output_result = {
+                'key': stored_keys[0],
+                'bucket': vector_bucket_name,
+                'index': index_name,
+                'model': model.model_id,
+                'contentType': processing_input.content_type,
+                'embeddingDimensions': index_dimensions,
+                'metadata': result.vectors[0]['metadata']
+            }
+
+        console.print_json(data=output_result)
+
+    except Exception as e:
+        raise click.ClickException(str(e))
+
+
+def _process_streaming_batch(file_path, content_type, vector_bucket_name, index_name, model, metadata_dict, user_ark_params, processor, console, output, max_workers, batch_size, index_dimensions, filename_as_key, key_prefix):
+    """Process wildcard pattern using streaming batch orchestrator."""
+
+    try:
+        # Create streaming batch orchestrator
+        streaming_orchestrator = StreamingBatchOrchestrator(
+            processor,
+            max_workers,
+            batch_size
+        )
+
+        console.print(f"Starting streaming batch processing: {file_path}")
+
+        # Process using streaming approach (no pre-loading of file paths)
+        batch_result = streaming_orchestrator.process_streaming_batch(
+            file_path,
+            content_type,
+            vector_bucket_name,
+            index_name,
+            model,
+            metadata_dict,
+            user_ark_params,
+            index_dimensions,
+            filename_as_key,
+            key_prefix
+        )
+
+        # Display results
+        result_dict = {
+            "type": "streaming_batch",
+            "bucket": vector_bucket_name,
+            "index": index_name,
+            "model": model.model_id,
+            "contentType": content_type,
+            "totalFiles": (
+                batch_result.processed_count + batch_result.failed_count
+            ),
+            "processedFiles": batch_result.processed_count,
+            "failedFiles": batch_result.failed_count,
+            "totalVectors": len(batch_result.processed_keys),
+            # Show first 10
+            "vectorKeys": (
+                batch_result.processed_keys[:10]
+                if batch_result.processed_keys else []
+            )
+        }
+
+        if batch_result.errors:
+            result_dict["errors"] = batch_result.errors[:10]
+
+        if output == "table":
+            _display_batch_table(result_dict, console)
+        else:
+            console.print_json(data=result_dict)
+
+        # Print display limit messages after output
+        if len(batch_result.processed_keys) > 10:
+            console.print(
+                f"[dim]Note: Showing first 10 of "
+                f"{len(batch_result.processed_keys)} vector keys[/dim]"
+            )
+
+        if batch_result.errors and len(batch_result.errors) > 10:
+            console.print(
+                f"[dim]Note: Showing first 10 of "
+                f"{len(batch_result.errors)} errors[/dim]"
+            )
+
+        return result_dict if output == "json" else None
+
+    except Exception as e:
+        console.print(
+            f"[red]Streaming batch processing failed: {str(e)}[/red]"
+        )
+        raise click.ClickException(
+            f"Streaming batch processing failed: {str(e)}"
+        )
+
+
+def _display_batch_table(result, console):
+    """Display batch results in table format."""
+    from rich.table import Table
+
+    table = Table(title="Batch Processing Results")
+    table.add_column("Metric", style="cyan")
+    table.add_column("Value", style="green")
+
+    table.add_row("Total Files", str(result["totalFiles"]))
+    table.add_row("Processed Files", str(result["processedFiles"]))
+    table.add_row("Failed Files", str(result["failedFiles"]))
+    table.add_row("Total Vectors", str(result["totalVectors"]))
+    table.add_row("Model", result["model"])
+    table.add_row("Content Type", result["contentType"])
+
+    console.print(table)