Files
playbook/antigravity-awesome-skills/skills/weaviate/scripts/import.py
T
2026-06-29 16:09:10 +00:00

702 lines
24 KiB
Python
Executable File

#!/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()