Vector Search at Scale: Building Semantic Search Systems

#!/usr/bin/env python3
"""Embedding generation with batching support."""

from openai import OpenAI
from typing import List, Dict


class EmbeddingGenerator:
    """Generate text embeddings via OpenAI or a drop-in compatible API."""

    def __init__(self, model: str = "text-embedding-3-small"):
        self.client = OpenAI()
        self.model = model

    def embed_text(self, text: str) -> List[float]:
        """Embed a single string."""
        response = self.client.embeddings.create(model=self.model, input=text)
        return response.data[0].embedding

    def embed_batch(self, texts: List[str], batch_size: int = 100) -> List[List[float]]:
        """Embed a list of strings in batches to stay within API limits."""
        all_embeddings = []
        for i in range(0, len(texts), batch_size):
            batch = texts[i : i + batch_size]
            response = self.client.embeddings.create(model=self.model, input=batch)
            all_embeddings.extend(item.embedding for item in response.data)
        return all_embeddings

    def embed_documents(self, documents: List[Dict]) -> List[Dict]:
        """Embed a list of documents, concatenating title + content."""
        texts = [f"{d.get('title', '')} {d.get('content', '')}" for d in documents]
        embeddings = self.embed_batch(texts)
        return [
            {"id": d["id"], "embedding": emb, "text": texts[i], "metadata": {k: v for k, v in d.items() if k not in ("id", "content")}}
            for i, (d, emb) in enumerate(zip(documents, embeddings))
        ]

from sentence_transformers import SentenceTransformer
from typing import List


class OpenSourceEmbedder:
    """Drop-in embedder using sentence-transformers (no API key required)."""

    def __init__(self, model_name: str = "BAAI/bge-base-en-v1.5"):
        self.model = SentenceTransformer(model_name)

    def embed_text(self, text: str) -> List[float]:
        return self.model.encode(text).tolist()

    def embed_batch(self, texts: List[str]) -> List[List[float]]:
        return self.model.encode(texts, batch_size=64, show_progress_bar=True).tolist()

#!/usr/bin/env python3
"""Pinecone vector store operations."""

from pinecone import Pinecone, ServerlessSpec
from typing import List, Dict, Optional


class PineconeVectorStore:
    """Thin wrapper around the Pinecone client."""

    def __init__(self, api_key: str, index_name: str, dimension: int = 1536):
        self.client = Pinecone(api_key=api_key)
        self.index_name = index_name
        self._ensure_index(dimension)
        self.index = self.client.Index(index_name)

    def _ensure_index(self, dimension: int):
        if self.index_name not in self.client.list_indexes().names():
            self.client.create_index(
                name=self.index_name,
                dimension=dimension,
                metric="cosine",
                spec=ServerlessSpec(cloud="aws", region="us-east-1"),
            )

    def upsert(self, vectors: List[Dict]):
        """Upsert a list of {id, embedding, metadata} dicts."""
        self.index.upsert(
            vectors=[{"id": v["id"], "values": v["embedding"], "metadata": v.get("metadata", {})} for v in vectors]
        )

    def search(self, query_vector: List[float], top_k: int = 10, filter_dict: Optional[Dict] = None) -> List[Dict]:
        """Return top-K results, optionally filtered by metadata."""
        results = self.index.query(vector=query_vector, top_k=top_k, filter=filter_dict, include_metadata=True, include_values=False)
        return [{"id": m["id"], "score": m["score"], "metadata": m.get("metadata", {})} for m in results["matches"]]

#!/usr/bin/env python3
"""Milvus vector store with IVF_FLAT index."""

from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType
from typing import List, Dict


class MilvusVectorStore:
    """Milvus collection wrapper."""

    def __init__(self, collection_name: str, dimension: int = 1536, host: str = "localhost", port: int = 19530):
        connections.connect(host=host, port=port)
        self.collection = self._get_or_create(collection_name, dimension)

    def _get_or_create(self, name: str, dimension: int) -> Collection:
        from pymilvus import utility
        if utility.has_collection(name):
            return Collection(name)

        schema = CollectionSchema([
            FieldSchema("id", DataType.INT64, is_primary=True, auto_id=False),
            FieldSchema("vector", DataType.FLOAT_VECTOR, dim=dimension),
            FieldSchema("text", DataType.VARCHAR, max_length=65535),
        ])
        col = Collection(name, schema)
        col.create_index("vector", {"metric_type": "IP", "index_type": "IVF_FLAT", "params": {"nlist": 128}})
        return col

    def insert(self, vectors: List[Dict]):
        self.collection.insert([
            [v["id"] for v in vectors],
            [v["embedding"] for v in vectors],
            [v["text"][:65535] for v in vectors],
        ])
        self.collection.flush()

    def search(self, query_vector: List[float], top_k: int = 10) -> List[Dict]:
        self.collection.load()
        hits = self.collection.search(
            data=[query_vector], anns_field="vector",
            param={"metric_type": "IP", "params": {"nprobe": 10}},
            limit=top_k, output_fields=["text"],
        )
        return [{"id": h.id, "score": h.distance, "text": h.entity.get("text")} for h in hits[0]]

#!/usr/bin/env python3
"""Semantic search with optional cross-encoder reranking."""

from sentence_transformers import CrossEncoder
from typing import List, Dict


class SemanticSearchApp:
    """Combines an embedder and a vector store into a search interface."""

    def __init__(self, embedder, vector_store):
        self.embedder = embedder
        self.store = vector_store
        self._reranker = None

    def index(self, documents: List[Dict]) -> int:
        embedded = self.embedder.embed_documents(documents)
        self.store.upsert(embedded)
        return len(embedded)

    def search(self, query: str, top_k: int = 5) -> List[Dict]:
        return self.store.search(self.embedder.embed_text(query), top_k=top_k)

    def search_with_rerank(self, query: str, candidates: int = 20, top_k: int = 5) -> List[Dict]:
        """Retrieve a wider candidate set, then rerank with a cross-encoder."""
        if self._reranker is None:
            self._reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

        results = self.search(query, top_k=candidates)
        scores = self._reranker.predict([(query, r["text"]) for r in results])

        for r, score in zip(results, scores):
            r["rerank_score"] = float(score)

        return sorted(results, key=lambda x: x["rerank_score"], reverse=True)[:top_k]

class HybridSearchEngine:
    """Combine dense vector search with sparse keyword search."""

    def __init__(self, dense_index, sparse_index, alpha=0.7):
        this.dense = dense_index     # Vector DB (semantic)
        this.sparse = sparse_index   # BM25 index (keyword)
        this.alpha = alpha           # Dense weight (0-1)

    def search(self, query: str, top_k: int = 10) -> list:
        dense_results = this.dense.search(query, top_k)
        sparse_results = this.sparse.search(query, top_k)

        # Reciprocal Rank Fusion (RRF)
        scores = {}
        for rank, r in enumerate(dense_results):
            scores[r["id"]] = this.alpha * (1.0 / (rank + 60))
        for rank, r in enumerate(sparse_results):
            scores[r["id"]] = scores.get(r["id"], 0) + (1 - this.alpha) * (1.0 / (rank + 60))

        return sorted(scores.items(), key=lambda x: -x[1])[:top_k]

Client → API Gateway → Embedding Service (BERT/BGE)
                        ↓
                  Vector DB Cluster
                  ├── Primary (writes)
                  ├── Replica 1 (reads)
                  ├── Replica 2 (reads)
                  └── Replica N (reads)
                        ↓
                  Reranker (Cross-encoder)
                        ↓
                  LLM (for RAG applications)

#!/usr/bin/env python3
"""Fan-out search across multiple vector store shards."""

from typing import List, Dict


def federated_search(query_vector: List[float], shards: list, top_k: int = 10) -> List[Dict]:
    """Query all shards in parallel and return the global top-K."""
    import concurrent.futures

    all_results: List[Dict] = []
    with concurrent.futures.ThreadPoolExecutor() as pool:
        futures = [pool.submit(shard.search, query_vector, top_k) for shard in shards]
        for f in concurrent.futures.as_completed(futures):
            all_results.extend(f.result())

    # Deduplicate by ID and keep highest score
    seen: Dict[str, Dict] = {}
    for r in all_results:
        if r["id"] not in seen or r["score"] > seen[r["id"]]["score"]:
            seen[r["id"]] = r

    return sorted(seen.values(), key=lambda x: x["score"], reverse=True)[:top_k]

import pinecone

pc = pinecone.Pinecone(api_key="your-api-key")

# Create serverless index
pc.create_index(
    name="semantic-search",
    dimension=1024,
    metric="cosine",
    spec=pinecone.ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)

index = pc.Index("semantic-search")

# Upsert vectors
index.upsert(vectors=[
    ("id1", [0.1, 0.2, ...], {"text": "document content", "source": "pdf"}),
    ("id2", [0.3, 0.4, ...], {"text": "another document", "source": "web"}),
])

# Query
results = index.query(
    vector=query_embedding,
    top_k=10,
    filter={"source": {"$eq": "pdf"}},
    include_metadata=True
)

from qdrant_client import QdrantClient
from qdrant_client.http.models import Distance, VectorParams

client = QdrantClient(host="localhost", port=6333)

# Create collection with payload indexes
client.create_collection(
    collection_name="documents",
    vectors_config=VectorParams(size=1024, distance=Distance.COSINE),
)
client.create_payload_index(
    collection_name="documents",
    field_name="category",
    field_type="keyword",
)
client.create_payload_index(
    collection_name="documents",
    field_name="created_at",
    field_type="integer",
)

# Search with filtering
results = client.search(
    collection_name="documents",
    query_vector=query_embedding,
    limit=10,
    query_filter=models.Filter(
        must=[
            models.FieldCondition(key="category", match=models.MatchValue(value="technical")),
            models.FieldCondition(key="created_at", range=models.Range(gte=1700000000)),
        ]
    ),
    with_payload=True,
)

from sentence_transformers import CrossEncoder

class Reranker:
    """Cross-encoder reranker for improving search relevance."""

    def __init__(self, model_name="cross-encoder/ms-marco-MiniLM-L-6-v2"):
        this.model = CrossEncoder(model_name)

    def rerank(self, query: str, candidates: list[dict], top_k: int = 10) -> list[dict]:
        pairs = [(query, c["text"]) for c in candidates]
        scores = this.model.predict(pairs)
        for c, s in zip(candidates, scores):
            c["rerank_score"] = float(s)
        return sorted(candidates, key=lambda x: x["rerank_score"], reverse=True)[:top_k]

Stage 1: Vector search → 100 candidates (fast, approximate)
Stage 2: Reranker → 20 candidates (accurate, pairwise cross-encoder)
Stage 3: LLM rerank → top 5 (context-aware ranking)

# Pinecone: List indexes
curl -s https://api.pinecone.io/indexes -H "Api-Key: $PINECONE_API_KEY" | jq .

# Qdrant: Create collection
curl -X PUT http://localhost:6333/collections/my_collection \
  -H 'Content-Type: application/json' \
  -d '{"vectors": {"size": 1024, "distance": "Cosine"}}'

# Milvus: Load collection (required before search)
python3 -c "from pymilvus import Collection; Collection('my_collection').load()"

# Weaviate: Search with hybrid
curl -X POST http://localhost:8080/v1/graphql \
  -H 'Content-Type: application/json' \
  -d '{"query": "{ Get { Document(hybrid: {query: \"search text\"}) { title score } } }"}'

def estimate_vector_search_cost(
    num_vectors: int,
    dimension: int = 1024,
    queries_per_month: int = 1000000,
    managed: bool = True
) -> dict:
    """Estimate monthly cost for vector search infrastructure."""
    bytes_per_vector = dimension * 4  # FP32
    raw_storage_gb = num_vectors * bytes_per_vector / (1024**3)
    indexed_storage_gb = raw_storage_gb * 2.5  # HNSW overhead

    if managed:
        storage_cost = indexed_storage_gb * 0.50 * 30  # $0.50/GB/month
        query_cost = queries_per_month * 0.000005  # $5 per 1M queries
        total = storage_cost + query_cost
    else:
        # Self-hosted: 3x A100 GPUs for 10M+ vectors
        infra_cost = 3 * 3.50 * 24 * 30  # $7,560/month
        total = infra_cost

    return {
        "vectors": f"{num_vectors:,}",
        "raw_storage_gb": round(raw_storage_gb, 1),
        "indexed_storage_gb": round(indexed_storage_gb, 1),
        "monthly_cost": f"${total:,.0f}",
        "cost_per_query": f"${total / queries_per_month:.6f}"
    }

# Example: 10M vectors, 1024d, managed
print(estimate_vector_search_cost(10_000_000))

Documents → Chunk → Embed → Index → Query → Rerank → Generate (RAG)
            ^        ^       ^        ^        ^          ^
           512t     Model    DB +   Hybrid   Cross-     LLM with
           chunks   BGE-M3   HNSW   search   encoder    context