Lesson 16.14: Embeddings & Production Indexing Pipelines

[
  -0.008432194590568542,
   0.02498214691877365,
  -0.0156728345900774,
   0.04821908473968506,
  -0.03109482116997242,
  ...  # Exactly 1,536 floating-point values
]

from math import sqrt
from typing import List

def cosine_similarity(a: List[float], b: List[float]) -> float:
    """
    Computes cosine similarity between two float vectors.
    Returns a score ranging from -1.0 to +1.0.
    +1.0 indicates identical orientation; 0.0 indicates orthogonal/unrelated.
    """
    if len(a) != len(b):
        raise ValueError(f"Vector dimension mismatch: {len(a)} vs {len(b)}")
        
    dot_product = sum(x * y for x, y in zip(a, b))
    norm_a = sqrt(sum(x * x for x in a))
    norm_b = sqrt(sum(y * y for y in b))
    
    if norm_a == 0.0 or norm_b == 0.0:
        raise ValueError("Cosine similarity undefined for zero vectors.")
        
    return dot_product / (norm_a * norm_b)

pip install openai python-dotenv tiktoken qdrant-client pytest

import os
from openai import OpenAI
from dotenv import load_dotenv

load_dotenv()

client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

text = "Employees receive 20 days of paid annual leave."
response = client.embeddings.create(
    model="text-embedding-3-small",
    input=text
)

vector = response.data[0].embedding
print("Vector dimensions:", len(vector))
print("First 3 numbers:", vector[:3])
print("Tokens used:", response.usage.total_tokens)

from typing import Iterator, List, Dict, Any
import tiktoken

encoding = tiktoken.get_encoding("cl100k_base")

MAX_INPUT_TOKENS = 7500   # Cap per discrete chunk
MAX_BATCH_TOKENS = 70000  # Cumulative batch token limit
MAX_BATCH_ITEMS = 64      # Maximum chunks per batch API call

def count_tokens(text: str) -> int:
    return len(encoding.encode(text))

def iter_embedding_batches(records: List[Dict[str, Any]]) -> Iterator[List[Dict[str, Any]]]:
    """
    Partitions input records into memory- and token-safe batches.
    Discards empty strings and flags oversized chunks.
    """
    batch = []
    batch_tokens = 0
    
    for record in records:
        text = record["text"].strip()
        if not text:
            continue
            
        token_count = count_tokens(text)
        if token_count > MAX_INPUT_TOKENS:
            raise ValueError(f"Chunk exceeds token limit ({token_count} > {MAX_INPUT_TOKENS}): {record.get('chunk_id')}")
            
        if batch and (len(batch) >= MAX_BATCH_ITEMS or batch_tokens + token_count > MAX_BATCH_TOKENS):
            yield batch
            batch = []
            batch_tokens = 0
            
        batch.append({**record, "text": text, "token_count": token_count})
        batch_tokens += token_count
        
    if batch:
        yield batch

from openai import OpenAI
import os

client = OpenAI(
    api_key=os.getenv("OPENAI_API_KEY"),
    timeout=45.0,
    max_retries=3
)

MODEL_NAME = os.getenv("OPENAI_EMBEDDING_MODEL", "text-embedding-3-small")

def embed_texts(texts: List[str]) -> List[List[float]]:
    """
    Accepts an ordered list of strings and returns vectors matching the input order.
    """
    if not texts:
        return []
        
    response = client.embeddings.create(
        model=MODEL_NAME,
        input=texts
    )
    
    # Sort strictly by API index
    sorted_items = sorted(response.data, key=lambda x: x.index)
    vectors = [item.embedding for item in sorted_items]
    
    if len(vectors) != len(texts):
        raise RuntimeError(f"Input count ({len(texts)}) does not match output vector count ({len(vectors)})!")
        
    return vectors

from qdrant_client import QdrantClient, models

COLLECTION_NAME = "enterprise_knowledge"
VECTOR_SIZE = 1536  # Dimension for text-embedding-3-small

def ensure_collection(qdrant: QdrantClient) -> None:
    """
    Creates collection if absent, or verifies vector configuration compatibility.
    """
    if not qdrant.collection_exists(COLLECTION_NAME):
        qdrant.create_collection(
            collection_name=COLLECTION_NAME,
            vectors_config=models.VectorParams(
                size=VECTOR_SIZE,
                distance=models.Distance.COSINE
            )
        )
        return
        
    info = qdrant.get_collection(COLLECTION_NAME)
    cfg = info.config.params.vectors
    if not isinstance(cfg, models.VectorParams) or cfg.size != VECTOR_SIZE:
        raise ValueError(f"Collection dimensions ({cfg.size}) conflict with model dimensions ({VECTOR_SIZE})!")

from uuid import uuid5, NAMESPACE_URL
from app.embedding_service import embed_texts, iter_embedding_batches

def to_qdrant_uuid(chunk_id: str) -> str:
    """Maps string chunk identifiers to deterministic UUIDv5 values."""
    return str(uuid5(NAMESPACE_URL, chunk_id))

def index_chunks(chunks: List[Dict[str, Any]], qdrant: QdrantClient) -> int:
    """
    Consumes chunk dictionaries, embeds batches, and upserts points with metadata.
    """
    if not chunks:
        return 0
        
    ensure_collection(qdrant)
    total_indexed = 0
    
    for batch in iter_embedding_batches(chunks):
        texts = [item["text"] for item in batch]
        vectors = embed_texts(texts)
        
        points = []
        for item, vector in zip(batch, vectors):
            point_id = to_qdrant_uuid(item["chunk_id"])
            payload = {
                **item.get("metadata", {}),
                "chunk_id": item["chunk_id"],
                "document_id": item["document_id"],
                "tenant_id": item["tenant_id"],
                "text": item["text"],
                "chunking_version": item.get("chunking_version", "v1"),
                "embedding_model": "text-embedding-3-small"
            }
            points.append(
                models.PointStruct(
                    id=point_id,
                    vector=vector,
                    payload=payload
                )
            )
            
        qdrant.upsert(collection_name=COLLECTION_NAME, points=points, wait=True)
        total_indexed += len(points)
        
    return total_indexed

import pytest
from app.similarity import cosine_similarity
from app.indexing import to_qdrant_uuid
from app.embedding_service import iter_embedding_batches

def test_cosine_similarity_identity_and_orthogonality():
    v1 = [1.0, 0.0, 0.0]
    v2 = [1.0, 0.0, 0.0]
    v3 = [0.0, 1.0, 0.0]
    
    assert pytest.approx(cosine_similarity(v1, v2)) == 1.0
    assert pytest.approx(cosine_similarity(v1, v3)) == 0.0

def test_deterministic_qdrant_uuid():
    id_a = "chk_8a4f91e32d0b"
    uuid1 = to_qdrant_uuid(id_a)
    uuid2 = to_qdrant_uuid(id_a)
    assert uuid1 == uuid2
    assert len(uuid1) == 36  # Standard UUID string representation

def test_batch_generator_empty_filtering():
    records = [
        {"chunk_id": "c1", "text": "Valid text string."},
        {"chunk_id": "c2", "text": "   "},
        {"chunk_id": "c3", "text": "Another valid string."}
    ]
    batches = list(iter_embedding_batches(records))
    assert len(batches) == 1
    assert len(batches[0]) == 2
    assert batches[0][0]["chunk_id"] == "c1"
    assert batches[0][1]["chunk_id"] == "c3"