┌─────────────────────────────────────────────────────────────────────────────┐
│                          TRADITIONAL RAG                                    │
│                                                                             │
│  All Documents ───► Single Generation ───► One Answer                       │
│       │                    │                                                │
│       └── 10 docs all      └── Position bias: LLM favors start/end          │
│           at once              Information overload: key points diluted     │
│                                Single attempt: may miss important context   │
└─────────────────────────────────────────────────────────────────────────────┘

┌─────────────────────────────────────────────────────────────────────────────┐
│                          SPECULATIVE RAG                                    │
│                                                                             │
│  Retrieved Documents                                                        │
│         │                                                                   │
│         ▼                                                                   │
│  ┌─────────────────┐                                                        │
│  │ Cluster &       │                                                        │
│  │ Partition       │ (Ensure diverse subsets)                               │
│  └────────┬────────┘                                                        │
│           │                                                                 │
│    ┌──────┼──────┐                                                          │
│    ▼      ▼      ▼                                                          │
│  ┌────┐ ┌────┐ ┌────┐                                                       │
│  │Sub1│ │Sub2│ │Sub3│  (Different doc combinations)                         │
│  └──┬─┘ └──┬─┘ └──┬─┘                                                       │
│     │      │      │                                                         │
│     ▼      ▼      ▼      (Parallel execution - same wall-clock time)        │
│  ┌──────┬──────┬──────┐                                                     │
│  │Draft1│Draft2│Draft3│  GPT-4o-mini (cheap, fast)                          │
│  └──┬───┴──┬───┴──┬───┘                                                     │
│     │      │      │                                                         │
│     └──────┼──────┘                                                         │
│            ▼                                                                │
│     ┌─────────────┐                                                         │
│     │  Verifier   │  GPT-4o (accurate, evaluates all drafts)                │
│     │  (Large LM) │                                                         │
│     └──────┬──────┘                                                         │
│            ▼                                                                │
│       Best Answer                                                           │
│                                                                             │
│  Result: 13% better accuracy + 50% lower latency                            │
└─────────────────────────────────────────────────────────────────────────────┘

speculative-rag/
├── config.py                # Configuration
├── retriever.py            # Document retrieval
├── partitioner.py          # Document clustering & partitioning
├── drafter.py              # Parallel draft generation
├── verifier.py             # Draft verification & selection
├── speculative_rag.py      # Main orchestration
├── app.py                  # FastAPI application
└── requirements.txt

# config.py
from pydantic_settings import BaseSettings
from pydantic import Field
from functools import lru_cache


class Settings(BaseSettings):
    """Application configuration."""

    openai_api_key: str

    # Model settings
    embedding_model: str = "text-embedding-3-small"
    drafter_model: str = "gpt-4o-mini"      # Fast, cheap for drafting
    verifier_model: str = "gpt-4o"           # Accurate for verification

    # Retrieval settings
    retrieval_k: int = 12                    # Total docs to retrieve
    num_drafts: int = 3                      # Number of parallel drafts
    docs_per_draft: int = 4                  # Docs per draft subset

    # Clustering settings
    num_clusters: int = 4                    # For document diversity
    min_cluster_size: int = 2

    # Generation settings
    drafter_temperature: float = 0.7
    drafter_max_tokens: int = 512
    verifier_temperature: float = 0.3

    # ChromaDB
    chroma_persist_dir: str = "./chroma_db"
    collection_name: str = "speculative_rag_docs"

    class Config:
        env_file = ".env"


@lru_cache
def get_settings() -> Settings:
    return Settings()

# retriever.py
import chromadb
from chromadb.utils import embedding_functions
from openai import OpenAI
from pydantic import BaseModel
from config import get_settings


class RetrievedDocument(BaseModel):
    """A retrieved document with embedding."""
    id: str
    content: str
    source: str
    embedding: list[float]
    distance: float


class Retriever:
    """Document retriever with embeddings."""

    def __init__(self):
        settings = get_settings()

        self.openai = OpenAI(api_key=settings.openai_api_key)
        self.embedding_model = settings.embedding_model

        # ChromaDB setup
        self.chroma = chromadb.PersistentClient(
            path=settings.chroma_persist_dir
        )
        self.embedding_fn = embedding_functions.OpenAIEmbeddingFunction(
            api_key=settings.openai_api_key,
            model_name=settings.embedding_model
        )
        self.collection = self.chroma.get_or_create_collection(
            name=settings.collection_name,
            embedding_function=self.embedding_fn,
            metadata={"hnsw:space": "cosine"}
        )

    def retrieve(
        self,
        query: str,
        k: int | None = None
    ) -> list[RetrievedDocument]:
        """Retrieve documents with their embeddings."""
        settings = get_settings()
        k = k or settings.retrieval_k

        results = self.collection.query(
            query_texts=[query],
            n_results=k,
            include=["documents", "metadatas", "distances", "embeddings"]
        )

        documents = []
        for i in range(len(results["documents"][0])):
            doc = RetrievedDocument(
                id=results["ids"][0][i],
                content=results["documents"][0][i],
                source=results["metadatas"][0][i].get("source", "unknown"),
                embedding=results["embeddings"][0][i],
                distance=results["distances"][0][i]
            )
            documents.append(doc)

        return documents

    def add_documents(
        self,
        documents: list[str],
        sources: list[str]
    ):
        """Add documents to the collection."""
        ids = [f"doc_{i}" for i in range(len(documents))]
        self.collection.add(
            documents=documents,
            ids=ids,
            metadatas=[{"source": src} for src in sources]
        )

# partitioner.py
import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics.pairwise import cosine_similarity
from pydantic import BaseModel
from retriever import RetrievedDocument
from config import get_settings


class DocumentSubset(BaseModel):
    """A subset of documents for draft generation."""
    subset_id: int
    documents: list[RetrievedDocument]
    diversity_score: float  # Higher = more diverse


class DocumentPartitioner:
    """Partitions documents into diverse subsets for drafting."""

    def __init__(self):
        settings = get_settings()
        self.num_clusters = settings.num_clusters
        self.num_drafts = settings.num_drafts
        self.docs_per_draft = settings.docs_per_draft

    def partition(
        self,
        documents: list[RetrievedDocument]
    ) -> list[DocumentSubset]:
        """
        Partition documents into diverse subsets.

        Strategy:
        1. Cluster documents by content similarity
        2. For each subset, sample one document from each cluster
        3. This maximizes diversity while minimizing redundancy
        """
        if len(documents) < self.num_drafts:
            # Not enough documents - duplicate
            return [
                DocumentSubset(
                    subset_id=0,
                    documents=documents,
                    diversity_score=1.0
                )
            ]

        # Get embeddings matrix
        embeddings = np.array([doc.embedding for doc in documents])

        # Cluster documents
        n_clusters = min(self.num_clusters, len(documents))
        kmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=10)
        cluster_labels = kmeans.fit_predict(embeddings)

        # Group documents by cluster
        clusters: dict[int, list[int]] = {}
        for idx, label in enumerate(cluster_labels):
            if label not in clusters:
                clusters[label] = []
            clusters[label].append(idx)

        # Create diverse subsets
        subsets = []
        for subset_id in range(self.num_drafts):
            subset_docs = []
            used_clusters = set()

            # Sample one from each cluster (round-robin with offset)
            cluster_ids = list(clusters.keys())
            for i in range(self.docs_per_draft):
                cluster_idx = (subset_id + i) % len(cluster_ids)
                cluster_id = cluster_ids[cluster_idx]

                if cluster_id in used_clusters and len(clusters[cluster_id]) <= 1:
                    # Skip if already used and no more docs in cluster
                    continue

                # Get document from cluster
                doc_indices = clusters[cluster_id]
                # Use different doc for different subsets
                doc_idx = doc_indices[subset_id % len(doc_indices)]
                subset_docs.append(documents[doc_idx])
                used_clusters.add(cluster_id)

                if len(subset_docs) >= self.docs_per_draft:
                    break

            # Calculate diversity score
            diversity = self._calculate_diversity(subset_docs)

            subsets.append(DocumentSubset(
                subset_id=subset_id,
                documents=subset_docs,
                diversity_score=diversity
            ))

        return subsets

    def _calculate_diversity(
        self,
        documents: list[RetrievedDocument]
    ) -> float:
        """Calculate diversity score for a subset."""
        if len(documents) < 2:
            return 0.0

        embeddings = np.array([doc.embedding for doc in documents])
        similarities = cosine_similarity(embeddings)

        # Average pairwise similarity (lower = more diverse)
        n = len(documents)
        total_sim = 0
        count = 0
        for i in range(n):
            for j in range(i + 1, n):
                total_sim += similarities[i][j]
                count += 1

        avg_similarity = total_sim / count if count > 0 else 0

        # Convert to diversity (1 - similarity)
        return 1 - avg_similarity

# drafter.py
import asyncio
from openai import AsyncOpenAI
from pydantic import BaseModel
from partitioner import DocumentSubset
from config import get_settings


class Draft(BaseModel):
    """A generated draft answer."""
    draft_id: int
    content: str
    confidence: float
    sources_used: list[str]
    subset_diversity: float
    generation_time_ms: float


class DraftRationale(BaseModel):
    """Rationale for draft generation."""
    key_points: list[str]
    evidence_used: list[str]
    confidence_reason: str


class ParallelDrafter:
    """Generates multiple drafts in parallel from document subsets."""

    def __init__(self):
        settings = get_settings()
        self.client = AsyncOpenAI(api_key=settings.openai_api_key)
        self.model = settings.drafter_model
        self.temperature = settings.drafter_temperature
        self.max_tokens = settings.drafter_max_tokens

    async def generate_drafts(
        self,
        query: str,
        subsets: list[DocumentSubset]
    ) -> list[Draft]:
        """Generate drafts in parallel from document subsets."""

        # Create tasks for parallel execution
        tasks = [
            self._generate_single_draft(query, subset)
            for subset in subsets
        ]

        # Execute in parallel
        drafts = await asyncio.gather(*tasks)

        return drafts

    async def _generate_single_draft(
        self,
        query: str,
        subset: DocumentSubset
    ) -> Draft:
        """Generate a single draft from a document subset."""
        import time
        start = time.time()

        # Build context from subset documents
        context = "\n\n---\n\n".join([
            f"Source: {doc.source}\n{doc.content}"
            for doc in subset.documents
        ])

        system_prompt = """You are a precise answer generator. Generate a focused answer
based ONLY on the provided documents.

Guidelines:
1. Use only information from the provided sources
2. Be concise but comprehensive
3. Cite sources when making claims
4. If documents don't contain enough information, say so
5. Rate your confidence (0.0-1.0) based on source quality

Return JSON:
{
    "answer": "Your answer here",
    "confidence": 0.85,
    "key_points": ["point1", "point2"],
    "sources_cited": ["source1", "source2"]
}"""

        response = await self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": f"Documents:\n{context}\n\nQuestion: {query}"}
            ],
            temperature=self.temperature,
            max_tokens=self.max_tokens,
            response_format={"type": "json_object"}
        )

        import json
        result = json.loads(response.choices[0].message.content)

        generation_time = (time.time() - start) * 1000

        return Draft(
            draft_id=subset.subset_id,
            content=result.get("answer", ""),
            confidence=result.get("confidence", 0.5),
            sources_used=result.get("sources_cited", []),
            subset_diversity=subset.diversity_score,
            generation_time_ms=generation_time
        )

# verifier.py
from openai import OpenAI
from pydantic import BaseModel, Field
from drafter import Draft
from retriever import RetrievedDocument
from config import get_settings


class VerificationResult(BaseModel):
    """Result of draft verification."""
    selected_draft_id: int
    selected_answer: str
    confidence: float
    reasoning: str
    draft_scores: dict[int, float]
    verification_time_ms: float


class DraftEvaluation(BaseModel):
    """Evaluation of a single draft."""
    draft_id: int
    accuracy_score: float = Field(ge=0, le=1)
    completeness_score: float = Field(ge=0, le=1)
    coherence_score: float = Field(ge=0, le=1)
    overall_score: float = Field(ge=0, le=1)
    issues: list[str]


class Verifier:
    """Verifies and selects the best draft."""

    def __init__(self):
        settings = get_settings()
        self.client = OpenAI(api_key=settings.openai_api_key)
        self.model = settings.verifier_model
        self.temperature = settings.verifier_temperature

    def verify(
        self,
        query: str,
        drafts: list[Draft],
        all_documents: list[RetrievedDocument]
    ) -> VerificationResult:
        """
        Verify drafts and select the best one.

        The verifier:
        1. Evaluates each draft for accuracy, completeness, coherence
        2. Cross-references claims with source documents
        3. Selects the best draft or synthesizes if needed
        """
        import time
        start = time.time()

        # Build verification context
        drafts_text = ""
        for draft in drafts:
            drafts_text += f"\n--- Draft {draft.draft_id} ---\n"
            drafts_text += f"Content: {draft.content}\n"
            drafts_text += f"Self-reported confidence: {draft.confidence}\n"
            drafts_text += f"Sources used: {', '.join(draft.sources_used)}\n"

        sources_text = "\n\n".join([
            f"[{doc.source}]: {doc.content}"
            for doc in all_documents
        ])

        system_prompt = """You are an expert answer verifier. Your job is to:
1. Evaluate each draft answer for accuracy, completeness, and coherence
2. Cross-reference claims with the source documents
3. Select the best draft OR synthesize a better answer if all drafts have issues

Scoring criteria:
- Accuracy: Are claims supported by sources? (0-1)
- Completeness: Does it fully answer the question? (0-1)
- Coherence: Is it well-organized and clear? (0-1)

Return JSON:
{
    "evaluations": [
        {
            "draft_id": 0,
            "accuracy_score": 0.9,
            "completeness_score": 0.8,
            "coherence_score": 0.85,
            "overall_score": 0.85,
            "issues": ["Minor issue 1"]
        }
    ],
    "selected_draft_id": 0,
    "final_answer": "The selected or improved answer",
    "reasoning": "Why this draft was selected",
    "confidence": 0.9
}"""

        response = self.client.chat.completions.create(
            model=self.model,
            messages=[
                {"role": "system", "content": system_prompt},
                {
                    "role": "user",
                    "content": f"Question: {query}\n\n"
                               f"Drafts to evaluate:\n{drafts_text}\n\n"
                               f"Source documents for verification:\n{sources_text}"
                }
            ],
            temperature=self.temperature,
            response_format={"type": "json_object"}
        )

        import json
        result = json.loads(response.choices[0].message.content)

        verification_time = (time.time() - start) * 1000

        # Build draft scores dict
        draft_scores = {
            eval_data["draft_id"]: eval_data["overall_score"]
            for eval_data in result.get("evaluations", [])
        }

        return VerificationResult(
            selected_draft_id=result.get("selected_draft_id", 0),
            selected_answer=result.get("final_answer", drafts[0].content),
            confidence=result.get("confidence", 0.5),
            reasoning=result.get("reasoning", ""),
            draft_scores=draft_scores,
            verification_time_ms=verification_time
        )


class SelfConsistencyVerifier:
    """Alternative verifier using self-consistency (majority voting)."""

    def verify(
        self,
        drafts: list[Draft]
    ) -> VerificationResult:
        """
        Select best draft using self-consistency.

        Compares drafts pairwise and selects the one most
        consistent with others (majority voting).
        """
        import time
        start = time.time()

        if len(drafts) == 1:
            return VerificationResult(
                selected_draft_id=0,
                selected_answer=drafts[0].content,
                confidence=drafts[0].confidence,
                reasoning="Only one draft available",
                draft_scores={0: 1.0},
                verification_time_ms=0
            )

        # Score based on self-reported confidence and diversity
        scores = {}
        for draft in drafts:
            # Combine confidence with diversity bonus
            score = draft.confidence * 0.7 + draft.subset_diversity * 0.3
            scores[draft.draft_id] = score

        # Select highest scoring
        best_id = max(scores, key=scores.get)
        best_draft = next(d for d in drafts if d.draft_id == best_id)

        verification_time = (time.time() - start) * 1000

        return VerificationResult(
            selected_draft_id=best_id,
            selected_answer=best_draft.content,
            confidence=best_draft.confidence,
            reasoning=f"Selected based on confidence ({best_draft.confidence:.2f}) "
                      f"and diversity ({best_draft.subset_diversity:.2f})",
            draft_scores=scores,
            verification_time_ms=verification_time
        )

# speculative_rag.py
import asyncio
from pydantic import BaseModel
from retriever import Retriever, RetrievedDocument
from partitioner import DocumentPartitioner, DocumentSubset
from drafter import ParallelDrafter, Draft
from verifier import Verifier, VerificationResult
from config import get_settings


class SpeculativeRAGMetrics(BaseModel):
    """Performance metrics."""
    total_latency_ms: float
    retrieval_latency_ms: float
    partition_latency_ms: float
    drafting_latency_ms: float
    verification_latency_ms: float
    num_drafts: int
    docs_retrieved: int
    parallel_speedup: float  # vs sequential


class SpeculativeRAGResponse(BaseModel):
    """Response from Speculative RAG."""
    answer: str
    confidence: float
    selected_draft_id: int
    all_drafts: list[Draft]
    draft_scores: dict[int, float]
    verification_reasoning: str
    sources: list[str]
    metrics: SpeculativeRAGMetrics


class SpeculativeRAG:
    """Speculative RAG with parallel drafting and verification."""

    def __init__(self):
        self.retriever = Retriever()
        self.partitioner = DocumentPartitioner()
        self.drafter = ParallelDrafter()
        self.verifier = Verifier()

    async def query(self, question: str) -> SpeculativeRAGResponse:
        """
        Process query with speculative RAG.

        Pipeline:
        1. Retrieve documents
        2. Partition into diverse subsets
        3. Generate drafts in parallel
        4. Verify and select best draft
        """
        import time
        total_start = time.time()

        # Step 1: Retrieve
        retrieval_start = time.time()
        documents = self.retriever.retrieve(question)
        retrieval_time = (time.time() - retrieval_start) * 1000

        # Step 2: Partition
        partition_start = time.time()
        subsets = self.partitioner.partition(documents)
        partition_time = (time.time() - partition_start) * 1000

        # Step 3: Generate drafts in parallel
        drafting_start = time.time()
        drafts = await self.drafter.generate_drafts(question, subsets)
        drafting_time = (time.time() - drafting_start) * 1000

        # Calculate what sequential would have been
        sequential_time = sum(d.generation_time_ms for d in drafts)
        parallel_speedup = sequential_time / drafting_time if drafting_time > 0 else 1

        # Step 4: Verify and select
        verification_start = time.time()
        verification = self.verifier.verify(question, drafts, documents)
        verification_time = (time.time() - verification_start) * 1000

        total_time = (time.time() - total_start) * 1000

        # Collect all sources used
        all_sources = list(set(
            src for draft in drafts for src in draft.sources_used
        ))

        metrics = SpeculativeRAGMetrics(
            total_latency_ms=total_time,
            retrieval_latency_ms=retrieval_time,
            partition_latency_ms=partition_time,
            drafting_latency_ms=drafting_time,
            verification_latency_ms=verification_time,
            num_drafts=len(drafts),
            docs_retrieved=len(documents),
            parallel_speedup=parallel_speedup
        )

        return SpeculativeRAGResponse(
            answer=verification.selected_answer,
            confidence=verification.confidence,
            selected_draft_id=verification.selected_draft_id,
            all_drafts=drafts,
            draft_scores=verification.draft_scores,
            verification_reasoning=verification.reasoning,
            sources=all_sources,
            metrics=metrics
        )

    def add_documents(self, documents: list[str], sources: list[str]):
        """Add documents to the knowledge base."""
        self.retriever.add_documents(documents, sources)


async def compare_with_traditional(
    question: str,
    speculative_rag: SpeculativeRAG
) -> dict:
    """Compare Speculative RAG with traditional single-generation RAG."""
    from openai import OpenAI
    import time

    settings = get_settings()
    client = OpenAI(api_key=settings.openai_api_key)

    # Run speculative RAG
    spec_result = await speculative_rag.query(question)

    # Run traditional RAG (single generation with all docs)
    trad_start = time.time()
    documents = speculative_rag.retriever.retrieve(question)
    context = "\n\n".join([d.content for d in documents])

    trad_response = client.chat.completions.create(
        model=settings.verifier_model,  # Use same quality model
        messages=[
            {"role": "system", "content": "Answer based on the context."},
            {"role": "user", "content": f"Context:\n{context}\n\nQuestion: {question}"}
        ]
    )
    trad_time = (time.time() - trad_start) * 1000

    return {
        "speculative": {
            "answer": spec_result.answer,
            "latency_ms": spec_result.metrics.total_latency_ms,
            "confidence": spec_result.confidence,
            "parallel_speedup": spec_result.metrics.parallel_speedup
        },
        "traditional": {
            "answer": trad_response.choices[0].message.content,
            "latency_ms": trad_time
        },
        "speedup": trad_time / spec_result.metrics.total_latency_ms
    }

# app.py
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from contextlib import asynccontextmanager

from speculative_rag import (
    SpeculativeRAG,
    SpeculativeRAGResponse,
    compare_with_traditional
)


# Global
speculative_rag: SpeculativeRAG | None = None


@asynccontextmanager
async def lifespan(app: FastAPI):
    global speculative_rag
    speculative_rag = SpeculativeRAG()

    # Add sample documents
    sample_docs = [
        "Machine learning is a subset of AI that enables systems to learn from data. Supervised learning uses labeled data, unsupervised learning finds patterns in unlabeled data.",
        "Deep learning uses neural networks with multiple layers. CNNs excel at image processing, RNNs handle sequential data, and Transformers power modern NLP.",
        "Transfer learning allows pre-trained models to be adapted for new tasks. This is efficient when labeled data is scarce. BERT and GPT are commonly fine-tuned.",
        "RAG (Retrieval-Augmented Generation) combines retrieval with generation. It reduces hallucination by grounding responses in retrieved documents.",
        "Vector databases store embeddings for similarity search. Popular options include Pinecone, Weaviate, Milvus, and ChromaDB.",
        "Prompt engineering optimizes LLM outputs through careful prompt design. Techniques include few-shot learning, chain-of-thought, and role prompting.",
        "LLM evaluation metrics include BLEU, ROUGE for text similarity, and human evaluation for quality. Automated metrics often correlate poorly with human judgment.",
        "Fine-tuning adapts pre-trained models to specific tasks. Full fine-tuning updates all weights, while LoRA only updates low-rank adapters.",
        "Quantization reduces model size by using lower precision. INT8 and INT4 quantization can reduce memory 4x with minimal quality loss.",
        "Model serving requires consideration of latency, throughput, and cost. Batching, caching, and model optimization are key techniques.",
        "MLOps encompasses the practices for deploying and maintaining ML systems. It includes CI/CD, monitoring, versioning, and reproducibility.",
        "LangChain is a framework for building LLM applications. It provides abstractions for chains, agents, memory, and retrieval."
    ]
    sources = [
        "ml_basics", "deep_learning", "transfer_learning", "rag_overview",
        "vector_dbs", "prompt_engineering", "llm_evaluation", "fine_tuning",
        "quantization", "model_serving", "mlops", "langchain"
    ]
    speculative_rag.add_documents(sample_docs, sources)

    yield

    speculative_rag = None


app = FastAPI(
    title="Speculative RAG API",
    description="Parallel draft generation with verification for faster, more accurate RAG",
    lifespan=lifespan
)


class QueryRequest(BaseModel):
    query: str


class DocumentsRequest(BaseModel):
    documents: list[str]
    sources: list[str]


@app.post("/query", response_model=SpeculativeRAGResponse)
async def query(request: QueryRequest):
    """Query with Speculative RAG."""
    if not speculative_rag:
        raise HTTPException(status_code=503, detail="Service not initialized")

    result = await speculative_rag.query(request.query)
    return result


@app.post("/compare")
async def compare(request: QueryRequest):
    """Compare Speculative RAG with traditional RAG."""
    if not speculative_rag:
        raise HTTPException(status_code=503, detail="Service not initialized")

    result = await compare_with_traditional(request.query, speculative_rag)
    return result


@app.post("/documents")
async def add_documents(request: DocumentsRequest):
    """Add documents to the knowledge base."""
    if not speculative_rag:
        raise HTTPException(status_code=503, detail="Service not initialized")

    if len(request.documents) != len(request.sources):
        raise HTTPException(
            status_code=400,
            detail="Documents and sources must have same length"
        )

    speculative_rag.add_documents(request.documents, request.sources)
    return {"status": "success", "documents_added": len(request.documents)}


@app.get("/health")
async def health():
    return {"status": "healthy", "service": "speculative-rag"}

# requirements.txt
openai>=1.12.0
chromadb>=0.4.22
scikit-learn>=1.4.0
numpy>=1.24.0
pydantic>=2.0.0
pydantic-settings>=2.0.0
fastapi>=0.109.0
uvicorn>=0.27.0
python-dotenv>=1.0.0

from speculative_rag import SpeculativeRAG
import asyncio

async def main():
    rag = SpeculativeRAG()

    # Add documents
    rag.add_documents(
        documents=["Your content here..."],
        sources=["source_name"]
    )

    # Query
    result = await rag.query("What is the difference between RAG and fine-tuning?")

    print(f"Answer: {result.answer}")
    print(f"Confidence: {result.confidence:.2f}")
    print(f"Selected draft: {result.selected_draft_id}")
    print(f"Draft scores: {result.draft_scores}")
    print(f"\nMetrics:")
    print(f"  Total latency: {result.metrics.total_latency_ms:.2f}ms")
    print(f"  Parallel speedup: {result.metrics.parallel_speedup:.2f}x")

asyncio.run(main())

from speculative_rag import SpeculativeRAG, compare_with_traditional

async def benchmark():
    rag = SpeculativeRAG()
    # ... add documents ...

    comparison = await compare_with_traditional(
        "Explain how to deploy ML models in production",
        rag
    )

    print(f"Speculative RAG: {comparison['speculative']['latency_ms']:.2f}ms")
    print(f"Traditional RAG: {comparison['traditional']['latency_ms']:.2f}ms")
    print(f"Speedup: {comparison['speedup']:.2f}x")

# Start server
uvicorn app:app --reload

# Query
curl -X POST http://localhost:8000/query \
  -H "Content-Type: application/json" \
  -d '{"query": "What are the key components of MLOps?"}'

# Compare with traditional RAG
curl -X POST http://localhost:8000/compare \
  -H "Content-Type: application/json" \
  -d '{"query": "Explain transfer learning and fine-tuning"}'

┌─────────────────────────────────────────────────────────────────────────────┐
│                     SPECULATIVE RAG ARCHITECTURE                            │
├─────────────────────────────────────────────────────────────────────────────┤
│                                                                             │
│  1. RETRIEVAL                                                               │
│     Query ───► Retrieve K docs ───► Documents + Embeddings                  │
│                                                                             │
│  ───────────────────────────────────────────────────────────────────────── │
│                                                                             │
│  2. PARTITION                                                               │
│     Documents ───► Cluster by Similarity ───► Create Diverse Subsets       │
│                         │                           │                       │
│                    (K-means on                 (Each subset has            │
│                     embeddings)                different focus)            │
│                                                                             │
│  ───────────────────────────────────────────────────────────────────────── │
│                                                                             │
│  3. PARALLEL DRAFTING                                                       │
│     ┌────────────────┬────────────────┬────────────────┐                   │
│     │    Subset 1    │    Subset 2    │    Subset 3    │                   │
│     │       │        │       │        │       │        │                   │
│     │       ▼        │       ▼        │       ▼        │   (async)         │
│     │  GPT-4o-mini   │  GPT-4o-mini   │  GPT-4o-mini   │                   │
│     │       │        │       │        │       │        │                   │
│     │       ▼        │       ▼        │       ▼        │                   │
│     │    Draft 1     │    Draft 2     │    Draft 3     │                   │
│     └────────────────┴────────────────┴────────────────┘                   │
│                                                                             │
│  ───────────────────────────────────────────────────────────────────────── │
│                                                                             │
│  4. VERIFICATION                                                            │
│     All Drafts + Original Docs ───► GPT-4o Verifier                        │
│                                           │                                 │
│                                           ▼                                 │
│                                    Score Each Draft                         │
│                                    (support, relevance,                     │
│                                     completeness)                           │
│                                           │                                 │
│                                           ▼                                 │
│                                    Select Best Answer                       │
│                                                                             │
└─────────────────────────────────────────────────────────────────────────────┘