Production RAG Architectures & Grounded Answering

# Production reranking combines multiple signals beyond raw similarity.
# Weights are tunable per domain; freshness matters more for news,
# trust matters more for compliance docs.

def rank_candidate(candidate):
    # Semantic similarity from the embedding model (0-1)
    relevance = candidate["semantic_score"]
    # Trust score based on source authority (0-1)
    trust = candidate["source_trust"]
    # Freshness decay: newer docs score higher (0-1)
    freshness = candidate["freshness_score"]
    # Weighted combination - tune these for your domain
    return 0.65 * relevance + 0.20 * trust + 0.15 * freshness

# Sort candidates by composite score, best first
ranked = sorted(candidates, key=rank_candidate, reverse=True)

# Only pass top-k candidates that clear a minimum threshold
MIN_SCORE = 0.45
evidence = [c for c in ranked[:5] if rank_candidate(c) >= MIN_SCORE]

# Multi-hop retrieval: decompose, retrieve, refine, synthesize.
# Each hop uses previous results to formulate better queries.

def multi_hop_retrieve(question, retriever, llm, max_hops=3):
    # Step 1: Ask the LLM to decompose the question
    sub_questions = llm.decompose(question)
    all_evidence = []

    for i, sq in enumerate(sub_questions[:max_hops]):
        # Step 2: Retrieve evidence for each sub-question
        docs = retriever.search(sq, top_k=3)
        all_evidence.extend(docs)

        # Step 3: Optionally refine next query using current evidence
        if i < len(sub_questions) - 1:
            # Inject current findings into the next sub-question
            context_summary = llm.summarize(docs)
            sub_questions[i + 1] = llm.refine_query(
                sub_questions[i + 1], context_summary
            )

    # Step 4: Deduplicate and synthesize final answer
    unique_evidence = deduplicate(all_evidence)
    return llm.synthesize(question, unique_evidence)

# A simple grounding check: if the generated answer cannot be
# traced back to retrieved evidence, flag it for abstention.

def grounded_generate(question, evidence, llm):
    # Instruct the model to cite evidence and abstain if unsupported
    system_prompt = (
        "Answer the question using ONLY the provided evidence. "
        "Cite evidence by [doc_id]. If the evidence does not support "
        "a confident answer, respond with: 'I don't have enough "
        "information to answer this reliably.'"
    )

    response = llm.generate(
        system=system_prompt,
        user=f"Question: {question}\n\nEvidence:\n{format_evidence(evidence)}"
    )

    # Post-generation check: does the answer cite at least one source?
    if not has_citations(response):
        # No citations found - likely ungrounded, trigger fallback
        return {"answer": None, "status": "abstained", "reason": "no_citations"}

    return {"answer": response, "status": "grounded"}

# Generate an answer with structured citations.
# Each claim maps to the evidence passage that supports it.

import json

def generate_cited_answer(question, evidence_docs, llm):
    # Build a schema-enforced prompt for citation
    schema_instruction = """Respond in JSON:
{
  "answer": "your full answer text with [1], [2] markers",
  "citations": [
    {"id": 1, "doc_id": "...", "passage": "exact quote", "claim": "what it supports"}
  ],
  "confidence": "high | medium | low",
  "unsupported_claims": ["any claims you could not ground"]
}"""

    # Format evidence with doc IDs for reference
    evidence_text = "\n".join(
        f"[Doc {d['id']}]: {d['text']}"
        for d in evidence_docs
    )

    raw = llm.generate(
        system=schema_instruction,
        user=f"Question: {question}\n\nEvidence:\n{evidence_text}"
    )

    # Parse and validate the structured response
    result = json.loads(raw)
    # Flag answers where no citations were produced
    if not result.get("citations"):
        result["confidence"] = "low"
    return result

# Filter and score documents by freshness before ranking.
# Ensures stale evidence does not dominate the context window.

from datetime import datetime, timedelta

def freshness_score(doc, max_age_days=90):
    # Calculate how fresh the document is (0 = expired, 1 = brand new)
    age = (datetime.now() - doc["last_updated"]).days
    if age > max_age_days:
        return 0.0  # Beyond max age, treat as stale
    return 1.0 - (age / max_age_days)

def filter_stale(docs, min_freshness=0.1):
    # Remove documents below the freshness threshold
    fresh = [d for d in docs if freshness_score(d) >= min_freshness]
    if not fresh:
        # All docs are stale - signal uncertainty upstream
        return [], "all_evidence_stale"
    return fresh, "ok"

# Enforce document permissions at retrieval time.
# Never let the model see documents the user cannot access.

def permission_scoped_search(query, user, vector_store):
    # Get the user's permission groups
    user_groups = user["access_groups"]  # e.g., ["engineering", "public"]

    # Build a metadata filter that restricts to allowed documents
    acl_filter = {
        "access_groups": {"$in": user_groups}
    }

    # The vector search only sees documents matching the ACL filter
    results = vector_store.similarity_search(
        query=query,
        k=10,
        filter=acl_filter  # Pre-retrieval: restricted docs never enter the results
    )

    # Double-check: log any result without ACL metadata as anomalous
    for r in results:
        if not r.metadata.get("access_groups"):
            logger.warning(f"Doc {r.id} missing ACL metadata")

    return results

# Agentic RAG: the model decides which tools to use and
# whether evidence is sufficient before generating.

def agentic_rag(question, tools, llm, max_steps=5):
    # Available tools: vector_search, sql_query, api_lookup, etc.
    context = []
    plan = llm.plan(question, available_tools=tools)

    for step in plan.steps[:max_steps]:
        # The model chose which tool and query to use
        tool = tools[step.tool_name]
        result = tool.execute(step.query)
        context.append({"tool": step.tool_name, "result": result})

        # After each step, ask: do we have enough evidence?
        sufficiency = llm.evaluate_evidence(question, context)
        if sufficiency.is_sufficient:
            break  # Evidence is good enough, proceed to generation

    # Generate the final answer using all collected evidence
    return llm.generate(question, context)

# A permission-aware cache for RAG answers.
# Cache keys include user scope to prevent cross-user leaks.

import hashlib, time

class ScopedAnswerCache:
    def __init__(self, ttl_seconds=300):
        self.cache = {}
        self.ttl = ttl_seconds

    def _key(self, query, user_scope):
        # Include user permissions in the cache key
        scope_str = ",".join(sorted(user_scope))
        raw = f"{query}||{scope_str}"
        return hashlib.sha256(raw.encode()).hexdigest()

    def get(self, query, user_scope):
        key = self._key(query, user_scope)
        entry = self.cache.get(key)
        if entry and (time.time() - entry["ts"]) < self.ttl:
            return entry["answer"]  # Cache hit, within TTL
        return None  # Cache miss or expired

    def put(self, query, user_scope, answer):
        key = self._key(query, user_scope)
        self.cache[key] = {"answer": answer, "ts": time.time()}

# A minimal offline evaluation harness for RAG.
# Measures retrieval recall, groundedness, and answer quality.

def evaluate_rag(test_cases, rag_pipeline, judge_llm):
    results = []
    for tc in test_cases:
        # Run the RAG pipeline on the test question
        output = rag_pipeline.run(tc["question"])

        # Metric 1: Did we retrieve the expected source documents?
        retrieved_ids = {d["id"] for d in output["retrieved_docs"]}
        expected_ids = set(tc["expected_doc_ids"])
        recall = len(retrieved_ids & expected_ids) / max(len(expected_ids), 1)

        # Metric 2: Is the answer grounded in the evidence?
        groundedness = judge_llm.score_groundedness(
            answer=output["answer"],
            evidence=output["retrieved_docs"]
        )

        # Metric 3: Overall answer quality vs reference
        quality = judge_llm.score_quality(
            answer=output["answer"],
            reference=tc["reference_answer"]
        )

        results.append({
            "question": tc["question"],
            "retrieval_recall": recall,
            "groundedness": groundedness,
            "quality": quality
        })
    return results

# A simple decision function to route tasks to the right architecture.
# Not every LLM task needs RAG; use the simplest adequate approach.

def choose_architecture(task):
    # Check if the task is better served by deterministic tools
    if task.requires_computation:
        return "code_execution"  # Math, data transforms, logic

    if task.data_is_structured:
        return "text_to_sql"  # Tabular data, databases

    if task.workflow_is_static:
        return "template_engine"  # Fixed procedures, decision trees

    if task.corpus_too_noisy:
        return "curate_first"  # Clean data before building RAG

    if task.search_suffices:
        return "traditional_search"  # Snippets are enough

    # RAG is appropriate: unstructured knowledge, grounding needed
    if task.needs_grounded_answers and task.has_quality_corpus:
        return "rag"

    return "evaluate_further"  # Edge case - needs human judgment