opencode-skills-collection 3.1.9 → 3.1.11

package/bundled-skills/weaviate/scripts/import.py

@@ -0,0 +1,701 @@
+#!/usr/bin/env python3
+# /// script
+# dependencies = [
+#   "weaviate-client==4.19.2",
+#   "typer==0.21.0",
+#   "pdf2image>=1.17.0",
+#   "pillow>=10.0.0",
+# ]
+# ///
+"""
+Import data from CSV, JSON, JSONL, or PDF files to a Weaviate collection.
+
+Usage:
+    uv run import.py data.csv --collection "CollectionName" [options]
+    uv run import.py document.pdf --collection "CollectionName" [options]
+
+Environment Variables:
+    WEAVIATE_URL: Weaviate Cloud cluster URL
+    WEAVIATE_API_KEY: API key for authentication
+    + Any provider API keys (OPENAI_API_KEY, COHERE_API_KEY, etc.) - auto-detected
+"""
+
+import base64
+import csv
+import itertools
+import json
+import re
+import sys
+from collections.abc import Iterator
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+
+_DATE_RE = re.compile(r"^\d{4}-\d{2}-\d{2}$")
+_DATETIME_RE = re.compile(r"^\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}$")
+_RESERVED_FIELDS = {"id", "_additional"}
+
+import typer
+import weaviate
+from weaviate.classes.config import Configure, DataType, Property
+
+# Import shared connection utilities (local to this skill)
+from weaviate_conn import get_client
+
+# Types whose string values must never be JSON-parsed (already correct as strings)
+_KEEP_AS_STRING = {DataType.TEXT, DataType.UUID, DataType.BLOB}
+
+app = typer.Typer()
+
+
+def detect_file_format(file_path: Path) -> str:
+    """
+    Detect file format based on extension.
+
+    Args:
+        file_path: Path to the file
+
+    Returns:
+        File format: "csv", "json", or "jsonl"
+
+    Raises:
+        ValueError: If file format is not supported
+    """
+    extension = file_path.suffix.lower()
+
+    if extension == ".csv":
+        return "csv"
+    elif extension == ".json":
+        return "json"
+    elif extension == ".jsonl":
+        return "jsonl"
+    elif extension == ".pdf":
+        return "pdf"
+    else:
+        raise ValueError(
+            f"Unsupported file format: {extension}. "
+            f"Supported formats: .csv, .json, .jsonl, .pdf"
+        )
+
+
+def read_csv(
+    file_path: Path, mapping: dict[str, str] | None = None
+) -> Iterator[dict[str, Any]]:
+    """
+    Read data from CSV file with automatic dialect detection.
+
+    Yields rows one at a time — suitable for large files.
+
+    Args:
+        file_path: Path to CSV file
+        mapping: Optional column name mapping
+
+    Yields:
+        Row dictionaries with data
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        # Read a sample to detect the CSV dialect
+        sample = f.read(8192)
+        f.seek(0)
+
+        # Use Sniffer to detect the dialect (delimiter, quoting, etc.)
+        sniffer = csv.Sniffer()
+        try:
+            dialect = sniffer.sniff(sample)
+        except csv.Error:
+            dialect = csv.excel
+
+        reader = csv.DictReader(f, dialect=dialect)
+
+        # Warn if the header row looks like data (all-numeric or JSON-like values
+        # suggest the file has no header row and the first data row was misread as one).
+        if reader.fieldnames:
+            suspicious = [
+                k
+                for k in reader.fieldnames
+                if k
+                and (
+                    k.lstrip("-").replace(".", "", 1).isdigit()
+                    or k.startswith(("[", "{"))
+                )
+            ]
+            if suspicious:
+                print(
+                    f"Warning: CSV column names look like data values: {suspicious}. "
+                    f"Ensure the first row is a header row with property names.",
+                    file=sys.stderr,
+                )
+
+        for row in reader:
+            # Apply mapping if provided
+            if mapping:
+                row = {mapping.get(k, k): v for k, v in row.items()}
+            yield row
+
+
+def read_json(
+    file_path: Path, mapping: dict[str, str] | None = None
+) -> list[dict[str, Any]]:
+    """
+    Read data from JSON file (expects array of objects).
+
+    Args:
+        file_path: Path to JSON file
+        mapping: Optional key name mapping
+
+    Returns:
+        List of dictionaries with data
+
+    Raises:
+        ValueError: If JSON is not an array
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        data = json.load(f)
+
+    if not isinstance(data, list):
+        raise ValueError(
+            f"JSON file must contain an array of objects, got {type(data).__name__}"
+        )
+
+    # Apply mapping if provided
+    if mapping:
+        data = [{mapping.get(k, k): v for k, v in obj.items()} for obj in data]
+
+    return data
+
+
+def read_jsonl(
+    file_path: Path, mapping: dict[str, str] | None = None
+) -> Iterator[dict[str, Any]]:
+    """
+    Read data from JSONL file (one JSON object per line).
+
+    Yields objects one at a time — suitable for large files.
+
+    Args:
+        file_path: Path to JSONL file
+        mapping: Optional key name mapping
+
+    Yields:
+        Object dictionaries with data
+    """
+    with open(file_path, "r", encoding="utf-8") as f:
+        for line_num, line in enumerate(f, 1):
+            line = line.strip()
+            if not line:
+                continue
+            try:
+                obj = json.loads(line)
+                # Apply mapping if provided
+                if mapping:
+                    obj = {mapping.get(k, k): v for k, v in obj.items()}
+                yield obj
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON on line {line_num}: {e}")
+
+
+def read_pdf(
+    file_path: Path, image_field: str = "doc_page"
+) -> Iterator[dict[str, Any]]:
+    """
+    Convert each page of a PDF to a base64-encoded JPEG and yield as objects.
+
+    Each page becomes one Weaviate object with the base64 image stored in
+    `image_field`, plus `page_number` and `file_name` metadata properties.
+    Page images are freed from memory after encoding.
+
+    Args:
+        file_path: Path to the PDF file
+        image_field: Name of the BLOB property to store the base64 image
+
+    Yields:
+        Dicts with image_field, page_number, and file_name keys
+
+    Raises:
+        RuntimeError: If poppler is not installed
+    """
+    try:
+        from pdf2image import convert_from_path
+
+        pages = convert_from_path(str(file_path))
+    except Exception as e:
+        if "poppler" in str(e).lower() or "pdftoppm" in str(e).lower():
+            raise RuntimeError(
+                f"Poppler is not installed or not in PATH. "
+                f"Install it with:\n"
+                f"  macOS:         brew install poppler\n"
+                f"  Ubuntu/Debian: sudo apt-get install poppler-utils\n"
+                f"Original error: {e}"
+            )
+        raise
+
+    for page_num, page_img in enumerate(pages, 1):
+        buffer = BytesIO()
+        page_img.save(buffer, format="JPEG")
+        img_base64 = base64.b64encode(buffer.getvalue()).decode("utf-8")
+        yield {
+            image_field: img_base64,
+            "page_number": page_num,
+            "file_name": file_path.stem,
+        }
+
+
+def create_pdf_collection(
+    client: weaviate.WeaviateClient, name: str, image_field: str
+) -> None:
+    """
+    Create a Weaviate collection with the standard multimodal PDF schema.
+
+    Properties: image_field (BLOB), page_number (INT), file_name (TEXT)
+    Vectorizer: multi2vec_weaviate with ModernVBERT/colmodernvbert + MUVERA encoding
+
+    Args:
+        client: Connected Weaviate client
+        name: Collection name
+        image_field: Name of the BLOB property to store base64 page images
+    """
+    client.collections.create(
+        name=name,
+        properties=[
+            Property(name=image_field, data_type=DataType.BLOB),
+            Property(name="page_number", data_type=DataType.INT),
+            Property(name="file_name", data_type=DataType.TEXT),
+        ],
+        vector_config=[
+            Configure.MultiVectors.multi2vec_weaviate(
+                name="doc_vector",
+                image_field=image_field,
+                model="ModernVBERT/colmodernvbert",
+                encoding=Configure.VectorIndex.MultiVector.Encoding.muvera(
+                    ksim=4,
+                    dprojections=16,
+                    repetitions=20,
+                ),
+            )
+        ],
+    )
+
+
+def convert_types(
+    obj: dict[str, Any],
+    prop_types: dict[str, DataType],
+) -> dict[str, Any]:
+    """
+    Prepare an object for insertion using the collection schema to guide conversion.
+
+    Non-string values (JSON/JSONL native types) pass through unchanged. String values
+    are cast to the type declared in prop_types. Fields not in the schema pass through
+    as-is. Reserved fields always pass through unchanged.
+
+    Args:
+        obj: Raw object from the file
+        prop_types: Map of property name → DataType from the collection schema
+
+    Returns:
+        Object ready for batch insertion
+    """
+    result = {}
+    for key, value in obj.items():
+        if value is None or value == "":
+            continue
+
+        # Reserved fields pass through as-is (will be dropped or renamed by caller)
+        if key in _RESERVED_FIELDS:
+            result[key] = value
+            continue
+
+        # String value: cast based on schema
+        target_type = prop_types.get(key)
+
+        # Non-string values already have the right native type, with one exception:
+        # date[] lists from JSON/JSONL may contain bare date strings needing RFC3339
+        if not isinstance(value, str):
+            if target_type == DataType.DATE_ARRAY and isinstance(value, list):
+                result[key] = [
+                    f"{d}T00:00:00Z"
+                    if isinstance(d, str) and _DATE_RE.match(d)
+                    else d.replace(" ", "T") + "Z"
+                    if isinstance(d, str) and _DATETIME_RE.match(d)
+                    else d
+                    for d in value
+                ]
+            else:
+                result[key] = value
+            continue
+
+        if target_type == DataType.INT:
+            try:
+                result[key] = int(value)
+            except (ValueError, TypeError):
+                result[key] = value
+        elif target_type == DataType.INT_ARRAY:
+            try:
+                result[key] = [int(x) for x in json.loads(value)]
+            except (ValueError, TypeError):
+                result[key] = value
+        elif target_type == DataType.NUMBER:
+            try:
+                result[key] = float(value)
+            except (ValueError, TypeError):
+                result[key] = value
+        elif target_type == DataType.NUMBER_ARRAY:
+            try:
+                result[key] = [float(x) for x in json.loads(value)]
+            except (ValueError, TypeError):
+                result[key] = value
+        elif target_type == DataType.BOOL:
+            if value.lower() in ("true", "false"):
+                result[key] = value.lower() == "true"
+            else:
+                result[key] = value
+        elif target_type == DataType.BOOL_ARRAY:
+            try:
+                parsed = json.loads(value)
+                result[key] = [
+                    b if isinstance(b, bool) else str(b).lower() == "true"
+                    for b in parsed
+                ]
+            except (ValueError, TypeError):
+                result[key] = value
+        elif target_type == DataType.DATE:
+            if _DATE_RE.match(value):
+                result[key] = f"{value}T00:00:00Z"
+            elif _DATETIME_RE.match(value):
+                result[key] = value.replace(" ", "T") + "Z"
+            else:
+                result[key] = value
+        elif target_type == DataType.DATE_ARRAY:
+            try:
+                parsed = json.loads(value)
+                result[key] = [
+                    f"{d}T00:00:00Z"
+                    if isinstance(d, str) and _DATE_RE.match(d)
+                    else d.replace(" ", "T") + "Z"
+                    if isinstance(d, str) and _DATETIME_RE.match(d)
+                    else d
+                    for d in parsed
+                ]
+            except (ValueError, TypeError):
+                result[key] = value
+        elif target_type is not None and target_type not in _KEEP_AS_STRING:
+            try:
+                result[key] = json.loads(value)
+            except (ValueError, TypeError):
+                result[key] = value
+        else:
+            # text, uuid, blob, or field not in schema — keep as string
+            result[key] = value
+
+    return result
+
+
+def import_objects(
+    coll: Any,
+    data: Iterator[dict[str, Any]],
+    prop_types: dict[str, DataType],
+    skip_set: set[str],
+    batch_size: int,
+) -> tuple[int, int, int, list[str]]:
+    """
+    Batch-insert objects from *data* into *coll*.
+
+    Returns:
+        (total_count, imported_count, failed_count, errors)
+    """
+    total_count = 0
+    imported_count = 0
+    failed_count = 0
+    errors: list[str] = []
+
+    with coll.batch.dynamic() as batch:
+        for i, obj in enumerate(data, 1):
+            total_count += 1
+            try:
+                converted_obj = convert_types(obj, prop_types)
+                if skip_set:
+                    converted_obj = {
+                        k: v for k, v in converted_obj.items() if k not in skip_set
+                    }
+                batch.add_object(properties=converted_obj)
+                imported_count += 1
+
+                if i % batch_size == 0:
+                    print(f"Progress: {i} objects processed", file=sys.stderr)
+
+            except Exception as e:
+                failed_count += 1
+                error_msg = f"Object {i}: {str(e)}"
+                if len(errors) < 10:
+                    errors.append(error_msg)
+                    if len(errors) <= 5:
+                        print(f"Warning: {error_msg}", file=sys.stderr)
+
+    # Check for server-side failures
+    server_failed = 0
+    for failed_obj in coll.batch.failed_objects:
+        server_failed += 1
+        if len(errors) < 10:
+            errors.append(f"Batch error: {failed_obj.message}")
+
+    failed_count += server_failed
+    return total_count, imported_count, failed_count, errors
+
+
+@app.command()
+def main(
+    files: list[str] = typer.Argument(
+        ..., help="One or more CSV, JSON, JSONL, or PDF files"
+    ),
+    collection: str = typer.Option(
+        ..., "--collection", "-c", help="Target collection name"
+    ),
+    mapping: str = typer.Option(
+        None,
+        "--mapping",
+        "-m",
+        help="JSON object mapping file columns/keys to properties",
+    ),
+    tenant: str = typer.Option(
+        None, "--tenant", "-t", help="Tenant name for multi-tenant collections"
+    ),
+    batch_size: int = typer.Option(
+        100, "--batch-size", "-b", help="Number of objects per batch"
+    ),
+    image_field: str = typer.Option(
+        "doc_page",
+        "--image-field",
+        "-i",
+        help="BLOB property name to store base64 page images (PDF imports only)",
+    ),
+    skip_fields: str = typer.Option(
+        None,
+        "--skip-fields",
+        help="Comma-separated field names to exclude from import (e.g. 'id,created_at')",
+    ),
+    json_output: bool = typer.Option(False, "--json", help="Output in JSON format"),
+):
+    """Import data from CSV, JSON, JSONL, or PDF files to a Weaviate collection."""
+    try:
+        # Validate all file paths up front
+        file_paths: list[Path] = []
+        for f in files:
+            fp = Path(f)
+            if not fp.exists():
+                print(f"Error: File not found: {f}", file=sys.stderr)
+                raise typer.Exit(1)
+            file_paths.append(fp)
+
+        # Parse mapping if provided
+        mapping_dict = None
+        if mapping:
+            try:
+                mapping_dict = json.loads(mapping)
+                if not isinstance(mapping_dict, dict):
+                    raise ValueError("Mapping must be a JSON object")
+            except json.JSONDecodeError as e:
+                print(f"Error: Invalid JSON in mapping: {e}", file=sys.stderr)
+                raise typer.Exit(1)
+
+        # Parse skip_fields
+        skip_set: set[str] = (
+            {f.strip() for f in skip_fields.split(",")} if skip_fields else set()
+        )
+
+        # Validate batch size
+        if batch_size < 1:
+            print("Error: Batch size must be at least 1", file=sys.stderr)
+            raise typer.Exit(1)
+
+        # Detect formats. CSV/JSON/JSONL can be mixed freely; PDF cannot be mixed with them.
+        try:
+            fmt_by_path = {fp: detect_file_format(fp) for fp in file_paths}
+        except ValueError as e:
+            print(f"Error: {e}", file=sys.stderr)
+            raise typer.Exit(1)
+
+        has_pdf = any(f == "pdf" for f in fmt_by_path.values())
+        has_non_pdf = any(f != "pdf" for f in fmt_by_path.values())
+        if has_pdf and has_non_pdf:
+            print(
+                "Error: PDF files cannot be mixed with CSV/JSON/JSONL files. "
+                "Import PDFs separately.",
+                file=sys.stderr,
+            )
+            raise typer.Exit(1)
+
+        # Connect to Weaviate once for all files
+        with get_client() as client:
+            # PDF: create collection if absent, append if it exists. CSV/JSON/JSONL: must already exist.
+            if has_pdf:
+                if not client.collections.exists(collection):
+                    print(
+                        f"Creating collection '{collection}' with multimodal PDF schema...",
+                        file=sys.stderr,
+                    )
+                    create_pdf_collection(client, collection, image_field)
+                    print(f"Collection '{collection}' created.", file=sys.stderr)
+                else:
+                    print(
+                        f"Collection '{collection}' exists — appending pages to it.",
+                        file=sys.stderr,
+                    )
+            else:
+                if not client.collections.exists(collection):
+                    print(
+                        f"Error: Collection '{collection}' does not exist. "
+                        f"Read `weaviate` skill's `create_collection.md` reference to create it first.",
+                        file=sys.stderr,
+                    )
+                    raise typer.Exit(1)
+
+            # Fetch schema once — used for multi-tenancy check and type-safe coercion
+            coll = client.collections.get(collection)
+            config = coll.config.get()
+            prop_types: dict[str, DataType] = {
+                p.name: p.data_type for p in config.properties
+            }
+            is_multi_tenant = (
+                config.multi_tenancy_config.enabled
+                if config.multi_tenancy_config
+                else False
+            )
+
+            # Validate tenant parameter
+            if is_multi_tenant and not tenant:
+                print(
+                    f"Error: Collection '{collection}' is multi-tenant, "
+                    f"--tenant parameter is required",
+                    file=sys.stderr,
+                )
+                raise typer.Exit(1)
+            elif not is_multi_tenant and tenant:
+                print(
+                    f"Warning: Collection '{collection}' is not multi-tenant, "
+                    f"--tenant parameter will be ignored",
+                    file=sys.stderr,
+                )
+                tenant = None
+
+            if tenant:
+                coll = coll.with_tenant(tenant)
+                print(f"Using tenant: {tenant}", file=sys.stderr)
+
+            # Process each file
+            grand_total = grand_imported = grand_failed = 0
+            all_errors: list[str] = []
+            file_results = []
+
+            for file_path in file_paths:
+                file_fmt = fmt_by_path[file_path]
+                print(
+                    f"\n[{file_fmt.upper()}] {file_path}",
+                    file=sys.stderr,
+                )
+
+                try:
+                    if file_fmt == "csv":
+                        data: Iterator[dict[str, Any]] = read_csv(
+                            file_path, mapping_dict
+                        )
+                    elif file_fmt == "json":
+                        data = iter(read_json(file_path, mapping_dict))
+                    elif file_fmt == "jsonl":
+                        data = read_jsonl(file_path, mapping_dict)
+                    elif file_fmt == "pdf":
+                        if mapping_dict:
+                            print(
+                                "Warning: --mapping is not supported for PDF imports and will be ignored.",
+                                file=sys.stderr,
+                            )
+                        data = read_pdf(file_path, image_field)
+                except Exception as e:
+                    print(f"Error reading file: {e}", file=sys.stderr)
+                    raise typer.Exit(1)
+
+                # Peek: validate non-empty and warn on reserved fields
+                first = next(data, None)
+                if first is None:
+                    print(
+                        f"Warning: No data found in {file_path}, skipping.",
+                        file=sys.stderr,
+                    )
+                    continue
+                if file_fmt != "pdf":
+                    reserved_found = (set(first.keys()) & _RESERVED_FIELDS) - skip_set
+                    if reserved_found:
+                        print(
+                            f"Warning: Reserved Weaviate field(s) detected in data: "
+                            f"{', '.join(sorted(reserved_found))}. "
+                            f"These will cause import failures. "
+                            f"Use --skip-fields to exclude or --mapping to rename them.",
+                            file=sys.stderr,
+                        )
+                data = itertools.chain([first], data)
+
+                print(
+                    f"Importing objects in batches of {batch_size}...", file=sys.stderr
+                )
+                total, imported, failed, errors = import_objects(
+                    coll, data, prop_types, skip_set, batch_size
+                )
+                grand_total += total
+                grand_imported += imported
+                grand_failed += failed
+                all_errors.extend(errors)
+                file_results.append(
+                    {
+                        "file": str(file_path),
+                        "format": file_fmt,
+                        "total_objects": total,
+                        "imported": total - failed,
+                        "failed": failed,
+                        **({"errors": errors[:10]} if errors else {}),
+                    }
+                )
+
+            grand_success = grand_imported - grand_failed
+
+            result = {
+                "collection": collection,
+                "tenant": tenant,
+                "total_objects": grand_total,
+                "imported": grand_success,
+                "failed": grand_failed,
+                "files": file_results,
+            }
+            if has_pdf:
+                result["image_field"] = image_field
+            if all_errors:
+                result["errors"] = all_errors[:10]
+
+            if json_output:
+                print(json.dumps(result, indent=2))
+            else:
+                print(f"\n✓ Import completed!", file=sys.stderr)
+                print(f"\n**Collection:** {collection}")
+                if tenant:
+                    print(f"**Tenant:** {tenant}")
+                if len(file_paths) > 1:
+                    print(f"**Files Processed:** {len(file_results)}")
+                print(f"**Total Objects:** {grand_total}")
+                print(f"**Successfully Imported:** {grand_success}")
+                if grand_failed > 0:
+                    print(f"**Failed:** {grand_failed}")
+                    if all_errors:
+                        print(f"\n**Sample Errors:**")
+                        for error in all_errors[:5]:
+                            print(f"  - {error}")
+
+            if grand_failed > 0:
+                raise typer.Exit(1)
+
+    except weaviate.exceptions.WeaviateConnectionError as e:
+        print(f"Error: Connection failed - {e}", file=sys.stderr)
+        raise typer.Exit(1)
+    except Exception as e:
+        print(f"Error: {e}", file=sys.stderr)
+        raise typer.Exit(1)
+
+
+if __name__ == "__main__":
+    app()