from openai import OpenAI

client = OpenAI()  # Uses OPENAI_API_KEY env var

def get_embedding(text: str, model: str = "text-embedding-3-large") -> list[float]:
    """Get embedding for a single text."""
    response = client.embeddings.create(
        input=text,
        model=model
    )
    return response.data[0].embedding

def get_embeddings_batch(
    texts: list[str],
    model: str = "text-embedding-3-large",
    batch_size: int = 2048  # OpenAI limit
) -> list[list[float]]:
    """Get embeddings for multiple texts efficiently."""
    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        response = client.embeddings.create(
            input=batch,
            model=model
        )
        # Preserve order (API may return out of order)
        sorted_embeddings = sorted(response.data, key=lambda x: x.index)
        all_embeddings.extend([e.embedding for e in sorted_embeddings])

    return all_embeddings

def get_embedding_reduced(
    text: str,
    model: str = "text-embedding-3-large",
    dimensions: int = 256  # Reduce from 3072 to 256
) -> list[float]:
    """Get reduced-dimension embedding (cheaper storage, similar quality)."""
    response = client.embeddings.create(
        input=text,
        model=model,
        dimensions=dimensions  # 256, 512, 1024, 1536, 3072
    )
    return response.data[0].embedding

from mistralai import Mistral

client = Mistral(api_key="YOUR_API_KEY")

def get_mistral_embedding(text: str) -> list[float]:
    """Get embedding using Mistral."""
    response = client.embeddings.create(
        model="mistral-embed",
        inputs=[text]
    )
    return response.data[0].embedding

# Batch processing
def get_mistral_embeddings_batch(texts: list[str]) -> list[list[float]]:
    """Batch embeddings with Mistral."""
    response = client.embeddings.create(
        model="mistral-embed",
        inputs=texts
    )
    return [e.embedding for e in response.data]

import cohere

co = cohere.Client("YOUR_API_KEY")

def get_cohere_embedding(
    texts: list[str],
    input_type: str = "search_document"  # or "search_query"
) -> list[list[float]]:
    """
    Get Cohere embeddings.

    input_type options:
    - "search_document": For documents to be searched
    - "search_query": For search queries
    - "classification": For classification tasks
    - "clustering": For clustering tasks
    """
    response = co.embed(
        texts=texts,
        model="embed-english-v3.0",
        input_type=input_type,
        truncate="END"  # or "START", "NONE"
    )
    return response.embeddings

pip install sentence-transformers
# For GPU support
pip install sentence-transformers[gpu]

from sentence_transformers import SentenceTransformer

# Load model (downloads on first use)
model = SentenceTransformer("BAAI/bge-large-en-v1.5")

# Single text
embedding = model.encode("Your text here")

# Batch processing (automatic batching)
texts = ["Text 1", "Text 2", "Text 3"]
embeddings = model.encode(texts, show_progress_bar=True)

import torch
from sentence_transformers import SentenceTransformer

# Auto-detect GPU
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer("BAAI/bge-large-en-v1.5", device=device)

# Encode with GPU
embeddings = model.encode(
    texts,
    batch_size=32,  # Adjust based on GPU memory
    show_progress_bar=True,
    convert_to_numpy=True,  # Return numpy array
    normalize_embeddings=True  # L2 normalize for cosine similarity
)

# Best quality (slow)
model = SentenceTransformer("BAAI/bge-m3")

# Good quality, balanced
model = SentenceTransformer("BAAI/bge-large-en-v1.5")
model = SentenceTransformer("intfloat/e5-large-v2")

# Fast inference
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Multilingual
model = SentenceTransformer("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# Long context (8K tokens)
model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True)

from FlagEmbedding import BGEM3FlagModel

model = BGEM3FlagModel("BAAI/bge-m3", use_fp16=True)

# Dense embeddings (default)
embeddings = model.encode(texts)["dense_vecs"]

# Sparse embeddings (for hybrid search)
sparse = model.encode(texts, return_sparse=True)["lexical_weights"]

# Both dense + sparse
output = model.encode(texts, return_dense=True, return_sparse=True)

def chunk_fixed_size(
    text: str,
    chunk_size: int = 500,
    overlap: int = 50
) -> list[str]:
    """Split text into fixed-size chunks with overlap."""
    words = text.split()
    chunks = []
    start = 0

    while start < len(words):
        end = start + chunk_size
        chunk = " ".join(words[start:end])
        chunks.append(chunk)
        start = end - overlap

    return chunks

import tiktoken

def chunk_by_tokens(
    text: str,
    max_tokens: int = 500,
    overlap_tokens: int = 50,
    model: str = "text-embedding-3-large"
) -> list[str]:
    """Split text by token count (more accurate)."""
    encoding = tiktoken.encoding_for_model(model)
    tokens = encoding.encode(text)
    chunks = []
    start = 0

    while start < len(tokens):
        end = min(start + max_tokens, len(tokens))
        chunk_tokens = tokens[start:end]
        chunk_text = encoding.decode(chunk_tokens)
        chunks.append(chunk_text)
        start = end - overlap_tokens

    return chunks

from langchain.text_splitter import RecursiveCharacterTextSplitter

def chunk_semantic(
    text: str,
    chunk_size: int = 1000,
    chunk_overlap: int = 200
) -> list[str]:
    """Split at natural boundaries (paragraphs, sentences)."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap,
        separators=["\n\n", "\n", ". ", " ", ""]
    )
    return splitter.split_text(text)

import nltk
nltk.download('punkt')

def chunk_by_sentences(
    text: str,
    sentences_per_chunk: int = 5,
    overlap_sentences: int = 1
) -> list[str]:
    """Split by sentences for better semantic coherence."""
    sentences = nltk.sent_tokenize(text)
    chunks = []
    start = 0

    while start < len(sentences):
        end = start + sentences_per_chunk
        chunk = " ".join(sentences[start:end])
        chunks.append(chunk)
        start = end - overlap_sentences

    return chunks

import hashlib
from functools import lru_cache
from typing import Tuple

def text_hash(text: str) -> str:
    """Create consistent hash for text."""
    return hashlib.sha256(text.encode()).hexdigest()[:16]

@lru_cache(maxsize=10000)
def get_embedding_cached(text_hash: str, text: str, model: str) -> Tuple[float, ...]:
    """Cache embeddings in memory."""
    embedding = get_embedding(text, model)
    return tuple(embedding)  # Tuples are hashable

import json
import hashlib
from pathlib import Path

class EmbeddingCache:
    def __init__(self, cache_dir: str = ".embedding_cache"):
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)

    def _get_key(self, text: str, model: str) -> str:
        content = f"{model}:{text}"
        return hashlib.sha256(content.encode()).hexdigest()

    def get(self, text: str, model: str) -> list[float] | None:
        key = self._get_key(text, model)
        cache_file = self.cache_dir / f"{key}.json"

        if cache_file.exists():
            return json.loads(cache_file.read_text())
        return None

    def set(self, text: str, model: str, embedding: list[float]):
        key = self._get_key(text, model)
        cache_file = self.cache_dir / f"{key}.json"
        cache_file.write_text(json.dumps(embedding))

    def get_or_compute(
        self,
        text: str,
        model: str,
        compute_fn
    ) -> list[float]:
        cached = self.get(text, model)
        if cached:
            return cached

        embedding = compute_fn(text, model)
        self.set(text, model, embedding)
        return embedding

import redis
import json
import hashlib

class RedisEmbeddingCache:
    def __init__(self, redis_url: str = "redis://localhost:6379"):
        self.client = redis.from_url(redis_url)
        self.ttl = 86400 * 30  # 30 days

    def _get_key(self, text: str, model: str) -> str:
        content = f"{model}:{text}"
        hash_val = hashlib.sha256(content.encode()).hexdigest()
        return f"emb:{hash_val}"

    def get(self, text: str, model: str) -> list[float] | None:
        key = self._get_key(text, model)
        data = self.client.get(key)
        return json.loads(data) if data else None

    def set(self, text: str, model: str, embedding: list[float]):
        key = self._get_key(text, model)
        self.client.setex(key, self.ttl, json.dumps(embedding))

    def get_batch(
        self,
        texts: list[str],
        model: str
    ) -> dict[str, list[float] | None]:
        """Get multiple embeddings from cache."""
        pipe = self.client.pipeline()
        keys = [self._get_key(t, model) for t in texts]

        for key in keys:
            pipe.get(key)

        results = pipe.execute()
        return {
            text: json.loads(data) if data else None
            for text, data in zip(texts, results)
        }

# Instead of 3072 dimensions ($0.13/1M)
# Use 1024 dimensions with minimal quality loss
embedding = get_embedding_reduced(text, dimensions=1024)
# Saves 66% storage, similar retrieval quality

# Bad: 1000 API calls
for text in texts:
    get_embedding(text)  # 1000 requests

# Good: 1 API call
get_embeddings_batch(texts)  # 1 request for up to 2048 texts

# Cache hit rate of 80% = 80% cost reduction
cache = RedisEmbeddingCache()
embedding = cache.get_or_compute(text, model, get_embedding)

# Two-stage retrieval:
# 1. Fast filter with small model
quick_results = search_with_model(query, "all-MiniLM-L6-v2", top_k=100)

# 2. Rerank with large model
final_results = rerank_with_model(query, quick_results, "text-embedding-3-large", top_k=10)

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity

def deduplicate_texts(texts: list[str], threshold: float = 0.95) -> list[str]:
    """Remove near-duplicate texts before embedding."""
    vectorizer = TfidfVectorizer()
    tfidf_matrix = vectorizer.fit_transform(texts)

    unique_texts = []
    for i, text in enumerate(texts):
        if not unique_texts:
            unique_texts.append(text)
            continue

        # Compare with existing
        current_vec = tfidf_matrix[i]
        existing_vecs = vectorizer.transform(unique_texts)
        similarities = cosine_similarity(current_vec, existing_vecs)[0]

        if max(similarities) < threshold:
            unique_texts.append(text)

    return unique_texts

import numpy as np
from typing import Callable
from sklearn.metrics.pairwise import cosine_similarity

def benchmark_retrieval(
    queries: list[str],
    documents: list[str],
    relevance: dict[int, list[int]],  # query_idx -> [relevant_doc_idxs]
    embed_fn: Callable[[list[str]], np.ndarray],
    k: int = 10
) -> dict[str, float]:
    """
    Benchmark retrieval quality.

    Returns:
    - Recall@K
    - MRR (Mean Reciprocal Rank)
    - Precision@K
    """
    # Embed all
    query_embeddings = embed_fn(queries)
    doc_embeddings = embed_fn(documents)

    # Compute similarities
    similarities = cosine_similarity(query_embeddings, doc_embeddings)

    recalls, mrrs, precisions = [], [], []

    for q_idx, relevant_docs in relevance.items():
        # Get top-k results
        scores = similarities[q_idx]
        top_k_idxs = np.argsort(scores)[::-1][:k]

        # Recall@K
        hits = len(set(top_k_idxs) & set(relevant_docs))
        recalls.append(hits / len(relevant_docs))

        # MRR
        for rank, doc_idx in enumerate(top_k_idxs, 1):
            if doc_idx in relevant_docs:
                mrrs.append(1 / rank)
                break
        else:
            mrrs.append(0)

        # Precision@K
        precisions.append(hits / k)

    return {
        "recall@k": np.mean(recalls),
        "mrr": np.mean(mrrs),
        "precision@k": np.mean(precisions)
    }

import time
import statistics

def benchmark_speed(
    texts: list[str],
    embed_fn: Callable,
    iterations: int = 5
) -> dict[str, float]:
    """Benchmark embedding speed."""
    times = []

    for _ in range(iterations):
        start = time.perf_counter()
        embed_fn(texts)
        elapsed = time.perf_counter() - start
        times.append(elapsed)

    return {
        "mean_time": statistics.mean(times),
        "std_time": statistics.stdev(times) if len(times) > 1 else 0,
        "texts_per_second": len(texts) / statistics.mean(times),
        "ms_per_text": (statistics.mean(times) / len(texts)) * 1000
    }

import asyncio
from openai import AsyncOpenAI

async_client = AsyncOpenAI()

async def get_embeddings_async(
    texts: list[str],
    model: str = "text-embedding-3-large",
    batch_size: int = 100,
    max_concurrent: int = 5
) -> list[list[float]]:
    """Process large volumes with controlled concurrency."""
    semaphore = asyncio.Semaphore(max_concurrent)

    async def process_batch(batch: list[str]) -> list[list[float]]:
        async with semaphore:
            response = await async_client.embeddings.create(
                input=batch,
                model=model
            )
            sorted_data = sorted(response.data, key=lambda x: x.index)
            return [e.embedding for e in sorted_data]

    # Create batches
    batches = [texts[i:i+batch_size] for i in range(0, len(texts), batch_size)]

    # Process concurrently
    results = await asyncio.gather(*[process_batch(b) for b in batches])

    # Flatten
    return [emb for batch in results for emb in batch]

import time
import random
from openai import RateLimitError, APIError

def get_embedding_with_retry(
    text: str,
    model: str = "text-embedding-3-large",
    max_retries: int = 5
) -> list[float]:
    """Robust embedding with retry logic."""
    for attempt in range(max_retries):
        try:
            return get_embedding(text, model)
        except RateLimitError:
            wait = (2 ** attempt) + random.random()
            print(f"Rate limited. Waiting {wait:.1f}s...")
            time.sleep(wait)
        except APIError as e:
            if attempt == max_retries - 1:
                raise
            wait = 1 + random.random()
            print(f"API error: {e}. Retrying in {wait:.1f}s...")
            time.sleep(wait)

    raise Exception("Max retries exceeded")

from dataclasses import dataclass
from typing import Optional

@dataclass
class EmbeddingConfig:
    model: str = "text-embedding-3-large"
    dimensions: Optional[int] = None
    chunk_size: int = 500
    chunk_overlap: int = 50
    batch_size: int = 100
    cache_enabled: bool = True

class EmbeddingPipeline:
    def __init__(self, config: EmbeddingConfig):
        self.config = config
        self.cache = EmbeddingCache() if config.cache_enabled else None

    def process_document(self, text: str) -> list[list[float]]:
        """Full pipeline: chunk -> cache check -> embed."""
        # 1. Chunk
        chunks = chunk_by_tokens(
            text,
            max_tokens=self.config.chunk_size,
            overlap_tokens=self.config.chunk_overlap
        )

        # 2. Check cache
        uncached = []
        cached_embeddings = {}

        if self.cache:
            for i, chunk in enumerate(chunks):
                cached = self.cache.get(chunk, self.config.model)
                if cached:
                    cached_embeddings[i] = cached
                else:
                    uncached.append((i, chunk))
        else:
            uncached = list(enumerate(chunks))

        # 3. Embed uncached
        if uncached:
            indices, texts = zip(*uncached)
            new_embeddings = get_embeddings_batch(
                list(texts),
                model=self.config.model
            )

            for idx, emb in zip(indices, new_embeddings):
                cached_embeddings[idx] = emb
                if self.cache:
                    self.cache.set(chunks[idx], self.config.model, emb)

        # 4. Return in order
        return [cached_embeddings[i] for i in range(len(chunks))]

def debug_similarity(
    query: str,
    documents: list[str],
    model: str = "text-embedding-3-large"
) -> None:
    """Debug why certain documents rank high/low."""
    query_emb = np.array(get_embedding(query, model))
    doc_embs = np.array(get_embeddings_batch(documents, model))

    similarities = cosine_similarity([query_emb], doc_embs)[0]

    print(f"Query: {query[:100]}...")
    print("-" * 50)

    for doc, sim in sorted(zip(documents, similarities), key=lambda x: -x[1]):
        print(f"Score: {sim:.4f} | {doc[:80]}...")

# Qdrant
from qdrant_client import QdrantClient
from qdrant_client.models import PointStruct

client = QdrantClient("localhost", port=6333)
embeddings = get_embeddings_batch(texts)

points = [
    PointStruct(id=i, vector=emb, payload={"text": text})
    for i, (emb, text) in enumerate(zip(embeddings, texts))
]
client.upsert("my_collection", points)

# pgvector
import psycopg2

conn = psycopg2.connect("postgresql://...")
cur = conn.cursor()

for text, embedding in zip(texts, embeddings):
    cur.execute(
        "INSERT INTO documents (content, embedding) VALUES (%s, %s)",
        (text, embedding)
    )
conn.commit()

# Qdrant
query_embedding = get_embedding(query)
results = client.search(
    collection_name="my_collection",
    query_vector=query_embedding,
    limit=10
)

# pgvector
cur.execute("""
    SELECT content, 1 - (embedding <=> %s) as similarity
    FROM documents
    ORDER BY embedding <=> %s
    LIMIT 10
""", (query_embedding, query_embedding))