Building retrieval-augmented generation (RAG) systems for statutory law and appellate jurisprudence is fundamentally different from building generic search over marketing documents.

In legal domains, precision is paramount: a single hallucinated citation or an imprecise statutory reference invalidates the entire analysis. Furthermore, appellate case records in Pakistani law are notoriously dense, spanning decades of High Court and Supreme Court precedents with complex linguistic conventions.

When designing the backend for InsafDost AI (live at insafdostai.abdurrafaylabs.tech), we had to solve three core engineering challenges:

  1. Preventing hallucinated citations through a strict two-stage retrieval pipeline.
  2. Balancing response latency with deep reasoning using dynamic multi-model routing.
  3. Streaming granular node-by-node execution states over Server-Sent Events (SSE) without blocking the FastAPI async event loop.

Here is the architectural breakdown and empirical evaluation behind the system.

1. System Topology: Dynamic Groq Routing

Instead of routing all requests through a single monolithic model, InsafDost partitions inference across a state graph using Groq-hosted LLaMA architectures:

[Inbound Client Request]
          │
          ▼
[Fail-Closed Classification Guardrail] ──► (Invalid? Reject with 400)
          │
          ▼
[Statutory Pre-Processor (Category & Query Expansion)]
          │
          ▼
[Two-Stage Retrieval: Qdrant Dense k=8 + BAAI Cross-Encoder]
          │
          ▼
[Appellate Reasoner Node (Groq 120B Legal Synthesis)]
          │
          ▼
[Grounding Auditor Node (Groq 20B Citation Verifier)]
          │
          ▼
[SSE Stream or JSON Response Serialization]

Why Multi-Model Routing?

  • Classification & Guardrails (Groq 20B): Rapid classification into civil, criminal, constitutional, or family jurisdictions requires low latency. Using a 20B model guarantees sub-300ms evaluation.
  • Complex Appellate Synthesis (Groq 120B): Generating arguments, analyzing statutory exceptions, and reconciling conflicting case precedents demands high reasoning capacity.
  • Grounding Audit (Groq 20B): Before any generated litigation strategy is sent to the client, a dedicated auditor node verifies that every cited section and precedent exists verbatim in the retrieved context.

2. Two-Stage Retrieval: +19.0% Relative Hit@1

Standard approximate nearest-neighbor (ANN) search on dense vector embeddings often struggles with legal queries. Two legal precedents may share broad semantic similarity while differing on a crucial operational nuance (e.g., whether an injunction requires proof of irreparable harm or mere balance of convenience).

To resolve this, we implemented a two-stage retrieval pipeline:

  1. First Stage (Dense Retrieval): Approximate nearest-neighbor search against Qdrant (pakistan_law collection) retrieving the top $k=8$ candidates.
  2. Second Stage (Cross-Encoder Reranking): Pairs the query with each candidate using BAAI/bge-reranker-base. Cross-encoders evaluate full cross-attention between query and document tokens, capturing subtle statutory dependencies that cosine distance overlooks.
import asyncio
from sentence_transformers import CrossEncoder

reranker = CrossEncoder("BAAI/bge-reranker-base")

async def rerank_candidates(query: str, documents: list[str], threshold: float = 0.40) -> list[str]:
    pairs = [[query, doc] for doc in documents]
    
    # Offload heavy CPU matrix math to a separate worker thread
    scores = await asyncio.to_thread(reranker.predict, pairs)
    
    # Map raw logits to probabilities via sigmoid and filter
    filtered_docs = [
        doc for doc, score in zip(documents, scores)
        if (1 / (1 + 2.71828 ** (-score))) >= threshold
    ]
    return filtered_docs

We evaluated the two-stage pipeline against raw dense vector retrieval on an annotated benchmark of 200 Pakistani appellate queries:

Retrieval Configuration Hit@1 MRR@3 Discrimination Gap
Dense ANN Alone (Qdrant k=8) 31.5% 0.383 0.540
Dense ANN + Cross-Encoder (Threshold 0.40) 37.5% 0.428 0.780
Relative Improvement +19.0% +11.7% +44.4%

Filtering candidates at a 0.40 sigmoid probability threshold pruned noisy citations while expanding the discrimination gap between authentic and hallucinated legal authorities to 0.780.

3. Real-Time Streaming via Server-Sent Events

Legal synthesis takes several seconds to process. To keep client interfaces interactive, the service provides an /analyze/stream endpoint that consumes LangGraph’s astream(stream_mode="updates") and emits granular event frames:

event: case_start
data: {"case_id": "c-8921", "jurisdiction": "criminal"}

event: node_complete
data: {"node": "retriever", "precedents_found": 4, "elapsed_ms": 312}

event: node_complete
data: {"node": "reasoner", "status": "synthesized"}

event: done
data: {"status": "success"}

Preventing Event Loop Starvation with asyncio.to_thread

In initial deployments, cross-encoder matrix operations caused the FastAPI event loop to freeze for 400–600ms. Because Python’s asyncio runs on a single thread by default, heavy linear algebra in PyTorch halted active SSE connections and triggered reverse-proxy gateway timeouts on Hugging Face Spaces.

By offloading reranker.predict to asyncio.to_thread(), CPU-bound tensor operations run inside worker thread pools, leaving the main loop free to flush SSE frames with X-Accel-Buffering: no headers.

4. Fail-Closed Guardrails (0.990 F1)

In legal AI, failing open (returning an erroneous answer when uncertain) is catastrophic. InsafDost implements fail-closed classification boundaries:

  • If an inbound prompt asks for illegal tactical advice, attempts prompt injection, or provides non-legal scenarios, the guardrail node sets is_valid = False and aborts graph progression before expensive LLM calls occur.
  • On our adversarial test set, the guardrail achieved a 0.990 F1-score with 100% adversarial rejection, ensuring that out-of-scope or manipulative inputs are rejected within 120ms.

Summary

High-stakes domain-specific RAG requires rigorous systems engineering over naive prompting:

  1. Dynamic Model Routing optimizes cost and latency across heterogeneous task tiers.
  2. Cross-Encoder Reranking provides double-digit accuracy leaps on complex domain terminology.
  3. Thread-Offloaded Async Streaming keeps event-driven user interfaces responsive under heavy tensor workloads.

The backend service is available on GitHub.