Vector Database Upserter for Pinecone

Vector Database Upserter for Pinecone

Problem Statement

Modern AI applications — semantic search, RAG pipelines, recommendation engines — need to store and query high-dimensional embeddings quickly. Pinecone provides a managed vector database, but getting data into it reliably is non-trivial: embeddings must be generated with the same model that will be used for queries, metadata must be attached for filtering, and upserts need to be batched to stay within API rate limits and payload size caps. Doing this ad-hoc leads to half-populated indexes, dimension mismatches, and silent failures that only surface when a user's query returns garbage.

This script solves the ingestion side of that problem. It reads a newline-delimited JSON (JSONL) corpus, generates embeddings via OpenAI in batches, and upserts them into a Pinecone index with deterministic IDs and attached metadata. It is idempotent (re-running the same corpus overwrites rather than duplicates), resumable in spirit (deterministic IDs mean partial runs can be safely repeated), and safe under rate limits (exponential backoff on transient errors).

How It Works

The script follows a straightforward ETL pipeline: Extract documents from a JSONL file, Transform each document's text into an embedding vector using OpenAI's text-embedding-3-small model, and Load the vectors plus metadata into a Pinecone index. The key design decisions are:

1. Deterministic IDs — Each document's ID is derived from a SHA-256 hash of its id field (or its text if no ID is supplied). This makes re-runs idempotent and prevents duplicate vectors.
2. Batched embedding calls — OpenAI's embeddings endpoint accepts up to 2048 inputs per request. We batch to reduce round-trips and stay well under the limit.
3. Batched upserts with backoff — Pinecone recommends batches of ~100 vectors. We wrap each upsert in a retry loop with exponential backoff for transient 5xx errors.
4. Dimension validation — Before upserting, we compare the embedding dimension against the index's configured dimension and fail fast if they mismatch.

The script uses a small Config dataclass to centralize environment variables, and a PineconeUpserter class to encapsulate the pipeline. This makes the code testable and easy to extend.

Full Script

#!/usr/bin/env python3
"""
Vector Database Upserter for Pinecone.

Reads a JSONL corpus, generates OpenAI embeddings for each document's text,
and upserts the resulting vectors plus metadata into a Pinecone index.

Typical usage:
    python3 pinecone_upsert.py --index my-index --docs corpus.jsonl
"""

from __future__ import annotations

import argparse
import hashlib
import json
import logging
import os
import sys
import time
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Iterable, Iterator

from openai import OpenAI, APIError, RateLimitError
from pinecone import Pinecone, ServerlessSpec
from pinecone.exceptions import PineconeApiException

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------


@dataclass(frozen=True)
class Config:
    """Runtime configuration loaded from environment variables."""

    openai_api_key: str
    pinecone_api_key: str
    embedding_model: str = "text-embedding-3-small"
    embedding_dim: int = 1536
    embed_batch_size: int = 100
    upsert_batch_size: int = 100
    max_retries: int = 5
    backoff_base: float = 2.0
    cloud: str = "aws"
    region: str = "us-east-1"

    @classmethod
    def from_env(cls) -> "Config":
        openai_key = os.environ.get("OPENAI_API_KEY")
        pinecone_key = os.environ.get("PINECONE_API_KEY")
        if not openai_key:
            raise EnvironmentError("OPENAI_API_KEY is not set.")
        if not pinecone_key:
            raise EnvironmentError("PINECONE_API_KEY is not set.")
        return cls(
            openai_api_key=openai_key,
            pinecone_api_key=pinecone_key,
            embedding_model=os.environ.get(
                "EMBEDDING_MODEL", "text-embedding-3-small"
            ),
            embedding_dim=int(os.environ.get("EMBEDDING_DIM", "1536")),
            embed_batch_size=int(os.environ.get("EMBED_BATCH_SIZE", "100")),
            upsert_batch_size=int(os.environ.get("UPSERT_BATCH_SIZE", "100")),
            max_retries=int(os.environ.get("MAX_RETRIES", "5")),
            backoff_base=float(os.environ.get("BACKOFF_BASE", "2.0")),
            cloud=os.environ.get("PINECONE_CLOUD", "aws"),
            region=os.environ.get("PINECONE_REGION", "us-east-1"),
        )


# ---------------------------------------------------------------------------
# Document loading
# ---------------------------------------------------------------------------


def load_documents(path: Path) -> list[dict[str, Any]]:
    """
    Load a JSONL file into a list of dicts.

    Each line must be a JSON object with a 'text' field (str). An optional
    'id' field (str) is used for deterministic vector IDs; if absent, the
    SHA-256 hash of 'text' is used. Any remaining keys become metadata.
    """
    if not path.exists():
        raise FileNotFoundError(f"Corpus file not found: {path}")

    docs: list[dict[str, Any]] = []
    with path.open("r", encoding="utf-8") as fh:
        for lineno, raw in enumerate(fh, start=1):
            raw = raw.strip()
            if not raw:
                continue
            try:
                obj = json.loads(raw)
            except json.JSONDecodeError as exc:
                raise ValueError(f"Invalid JSON on line {lineno}: {exc}") from exc

            if not isinstance(obj, dict) or "text" not in obj:
                raise ValueError(
                    f"Line {lineno}: expected object with a 'text' field."
                )
            if not isinstance(obj["text"], str) or not obj["text"].strip():
                raise ValueError(f"Line {lineno}: 'text' must be a non-empty string.")

            docs.append(obj)

    if not docs:
        raise ValueError(f"No documents found in {path}")
    return docs


def _stable_id(doc: dict[str, Any]) -> str:
    """Derive a deterministic vector ID from the document."""
    seed = doc.get("id") or doc["text"]
    return hashlib.sha256(seed.encode("utf-8")).hexdigest()[:32]


def _metadata(doc: dict[str, Any]) -> dict[str, Any]:
    """Extract metadata (everything except 'text' and 'id')."""
    return {k: v for k, v in doc.items() if k not in {"text", "id"}}


# ---------------------------------------------------------------------------
# Batching helpers
# ---------------------------------------------------------------------------


def _chunks(items: list[Any], size: int) -> Iterator[list[Any]]:
    """Yield successive chunks of `size` from `items`."""
    for i in range(0, len(items), size):
        yield items[i : i + size]


# ---------------------------------------------------------------------------
# Upserter
# ---------------------------------------------------------------------------


class PineconeUpserter:
    """End-to-end pipeline: JSONL -> embeddings -> Pinecone."""

    def __init__(self, cfg: Config) -> None:
        self.cfg = cfg
        self.openai = OpenAI(api_key=cfg.openai_api_key)
        self.pc = Pinecone(api_key=cfg.pinecone_api_key)

    # -- index management ---------------------------------------------------

    def ensure_index(self, name: str) -> None:
        """Create the index if it does not already exist."""
        existing = {idx.name for idx in self.pc.list_indexes()}
        if name in existing:
            logging.info("Index '%s' already exists.", name)
            return

        logging.info(
            "Creating index '%s' (dim=%d, metric=cosine).",
            name,
            self.cfg.embedding_dim,
        )
        self.pc.create_index(
            name=name,
            dimension=self.cfg.embedding_dim,
            metric="cosine",
            spec=ServerlessSpec(cloud=self.cfg.cloud, region=self.cfg.region),
        )
        # Wait for the index to become ready before returning.
        while not self.pc.describe_index(name).status["ready"]:
            time.sleep(2)

    # -- embeddings ---------------------------------------------------------

    def embed_batch(self, texts: list[str]) -> list[list[float]]:
        """Embed a batch of texts with retry/backoff."""
        for attempt in range(1, self.cfg.max_retries + 1):
            try:
                resp = self.openai.embeddings.create(
                    model=self.cfg.embedding_model,
                    input=texts,
                )
                return [item.embedding for item in resp.data]
            except RateLimitError:
                wait = self.cfg.backoff_base ** attempt
                logging.warning(
                    "OpenAI rate limit hit (attempt %d/%d). Sleeping %.1fs.",
                    attempt,
                    self.cfg.max_retries,
                    wait,
                )
                time.sleep(wait)
            except APIError as exc:
                if attempt == self.cfg.max_retries:
                    raise
                wait = self.cfg.backoff_base ** attempt
                logging.warning(
                    "OpenAI API error: %s. Retrying in %.1fs.", exc, wait
                )
                time.sleep(wait)
        raise RuntimeError("Exhausted retries while embedding batch.")

    # -- upsert -------------------------------------------------------------

    def upsert_vectors(self, index_name: str, vectors: list[dict[str, Any]]) -> int:
        """Upsert vectors in batches with retry/backoff. Returns count upserted."""
        index = self.pc.Index(index_name)
        total = 0
        for batch in _chunks(vectors, self.cfg.upsert_batch_size):
            for attempt in range(1, self.cfg.max_retries + 1):
                try:
                    index.upsert(vectors=batch)
                    total += len(batch)
                    break
                except PineconeApiException as exc:
                    if attempt == self.cfg.max_retries:
                        raise
                    wait = self.cfg.backoff_base ** attempt
                    logging.warning(
                        "Pinecone upsert failed (%s). Retry %d/%d in %.1fs.",
                        exc,
                        attempt,
                        self.cfg.max_retries,
                        wait,
                    )
                    time.sleep(wait)
        return total

    # -- main pipeline ------------------------------------------------------

    def run(self, index_name: str, docs: Iterable[dict[str, Any]]) -> int:
        """Full pipeline. Returns the number of vectors upserted."""
        docs = list(docs)
        logging.info("Loaded %d documents.", len(docs))

        self.ensure_index(index_name)

        all_vectors: list[dict[str, Any]] = []
        for batch in _chunks(docs, self.cfg.embed_batch_size):
            texts = [d["text"] for d in batch]
            embeddings = self.embed_batch(texts)

            # Validate the embedding dimension matches the index.
            dim = len(embeddings[0])
            if dim != self.cfg.embedding_dim:
                raise ValueError(
                    f"Embedding dimension mismatch: got {dim}, "
                    f"expected {self.cfg.embedding_dim}. "
                    "Check EMBEDDING_MODEL / EMBEDDING_DIM."
                )

            for doc, vec in zip(batch, embeddings):
                all_vectors.append(
                    {
                        "id": _stable_id(doc),
                        "values": vec,
                        "metadata": _metadata(doc),
                    }
                )
            logging.info("Embedded %d/%d documents.", len(all_vectors), len(docs))

        upserted = self.upsert_vectors(index_name, all_vectors)
        logging.info("Upserted %d vectors into '%s'.", upserted, index_name)
        return upserted


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Upsert JSONL documents into a Pinecone index."
    )
    parser.add_argument(
        "--index", required=True, help="Name of the Pinecone index."
    )
    parser.add_argument(
        "--docs",
        required=True,
        type=Path,
        help="Path to a JSONL corpus file.",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Parse and embed but do not upsert.",
    )
    parser.add_argument(
        "--verbose", "-v", action="store_true", help="Enable debug logging."
    )
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    logging.basicConfig(
        level=logging.DEBUG if args.verbose else logging.INFO,
        format="%(asctime)s [%(levelname)s] %(message)s",
    )

    try:
        cfg = Config.from_env()
        docs = load_documents(args.docs)

        if args.dry_run:
            logging.info("Dry run: %d documents parsed successfully.", len(docs))
            return 0

        upserter = PineconeUpserter(cfg)
        count = upserter.run(args.index, docs)
        logging.info("Done. %d vectors upserted.", count)
        return 0
    except EnvironmentError as exc:
        logging.error("Configuration error: %s", exc)
        return 2
    except FileNotFoundError as exc:
        logging.error("File error: %s", exc)
        return 3
    except ValueError as exc:
        logging.error("Data error: %s", exc)
        return 4
    except (APIError, PineconeApiException) as exc:
        logging.error("Remote API error: %s", exc)
        return 5
    except KeyboardInterrupt:
        logging.warning("Interrupted by user.")
        return 130
    except Exception as exc:  # noqa: BLE001 - top-level safety net
        logging.exception("Unexpected failure: %s", exc)
        return 1


if __name__ == "__main__":
    sys.exit(main())

Configuration

The script reads configuration from environment variables, with sensible defaults. Set these before running:

The JSONL corpus must contain one JSON object per line. Each object requires a text field; an optional id field controls the vector ID (otherwise derived from SHA-256 of the text). All other keys become vector metadata — keep them small (strings, numbers, booleans, or short lists), since Pinecone caps metadata at 40 KB per vector.

Example corpus.jsonl:

{"id": "doc-1", "text": "Pinecone is a managed vector database.", "source": "docs", "lang": "en"}
{"id": "doc-2", "text": "OpenAI provides embedding models.", "source": "blog", "lang": "en"}

Usage Examples

1. Basic upsert into a new or existing index:

export OPENAI_API_KEY=sk-...
export PINECONE_API_KEY=pc-...
python3 pinecone_upsert.py --index my-index --docs corpus.jsonl

Expected output:

2025-01-15 10:12:01 [INFO] Loaded 250 documents.
2025-01-15 10:12:02 [INFO] Creating index 'my-index' (dim=1536, metric=cosine).
2025-01-15 10:12:18 [INFO] Embedded 100/250 documents.
2025-01-15 10:12:20 [INFO] Embedded 200/250 documents.
2025-01-15 10:12:22 [INFO] Embedded 250/250 documents.
2025-01-15 10:12:24 [INFO] Upserted 250 vectors into 'my-index'.
2025-01-15 10:12:24 [INFO] Done. 250 vectors upserted.

2. Dry run to validate the corpus without touching Pinecone:

python3 pinecone_upsert.py --index my-index --docs corpus.jsonl --dry-run

Output:

2025-01-15 10:15:00 [INFO] Dry run: 250 documents parsed successfully.

3. Verbose run with a custom embedding model:

EMBEDDING_MODEL=text-embedding-3-large EMBEDDING_DIM=3072 \
  python3 pinecone_upsert.py --index big-index --docs corpus.jsonl -v

Error Handling

• Missing API keys — Config.from_env() raises EnvironmentError; main() catches it and exits with code 2.
• Missing or unreadable corpus file — FileNotFoundError is caught and returns exit code 3.
• Malformed JSON or missing text field — load_documents() raises ValueError with the offending line number; exit code 4.
• Embedding dimension mismatch — Detected before upsert; raises ValueError and exits 4. This prevents polluting an index with wrong-shaped vectors.
• OpenAI rate limits — RateLimitError triggers exponential backoff up to MAX_RETRIES attempts, then propagates.
• OpenAI transient API errors — Retried with backoff; final failure raises APIError, caught in main() with exit code 5.
• Pinecone upsert errors — PineconeApiException is retried per batch; final failure exits with code 5.
• Keyboard interrupt — Caught and exits cleanly with code 130 (standard for SIGINT).
• Unexpected exceptions — Logged with full traceback; exits with code 1.

Because vector IDs are deterministic hashes, a failed run can be safely retried from the beginning without creating duplicates — the same document always lands on the same ID.

Extending Ideas

• Parallel embedding — Replace the sequential embedding loop with concurrent.futures.ThreadPoolExecutor (or asyncio with the OpenAI async client) to overlap network round-trips. Keep batches independent so retries stay isolated.
• Namespacing — Add a --namespace flag and pass it to index.upsert(..., namespace=...). Namespaces let you isolate tenants or document versions within a single index without re-creating it.
• Incremental sync — Track a last_modified timestamp in metadata and skip documents whose content hash has not changed. Combine with index.fetch(ids=[...]) to avoid re-embedding unchanged text.
• Chunking long documents — For documents exceeding the embedding model's token limit (~8k tokens for text-embedding-3-small), split into overlapping chunks, embed each, and store the parent document ID in metadata so a query can retrieve the full source.
• Structured metadata filtering — Add a --filter flag that drops documents whose metadata does not match a predicate, so a single corpus can feed multiple indexes with different selective ingestion rules.