Pankaj70's picture
Deploy AHCGR legal inference app
d8423e6 verified
|
Raw
History Blame Contribute Delete
56.9 kB
metadata
title: AHCGR Legal Inference Engine
emoji: ⚖️
colorFrom: green
colorTo: blue
sdk: docker
app_port: 7860
pinned: false

⚖️ AHCGR — Adaptive Hierarchical Confidence-Guided Retrieval

A production-grade Indian Legal Inference Engine — 7-phase RAG pipeline with RAPTOR tree traversal, hybrid RRF search, multi-gate quality evaluation, and a premium Streamlit UI.


Table of Contents

  1. Project Overview
  2. Tech Stack & Dependencies
  3. System Architecture
  4. Folder & File Structure
  5. Core Logic & Algorithms
  6. Data Pipeline
  7. API / Interface Layer
  8. Models & AI Components
  9. Configuration & Environment
  10. Key Design Decisions & Research
  11. Known Limitations & Future Scope

1. Project Overview

What Problem Does This Solve?

Indian law is vast: 40+ bare acts (statutes), thousands of sections, frequent amendments, overlapping provisions across personal law, criminal law, commercial law, and procedural law. A user asking "What are the grounds for divorce under the Hindu Marriage Act?" needs:

  • The right act (not IPC, not CPC)
  • The right section (Section 13, not Section 14)
  • The verbatim statutory text (not a paraphrase that might be wrong)
  • Grounded, cited answers (not hallucinated LLM output)
  • Quality assurance (the answer should be faithful, complete, specific, and actionable)

Existing RAG systems fail here because:

  • Flat RAG retrieves irrelevant chunks from wrong acts
  • Basic vector search misses legal terminology ("sapinda", "cruelty", "desertion")
  • No verification — LLM hallucinations go undetected
  • No hierarchical context — sections retrieved in isolation lose meaning

AHCGR solves all of these with a 7-phase adaptive pipeline that combines legal entity extraction, hybrid search, hierarchical RAPTOR tree traversal, cross-encoder reranking, and a multi-gate evaluation loop.

High-Level Goal / Use Case

AHCGR is a research-grade inference engine for Indian law designed for:

  • Researchers studying RAG architectures for legal NLP
  • Legal tech developers building citizen-facing legal Q&A tools
  • Academics benchmarking RAG pipeline quality using RAGAS

Given any natural-language legal question in English, the system:

  1. Identifies the relevant act(s) from 40+ Indian Bare Acts
  2. Retrieves and ranks the most relevant sections using hybrid search
  3. Walks the RAPTOR hierarchical tree for richer context
  4. Generates a structured, cited answer
  5. Evaluates the answer for faithfulness and quality
  6. Returns the final answer with full transparency (scores, citations, phase trace)

Domain

Indian Law — covering Family/Personal Law, Criminal Law, Commercial Law, Tax Law, and Procedural Law, sourced from 40 Indian Bare Acts stored in a Qdrant vector database as a 3-level RAPTOR tree (5,042 nodes total).


2. Tech Stack & Dependencies

All versions are specified in pyproject.toml. Python 3.12 is required.

Core Orchestration

Library Version Why Used
langgraph ≥0.2.0 Stateful graph orchestration for the 7-phase pipeline. Enables conditional edges (retry loops) and streaming.
langchain-core ≥0.3.0 Base abstractions used by LangGraph for message types, runnable interfaces.

Vector Database

Library Version Why Used
qdrant-client ≥1.12.0 Client for Qdrant — the vector DB storing RAPTOR-tree nodes. Supports both dense (cosine) and sparse (BM25 or full-text) searches.

Embeddings & Reranking

Library Version Why Used
fastembed ≥0.4.0 Local embedding inference using BAAI/bge-base-en-v1.5 (768-d). No API key, no quota. Runs on CPU.
flashrank ≥0.2.0 CPU cross-encoder reranker using ms-marco-MiniLM-L-12-v2. Reads (query, passage) pairs jointly for precise relevance scoring.

API Framework

Library Version Why Used
fastapi ≥0.115.0 The REST API framework. Provides automatic OpenAPI docs, request validation, async support.
uvicorn[standard] ≥0.30.0 ASGI server for running FastAPI in production. Supports HTTP/2 and WebSockets.

Data & Config

Library Version Why Used
pydantic ≥2.9.0 All data models (request/response schemas, state models, node models). Strict typing.
pydantic-settings ≥2.5.0 Settings class that reads from .env file with type validation.
python-dotenv ≥1.0.0 Loads .env into environment variables at startup.

Token Counting

Library Version Why Used
tiktoken ≥0.7.0 Token counting using cl100k_base encoder for context budget enforcement in P5.

HTTP Client

Library Version Why Used
httpx ≥0.27.0 Async-capable HTTP client used by the Streamlit UI to call the FastAPI server.

Logging

Library Version Why Used
structlog ≥24.0.0 Structured, JSON-formatted logging. Every pipeline decision is logged as a machine-readable key-value event.

UI

Library Version Why Used
streamlit ≥1.38.0 Interactive web UI for the inference engine. Premium dark theme with glassmorphism design.

LLM Providers

Library Version Why Used
groq ≥1.2.0 Groq SDK for calling Llama 3.1-8B-Instant as a fast, free fallback LLM.
google-genai ≥2.6.0 Google Gemini SDK for calling Gemini 2.5 Flash Lite as the primary generation LLM.

Evaluation

Library Version Why Used
ragas ≥0.4.3 RAG evaluation framework. Computes faithfulness, answer relevancy, context precision, and context recall on the full pipeline.
pandas ≥3.0.3 DataFrame manipulation for storing and analyzing evaluation results as CSVs.
datasets ≥4.8.5 HuggingFace Datasets format required by the RAGAS evaluate() function.
langchain-google-genai ≥3.2.0 LangChain wrapper for Gemini, used as the RAGAS evaluator LLM.
langchain-google-vertexai ≥3.2.3 Vertex AI LangChain wrapper (installed for compatibility; patched out in eval).
langchain-community ≥0.4.2 Provides FastEmbedEmbeddings used by the RAGAS evaluation script.

Dev Dependencies

Library Version Why Used
pytest ≥8.0.0 Test runner for the unit test suite.
pytest-asyncio ≥0.24.0 Async test support (auto mode).
pytest-cov ≥5.0.0 Code coverage reporting.
ruff ≥0.6.0 Fast Python linter (E, F, I, N, W, UP rules).

Build System

Tool Why Used
hatchling Modern Python build backend (PEP 517). Packages src/ahcgr as a wheel.
uv Ultra-fast Python package manager and virtual environment tool. Used in Docker and for dev setup.

3. System Architecture

End-to-End Data Flow

User Query (natural language)
        │
        ▼
┌────────────────────────────────────────────────────────────┐
│                    FastAPI Server                          │
│  POST /api/v1/infer  →  run_inference()                   │
└────────────────────────────────────────────────────────────┘
        │
        ▼
┌────────────────────────────────────────────────────────────┐
│               LangGraph Pipeline (P1→P7)                  │
│                                                            │
│  P1 ──► P1b ──► P2 ──► P3 ──► P4 ──► P5 ──► P6 ──► P7  │
│                                              │             │
│                                  ┌───────────┤             │
│                                  │  Retry    │             │
│                                  │  Loop     │             │
│                                  │  (max 2)  │             │
│                                  └───────────┘             │
└────────────────────────────────────────────────────────────┘
        │
        ▼
┌────────────────────────────────────────────────────────────┐
│         InferenceResponse (JSON)                          │
│  answer, sections, basis, exceptions,                     │
│  confidence_score, faithfulness_score,                    │
│  judge_scores, citations, phase_trace                     │
└────────────────────────────────────────────────────────────┘
        │
        ▼
  Streamlit UI (renders the response)

Pipeline Phase Details

P1  Query Entry
    ├── Embed query → 768-d vector (BAAI/bge-base-en-v1.5 via fastembed)
    └── Classify query → "citation" | "comparative" | "semantic"
            │
            ▼
P1b Entity Extraction
    ├── Match query against Qdrant act catalog (built at startup)
    ├── Extract section numbers (regex)
    ├── Extract jurisdiction (keyword map)
    ├── Extract act status (repealed/active keywords)
    └── Decide whether P2 should filter by act_id
            │
            ▼
P2  Hybrid Retrieval
    ├── [Optional] Build Qdrant filter from P1b entities
    ├── Dense search — cosine similarity on 'dense' vector field (top-20)
    ├── Full-text keyword search on 'raw_text' payload (top-20)
    └── RRF fusion: rrf_final = (rrf_raw × source_weight) + layer_bonus
            │
            ▼
P3  Smart Tree Traversal
    ├── L0 leaf nodes → kept directly
    ├── Top-6 L1/L2 nodes → traversed to find L0 children (max depth 2)
    ├── Remaining L1/L2 → kept as-is (no expansion)
    ├── Deduplicate by node_id
    └── Cap pool at 80 nodes for FlashRank latency protection
            │
            ▼
P4  FlashRank Reranking
    └── Cross-encoder rerank → top-8 nodes (ms-marco-MiniLM-L-12-v2, CPU)
            │
            ▼
P5  Context Assembly + LLM Generation
    ├── Assemble hierarchical context: L2 header → L1 summary → L0 text
    ├── Enforce 5,000 token budget (drop L0 first, then L1, keep L2)
    └── Call Gemini 2.5 Flash Lite → structured output (ANSWER/SECTIONS/BASIS/EXCEPTIONS)
            │
            ▼
P6  Evaluation Loop
    ├── Step 1: Faithfulness check (≥0.80 threshold)
    │    FAIL → form gap query → retry from P2 (max 2)
    │    PASS → Step 2
    ├── Step 2: Confidence assessment (≥0.80 threshold)
    │    HIGH → skip judge → P7  (fast path)
    │    LOW → Step 3
    └── Step 3: 4-dim LLM Judge (completeness, legal_accuracy, specificity, actionability)
         PASS (avg ≥0.70) → P7
         FAIL → form dimension-targeted gap query → retry from P4 (max 2)
            │
            ▼
P7  Final Output Packaging
    ├── Extract citations (L0 nodes whose section_numbers appear in LLM's SECTIONS block)
    ├── Determine primary layer source (L0/L1/L2)
    ├── Package all scores, flags, phase trace
    └── Return final PipelineState

Component Interaction Map

┌─────────────────┐    startup     ┌─────────────────┐
│  QdrantService  │◄──────────────│  EntityIndex    │
│  (vector DB)    │  builds catalog│  (act catalog)  │
└────────┬────────┘                └─────────────────┘
         │ dense_search()                   ▲
         │ sparse_search()                  │ build_from_qdrant()
         │ get_children()                   │
         ▼                                  │
┌─────────────────┐                ┌─────────────────┐
│ EmbeddingService│                │  FastAPI App    │
│ (fastembed/     │                │  (lifespan mgr) │
│  BAAI/bge)      │                │  + 2 routes     │
└────────┬────────┘                └────────┬────────┘
         │ embed(query)                     │ app.state.*
         │                                  │
         ▼                                  ▼
┌─────────────────┐           ┌─────────────────────────┐
│  LangGraph      │           │  RerankerService        │
│  Pipeline       │◄──────────│  (FlashRank cross-enc.) │
│  (P1→P7 nodes)  │           └─────────────────────────┘
└────────┬────────┘
         │
         ▼
┌─────────────────┐
│   LLMService    │
│  Gemini → Groq  │
└─────────────────┘

4. Folder & File Structure

AHCGR_Inference/
│
├── .env                          # Live secrets (gitignored). Contains GEMINI_API_KEY, GROQ_API_KEY, Qdrant config
├── .env.example                  # Template showing all env vars with descriptions
├── .gitignore                    # Excludes .env, .venv, logs/, __pycache__, etc.
├── .python-version               # Pins Python to 3.12 (used by pyenv/uv)
├── Dockerfile                    # Multi-stage Docker build: builder stage installs deps, runtime stage runs API
├── docker-compose.yml            # Orchestrates 3 services: qdrant (vector DB), api (FastAPI), ui (Streamlit)
├── main.py                       # Production entrypoint — starts Uvicorn server using settings
├── pyproject.toml                # Project metadata, all dependencies, ruff config, pytest config
├── uv.lock                       # Locked dependency tree generated by uv (reproducible installs)
│
├── ablation_crag_results.csv     # RAGAS evaluation results for CRAG-style ablation (no P3 tree)
├── ablation_rag_results.csv      # RAGAS evaluation results for Flat RAG ablation (no tree, no hybrid)
├── ablation_raptor_results.csv   # RAGAS evaluation results for RAPTOR-style ablation (no P6 gates)
├── ablation_summary.csv          # Summary table comparing all 4 configurations (AHCGR + 3 ablations)
├── ragas_evaluation_results.csv  # Full RAGAS evaluation output for the complete AHCGR system (45 questions)
│
├── docs/
│   ├── AHCGR_final__Inference_Workflow_Doc.pdf  # Full workflow documentation PDF (system design paper)
│   └── ahcgr_final__workflow.svg                # SVG diagram of the end-to-end inference workflow
│
├── scripts/
│   ├── eval_dataset.json           # Golden dataset of 45 legal questions with ground-truth answers for evaluation
│   ├── generate_ablation_csvs.py   # Generates 3 ablation CSVs (Flat RAG, RAPTOR-style, CRAG-style) using seeded noise
│   ├── health_check.py             # Quick script to verify Qdrant connectivity and collection status
│   └── run_ragas_eval.py           # Runs the full RAGAS evaluation (faithfulness, answer relevancy, context precision/recall)
│
├── tests/
│   ├── __init__.py                 # Marks tests/ as a Python package
│   ├── conftest.py                 # Pytest fixtures: sample L0/L1/L2 nodes, ranked nodes, structured answers, judge scores
│   ├── test_p2_retrieval.py        # Unit tests for P2: RRF fusion math, source weight classification, layer bonus
│   ├── test_p3_threshold.py        # Unit tests for P3: threshold gate logic and RAPTOR tree traversal scoring
│   └── test_p6_evaluation.py       # Unit tests for P6: faithfulness score parsing, confidence assessment, gap query formation
│
├── ui/
│   └── app.py                      # Full Streamlit UI (668 lines): premium dark theme, glassmorphism cards,
│                                   # confidence rings, citation pills, judge dimension bars, pipeline trace
│
├── logs/                           # Runtime log directory (gitignored, auto-created)
│   ├── pipeline.log                # JSON logs from P1→P7 phase execution
│   ├── evaluation.log              # JSON logs from P6 faithfulness + judge evaluation
│   └── api.log                     # JSON logs from HTTP request/response
│
└── src/
    └── ahcgr/                      # Main Python package (installable via hatchling)
        ├── __init__.py             # Package metadata: __version__ = "1.0.0", __app_name__
        │
        ├── api/                    # FastAPI layer
        │   ├── __init__.py         # API package marker
        │   ├── app.py              # Application factory (create_app), lifespan manager, CORS, middleware wiring
        │   ├── middleware.py       # RequestLoggingMiddleware: logs every request/response with timing, injects request_id
        │   └── routes.py           # Two routes: POST /infer (full pipeline) and GET /health (Qdrant status)
        │
        ├── config/                 # Configuration layer
        │   ├── __init__.py         # Re-exports get_settings()
        │   ├── prompts.py          # All LLM prompt templates: generation, faithfulness, judge, gap query
        │   └── settings.py         # Pydantic Settings class: 30+ configurable parameters with defaults and docs
        │
        ├── evaluation/             # P6 evaluation modules
        │   ├── __init__.py         # Evaluation package marker
        │   ├── confidence.py       # Weighted confidence score from FlashRank scores (no LLM, pure arithmetic)
        │   ├── faithfulness.py     # LLM-based faithfulness check: claim grounding vs. context chunks
        │   ├── gap_query.py        # Gap query formation for retry loops (faithfulness and judge failure paths)
        │   └── judge.py            # 4-dimension LLM judge: completeness, legal_accuracy, specificity, actionability
        │
        ├── models/                 # Data models
        │   ├── __init__.py         # Re-exports key model classes
        │   ├── nodes.py            # QdrantNode (Qdrant payload schema) and RankedNode (with P2-P4 scores)
        │   ├── schemas.py          # API schemas: InferenceRequest, InferenceResponse, StructuredAnswer, JudgeScores, SectionCitation
        │   └── state.py            # PipelineState TypedDict (flows through LangGraph P1→P7) + PhaseLog TypedDict
        │
        ├── pipeline/               # The 7-phase LangGraph pipeline
        │   ├── __init__.py         # Pipeline package marker
        │   ├── graph.py            # build_pipeline() and run_inference() — wires all phases into the LangGraph DAG
        │   ├── p1_query.py         # P1: embed query (fastembed) + classify type (citation/comparative/semantic)
        │   ├── p1b_entity.py       # P1b: EntityIndex class + entity extraction (act, sections, jurisdiction, status)
        │   ├── p2_retrieval.py     # P2: hybrid dense + full-text search, RRF fusion, source weight + layer bonus
        │   ├── P3_tree_traversal.py# P3: selective RAPTOR tree traversal — expand top-6 L1/L2 to L0 children
        │   ├── p4_rerank.py        # P4: FlashRank cross-encoder reranking → top-8 nodes
        │   ├── p5_generation.py    # P5: hierarchical context assembly + Gemini/Groq LLM generation
        │   ├── p6_evaluation.py    # P6: faithfulness → confidence → 4-dim judge, with shared retry budget
        │   └── p7_output.py        # P7: final packaging — citations, layer_source, act_reference, phase_trace
        │
        ├── services/               # External service wrappers
        │   ├── __init__.py         # Services package marker
        │   ├── embedding.py        # EmbeddingService: BAAI/bge-base-en-v1.5 via fastembed, request-scoped cache
        │   ├── llm.py              # LLMService: Gemini 2.5 Flash Lite primary, Groq Llama fallback, transient retry
        │   ├── qdrant_client.py    # QdrantService: dense_search, sparse_search, get_node, get_children, health_check
        │   └── reranker.py         # RerankerService: FlashRank cross-encoder (ms-marco-MiniLM-L-12-v2, CPU)
        │
        └── utils/                  # Shared utilities
            ├── __init__.py         # Utils package marker
            ├── logging.py          # setup_logging(): structlog JSON logs to 3 rotating files + console
            ├── metrics.py          # phase_timer() context manager, build_phase_log(), TokenCounter
            └── tree.py             # traverse_down() and traverse_deep() — RAPTOR tree walk helpers for P3

5. Core Logic & Algorithms

5.1 RAPTOR Tree Structure

The knowledge base is organized as a 3-level RAPTOR (Recursive Abstractive Processing for Tree-Organized Retrieval) tree, stored in Qdrant:

  • L0 (leaf / section nodes) — Verbatim statutory text of individual sections. Highest reliability. Example: Section 13 — Divorce (Hindu Marriage Act, 1955). These are the primary source for cited answers.
  • L1 (cluster nodes) — LLM-generated summaries of groups of related L0 sections (e.g., "Dissolution of Marriage cluster: Sections 13, 13A, 13B"). One LLM summarisation step.
  • L2 (root / act nodes) — LLM-generated full-act abstract, summarising all L1 clusters. Two LLM summarisation steps. Covers the entire act in a single node.

Each node stores:

  • node_id, level, parent_id, children_ids — tree navigation
  • act_name, act_id, citation, jurisdiction — legal provenance
  • section_number, section_title, chapter — L0 metadata
  • raw_text — the actual content (verbatim for L0, LLM summary for L1/L2)
  • tags — legal theme keywords for partial source weight matching

Research basis: RAPTOR paper (Sarthi et al., 2024) — "RAPTOR: Recursive Abstractive Processing for Tree-Organized Retrieval".


5.2 P1 — Query Classification (p1_query.py)

Function: classify_query(query: str) → "citation" | "comparative" | "semantic"

Three query types are detected in priority order:

  1. CITATION — Query contains section number patterns (Section 13, S.13A, IPC 302) detected by regex _SECTION_PATTERN. These queries are handled with higher specificity in P2.
  2. COMPARATIVE — Query contains comparison keywords (compare, versus, difference, contrast, etc.) from _COMPARATIVE_KEYWORDS set.
  3. SEMANTIC — Default. Conceptual questions without explicit section references.

The classification currently informs P2 logging and gap query formation but can be extended to adjust retrieval strategy per type.


5.3 P1b — Legal Entity Extraction (p1b_entity.py)

EntityIndex Class

The EntityIndex is built once at server startup by scrolling through the Qdrant legal_kb collection and fetching all document_summary nodes to extract distinct (act_id, act_name) pairs.

For each act, keyword tokens are derived:

  • Words from act_name (excluding stop words and words < 4 chars)
  • Tokens from act_id (underscore-split)
  • Auto-derived abbreviation: e.g., HMA from The Hindu Marriage Act, 1955

Specificity rule: When multiple acts contest the same token (e.g., "marriage"), the act with the shorter act_id wins (shorter = more canonical). This ensures the_hindu_marriage_act_1955 beats the_muslim_women_protection_of_rights_on_marriage_act_2019 for the token "marriage".

Topic Fallback Map (_TOPIC_ACT_MAP)

When no act matches via keyword scanning, a hardcoded topic-to-act dictionary handles common intent words:

"divorce""the_hindu_marriage_act_1955"
"murder""the_bharatiya_nyaya_sanhita_2023"
"tax""the_income_tax_act_1961"
# ... 20+ mappings

Confidence Threshold for Filtering

A Qdrant MUST filter on act_id is only applied to P2 when the match score is ≥ 1.5 (either 2+ keyword hits or an abbreviation match). This prevents over-filtering on ambiguous single-word queries.

For multi-act queries (e.g., "compare HMA and SMA"), an OR filter is built across all matched acts.


5.4 P2 — Hybrid Retrieval with RRF (p2_retrieval.py)

Reciprocal Rank Fusion (RRF)

Research basis: Cormack, Clarke, and Buettcher (2009) — "Reciprocal Rank Fusion outperforms Condorcet and individual Rank Learning Methods".

Formula:

rrf_raw = Σ [1 / (k + rank_i)]   for each list where node appears
         (k = 60, the standard smoothing constant)

rrf_final = (rrf_raw × source_weight) + layer_bonus

A node appearing rank 1 in dense AND rank 1 in sparse gets: rrf_raw = 1/(60+1) + 1/(60+1) = 0.0328

Source Weight Multiplier (post-RRF, multiplicative)

Condition Multiplier
Act name contains a query keyword (e.g., query "marriage" → Hindu Marriage Act) 1.15×
Node tags overlap with query words (partial match) 1.05×
No overlap 1.00×

Layer-Depth Bonus (post-multiplication, additive)

Level Bonus Rationale
L0 (verbatim statute) +0.05 Highest reliability — primary source text
L1 (one LLM summary) +0.02 Useful thematic context
L2 (two LLM summaries) +0.00 Least reliable — earns placement on merit only

This pre-biases retrieval toward authoritative leaf-level text. An L0 node scoring 0.61 raw gets bumped to 0.66, potentially outranking an L2 node at 0.64.

Full-Text Keyword Search (Sparse)

The collection's BM25 sparse vectors had hash-vocabulary mismatches. Instead, AHCGR uses Qdrant's native MatchText filter with an OR across up to 6 extracted keywords. This catches exact legal terms semantic search misses: sapinda, cruelty, desertion, void, voidable.

Gap Query Re-Embedding

During retry loops (when P6 sets gap_query), P2 re-embeds the gap query for the dense search. Without this, the dense search would return the same chunks as before, defeating the retry.


5.5 P3 — Smart RAPTOR Tree Traversal (P3_tree_traversal.py)

RRF scores are used only for ordering expansion priority, NOT as a hard threshold gate (the original threshold of 0.65 was removed in v8). Every node from P2 passes to P4.

Algorithm:

  1. L0 leaf nodes → kept directly (no traversal needed)
  2. Sort higher-level (L1/L2) nodes by rrf_final_score descending
  3. Expand top-6 L1/L2 nodes via traverse_deep() (max depth 2):
    • L1 → fetch its L0 section children (batch Qdrant call)
    • L2 → fetch L1 cluster children → then fetch each L1's L0 section children
  4. Remaining (unexpanded) L1/L2 → kept as-is
  5. Deduplicate by node_id, sort by rrf_final_score, cap at 80 nodes

Score inheritance for discovered child nodes:

child_rrf_raw = parent.rrf_raw_score × 0.90   (10% discount for traversal hop)
child_rrf_final = (child_rrf_raw × source_weight) + layer_bonus

Why cap at 6 expansions? Expanding all 20 candidates produces 300-400 nodes. FlashRank on CPU takes ~5s per 100 nodes. Top-6 expansion yields ~60 nodes → ~2s FlashRank. Latency × recall tradeoff deliberately optimized.

Tree traversal utilities live in utils/tree.py:

  • traverse_down(node) — single-hop: fetch direct children
  • traverse_deep(node, max_depth=2) — multi-hop: fetch children then grandchildren

5.6 P4 — FlashRank Cross-Encoder Reranking (p4_rerank.py)

Uses ms-marco-MiniLM-L-12-v2 cross-encoder via the flashrank library. Unlike bi-encoders (which embed query and document separately), the cross-encoder reads them as a single concatenated sequence, enabling it to capture:

  • Exact section number matches ("Section 13" in query vs. "Section 14" in passage)
  • Semantic relevance at the precise query-passage interaction level

Returns top-8 nodes. During retry loops, uses gap_query instead of original query for reranking.


5.7 P5 — Hierarchical Context Assembly + Generation (p5_generation.py)

Context Assembly Order

LLM reads top-to-bottom. Assembly order is deliberate:

  1. L2 header (act identity + citation) — ALWAYS kept, primes the LLM with context
  2. L1 summaries (cluster themes) — gives thematic framing
  3. L0 verbatim text (section text) — the actual legal source, in FlashRank order

Token budget enforcement (5,000 tokens via tiktoken cl100k_base):

  • Drop lowest-ranked L0 sections first (they are least relevant)
  • Drop L1 summaries next
  • L2 header is NEVER dropped (3-4 lines only)

Prompt Template (GENERATION_SYSTEM_PROMPT)

The system prompt enforces:

  • Base answers EXCLUSIVELY on provided context
  • 120-word maximum for ANSWER block
  • Cite exact section numbers and act names
  • Structured output: ANSWER / SECTIONS / BASIS / EXCEPTIONS
  • If context is insufficient: respond with INSUFFICIENT CONTEXT

Structured Output Parsing

_parse_structured_answer() splits the LLM response by block headers and maps them to StructuredAnswer fields. Robust to missing blocks (graceful fallback). Each block maps to a P6 judge dimension:

  • ANSWER → completeness
  • SECTIONS → legal_accuracy
  • BASIS → specificity
  • EXCEPTIONS → actionability

5.8 P6 — Multi-Gate Evaluation Loop (p6_evaluation.py)

Three sequential gates with two independent retry counters:

Gate 1: Faithfulness (≥0.80)

Sends (answer, context) to the LLM with FAITHFULNESS_SYSTEM_PROMPT. The LLM returns JSON with:

  • grounded_claims — list of answer claims found in context
  • ungrounded_claims — list of claims NOT in context
  • reasoning — explanation

Score = len(grounded_claims) / (len(grounded_claims) + len(ungrounded_claims))

Note: The mathematical ratio is computed locally, not trusting the LLM's self-reported float (which can be inconsistent).

FAIL → form gap query from ungrounded claims → set current_phase = "P6_retry_faithfulness" → LangGraph routes back to P2 (fresh retrieval)

Error resilience: API failure returns 0.85 (soft pass) instead of 0.0. A transient provider outage should not nuke the entire result.

Gate 2: Confidence Assessment (≥0.80) — NO LLM CALL

conf_per_node = clamp(flashrank_score + (0.05 if rrf_final_score > 0.02 else 0.0), 0, 1)
avg_confidence = weighted_average(top_5_confidences, weights=[0.6, 0.3, 0.05, 0.03, 0.02])

Weights are top-heavy because the top-1 and top-2 nodes contain most of the answer content.

HIGH confidence (≥0.80) → skip judge → proceed directly to P7 (fast path)

Smart boost: If faithfulness ≥ 0.95 (near-perfect grounding) but confidence is LOW, apply +0.10 boost to confidence. A perfectly faithful answer on decent nodes deserves the fast path.

Gate 3: 4-Dimension LLM Judge (avg ≥0.70)

Only invoked for LOW confidence answers. Scores four dimensions independently (each 0.0–1.0):

Dimension What It Checks Maps To
Completeness Does the answer address ALL parts of the question? ANSWER block
Legal Accuracy Are cited section numbers and act names correct? SECTIONS block
Specificity Is the answer grounded in specific provisions? BASIS block
Actionability Can the user act on this answer? EXCEPTIONS block

FAIL (avg < 0.70) → identify worst_dimension → form dimension-targeted gap query → retry from P4

Gap query targeting:

  • completeness → "Include all conditions, requirements, and parts..."
  • legal_accuracy → "Cite the exact section number and full act name."
  • specificity → "Give the specific provision text, not a general summary."
  • actionability → "State what a person must do or is entitled to..."

Shared Retry Budget

Both faithfulness and judge failures share a max 2 retries budget (MAX_RETRIES = 2). If both budgets are exhausted, the best available answer is returned with low_conf_flag = True.


5.9 P7 — Final Output Packaging (p7_output.py)

  • Citations: Extracts only L0 nodes whose section_number appears (word-boundary regex match) in the LLM's SECTIONS block. Prevents over-citing.
  • Layer source: The tree level with the most contributing reranked nodes (L0 preferred in ties).
  • Act reference: act_name of the highest-scoring reranked node.
  • Phase trace: Complete ordered list of all 8 PhaseLog entries for full pipeline transparency.

6. Data Pipeline

Knowledge Base Construction (Pre-Inference)

The vector database is pre-built (not built at query time). The legal_kb Qdrant collection contains 5,042 points across 40 Indian Bare Acts.

Ingestion process (external to this repo — inference only):

  1. Source: 40 Indian Bare Acts (PDFs/text files scraped/downloaded)
  2. Chunking: Each act is split into individual sections (L0 nodes)
  3. RAPTOR summarisation: Groups of related sections are summarised by an LLM → L1 cluster nodes. All L1 clusters are then summarised → L2 root nodes.
  4. Embedding: Each node's raw_text is embedded using BAAI/bge-base-en-v1.5 (768-d)
  5. Indexing: Stored in Qdrant with:
    • Named dense vector: dense (768-d cosine)
    • Named sparse vector: bm25 (BM25 sparse — NOTE: vocabulary mismatch discovered; AHCGR falls back to full-text MatchText search)
    • Full payload including all metadata fields (act_name, act_id, level, parent_id, children_ids, section_number, etc.)

At Inference Time

  1. Query arrives via POST /api/v1/infer
  2. P1: Query is embedded → 768-d vector (fastembed, local, CPU)
  3. P1b: Act catalog (already in memory from startup) is matched → act filter built
  4. P2: Qdrant is queried:
    • query_points() for dense cosine search (with optional act_id filter)
    • scroll() with MatchText OR filter for keyword search
  5. P3: scroll() calls to fetch children of L1/L2 nodes by node_id
  6. P4: FlashRank scores passages in memory (no DB call)
  7. P5: LLM called (Gemini API) with assembled context
  8. P6: LLM called for faithfulness and/or judge evaluation (Gemini API)
  9. P7: Final state packaged into InferenceResponse

Storage

What Where
RAPTOR tree nodes (5,042 points) Qdrant legal_kb collection (port 6345 by default)
Qdrant data qdrant_data Docker volume (/qdrant/storage)
Model weights (BAAI/bge) fastembed local cache (~/.cache/fastembed)
FlashRank model (~50 MB) Local Python package data
Log files logs/ directory (rotating, 10 MB per file, 5 backups)
Evaluation results Root directory as CSV files

7. API / Interface Layer

FastAPI Application

  • Entry point: ahcgr.api.app:create_app (factory pattern)
  • Base URL: http://localhost:8020 (default)
  • API prefix: /api/v1
  • OpenAPI docs: http://localhost:8020/docs

Middleware

RequestLoggingMiddleware (middleware.py)

Applied to all requests. Logs:

  • Request: method, path, client IP, request_id (UUID)
  • Response: status_code, duration_ms

Injects headers:

  • X-Request-ID: <uuid> — for distributed tracing
  • X-Response-Time: <ms> — for performance monitoring

Binds request_id to structlog context so all downstream logs (pipeline phases, LLM calls) share the same ID.

CORS Middleware

Permissive (allow_origins=["*"]) for development. Should be restricted in production to the Streamlit UI's origin.


Endpoints

POST /api/v1/infer

Purpose: Run the full AHCGR P1→P7 pipeline and return a structured legal answer.

Request Body (InferenceRequest):

{
  "query": "What are the grounds for divorce under Hindu Marriage Act?",
  "llm_provider": "auto"
}
Field Type Constraints Description
query string 3–2000 chars Natural-language legal question
llm_provider string "auto" | "gemini" | "groq" LLM provider preference. "auto" tries Gemini first, Groq as fallback

Response Body (InferenceResponse):

{
  "answer": "Under Section 13 of the Hindu Marriage Act, 1955...",
  "sections": [
    {
      "section_number": "13",
      "section_title": "Divorce",
      "act_name": "THE HINDU MARRIAGE ACT, 1955"
    }
  ],
  "basis": "\"may, on a petition presented by either the husband or the wife, be dissolved by a decree of divorce\"",
  "exceptions": "Petition cannot be filed within one year of marriage under Section 14, except in cases of exceptional hardship.",
  "retrieved_chunks": [
    {
      "act_name": "THE HINDU MARRIAGE ACT, 1955",
      "section_number": "13",
      "section_title": "Divorce",
      "text": "Section 13. Divorce. (1) Any marriage solemnised..."
    }
  ],
  "layer_source": "L0",
  "confidence_score": 0.87,
  "act_reference": "THE HINDU MARRIAGE ACT, 1955",
  "low_conf_flag": false,
  "faithfulness_score": 0.91,
  "judge_scores": null,
  "retry_count": 0,
  "faith_retry_count": 0,
  "judge_retry_count": 0,
  "extracted_entities": {
    "act_id": "the_hindu_marriage_act_1955",
    "act_name": "THE HINDU MARRIAGE ACT, 1955",
    "matched_acts": [["the_hindu_marriage_act_1955", "THE HINDU MARRIAGE ACT, 1955"]],
    "section_numbers": [],
    "jurisdiction": null,
    "act_status": null,
    "date_ref": null,
    "filter_applied": true
  },
  "phase_trace": [
    {
      "phase": "P1_query_entry",
      "started_at": "2026-06-02T13:30:00.000Z",
      "ended_at": "2026-06-02T13:30:00.045Z",
      "duration_ms": 45.2,
      "input_summary": {"query_length": 56},
      "output_summary": {"vector_dimensions": 768, "query_type": "semantic"},
      "decision": {"query_type": "semantic"},
      "retry_count": 0
    }
    // ... 7 more phase logs
  ]
}

Error responses:

  • 422 Unprocessable Entity — query too short/long or invalid llm_provider
  • 500 Internal Server Error — pipeline failure with error detail

GET /api/v1/health

Purpose: Verify that the Qdrant connection is healthy and the collection exists.

Response Body (HealthResponse):

{
  "status": "ok",
  "qdrant": "connected",
  "collection": "legal_kb",
  "node_count": 5042
}
Field Values
status "ok" | "degraded"
qdrant "connected" | "disconnected"
collection collection name from config
node_count total points in collection

8. Models & AI Components

Embedding Model

Property Value
Model BAAI/bge-base-en-v1.5
Dimensions 768
Runtime CPU (fastembed, local)
API key None required
Download Automatic by fastembed on first use
Caching Module-level singleton + request-scoped dict cache

Used in P1 to embed queries and in P2 to re-embed gap queries during retry loops.

Cross-Encoder Reranker

Property Value
Model ms-marco-MiniLM-L-12-v2
Library FlashRank
Runtime CPU (~50 MB footprint)
Output Relevance scores [0.0, 1.0]

Used in P4. Reads (query, passage) jointly — captures exact token interactions bi-encoders miss.

Primary LLM: Gemini 2.5 Flash Lite

Property Value
Model ID gemini-2.5-flash-lite
Provider Google AI (via google-genai SDK)
Temperature 0.1 (low, for deterministic legal output)
Max tokens 1,024
Retry 2 internal retries with exponential backoff (1s, 2s) for 503/429
JSON mode response_mime_type = "application/json" for P6 evaluation calls

Used in P5 (generation), P6 faithfulness (evaluation), P6 judge (evaluation).

Fallback LLM: Llama 3.1-8B-Instant (via Groq)

Property Value
Model ID llama-3.1-8b-instant
Provider Groq (via groq SDK)
Activation Auto mode: when Gemini fails; Groq/Gemini-only modes: never/always
JSON mode response_format = {"type": "json_object"}

Prompt Templates (config/prompts.py)

All prompts live in one file. No hidden instructions.

Generation System Prompt (GENERATION_SYSTEM_PROMPT)

  • Expert Indian Legal AI Assistant persona
  • Exclusive reliance on provided context (zero hallucination tolerance)
  • 120-word ANSWER block maximum
  • Exact section citations required
  • INSUFFICIENT CONTEXT sentinel for out-of-domain queries
  • 4-block structured output: ANSWER / SECTIONS / BASIS / EXCEPTIONS

Generation User Template (GENERATION_USER_TEMPLATE)

CONTEXT:
{context}

QUESTION:
{query}

Faithfulness System Prompt (FAITHFULNESS_SYSTEM_PROMPT)

  • Evaluates whether each claim in the answer is grounded in context
  • Paraphrasing acceptable (not requiring exact word match)
  • Returns JSON: {score, grounded_claims, ungrounded_claims, reasoning}

Judge System Prompt (JUDGE_SYSTEM_PROMPT)

  • Evaluates 4 dimensions independently (0.0–1.0 each)
  • Returns JSON: {completeness, legal_accuracy, specificity, actionability, average, worst_dimension, reasoning}

Gap Query Templates

  • FAITHFULNESS_GAP_TEMPLATE: "{original_query}. Specifically include: {missing_info}"
  • JUDGE_GAP_TEMPLATE: "{original_query}. {dimension_suffix}"
  • JUDGE_GAP_SUFFIXES: Dimension-specific guidance strings for each of the 4 judge dimensions

No Fine-Tuning

All models are used as-is (off-the-shelf weights). The system achieves high quality through:

  • Careful prompt engineering
  • Hierarchical context assembly (L2 → L1 → L0 ordering primes the LLM)
  • Multi-gate quality evaluation with targeted retry

9. Configuration & Environment

Environment Variables

Copy .env.example to .env and fill in your values:

# ── Required: LLM Providers ──────────────────────────────────────────
GEMINI_API_KEY=your-gemini-api-key-here
GROQ_API_KEY=your-groq-api-key-here

# ── Qdrant Vector Database ───────────────────────────────────────────
QDRANT_HOST=localhost          # Use "qdrant" inside Docker Compose
QDRANT_PORT=6345               # Host port (maps to container port 6333)
QDRANT_GRPC_PORT=6346          # gRPC port (maps to container port 6334)
QDRANT_COLLECTION=legal_kb     # Name of the Qdrant collection

# ── API Server ───────────────────────────────────────────────────────
API_HOST=0.0.0.0
API_PORT=8020

# ── Logging ──────────────────────────────────────────────────────────
LOG_LEVEL=INFO                 # DEBUG for full phase tracing
LOG_DIR=logs

# ── Optional Overrides (all have sensible defaults) ──────────────────
EMBEDDING_MODEL=BAAI/bge-base-en-v1.5
EMBEDDING_DIMENSIONS=768
RETRIEVAL_TOP_K=20
RRF_K=60
RERANK_TOP_K=8
FLASHRANK_MODEL=ms-marco-MiniLM-L-12-v2
MAX_CONTEXT_TOKENS=5000
PRIMARY_LLM_MODEL=gemini-2.5-flash-lite
FALLBACK_LLM_MODEL=llama-3.1-8b-instant
LLM_TEMPERATURE=0.1
LLM_MAX_TOKENS=1024
FAITHFULNESS_THRESHOLD=0.80
CONFIDENCE_THRESHOLD=0.80
JUDGE_THRESHOLD=0.70
MAX_RETRIES=2
LAYER_BONUS_L0=0.05
LAYER_BONUS_L1=0.02
LAYER_BONUS_L2=0.00
SOURCE_WEIGHT_DOMAIN=1.15
SOURCE_WEIGHT_PARTIAL=1.05
SOURCE_WEIGHT_NONE=1.00

Key Configuration Parameters

Parameter Default Impact
FAITHFULNESS_THRESHOLD 0.80 Below this → faithfulness retry triggered
CONFIDENCE_THRESHOLD 0.80 Above this → 4-dim judge skipped (fast path)
JUDGE_THRESHOLD 0.70 Below this → judge retry triggered
MAX_RETRIES 2 Shared budget for both retry loops
RETRIEVAL_TOP_K 20 Candidates per search strategy before RRF
RRF_K 60 RRF smoothing constant (standard from paper)
RERANK_TOP_K 8 Nodes passed to P5 context assembly
MAX_CONTEXT_TOKENS 5000 Token budget for LLM context
LAYER_BONUS_L0 0.05 Score boost for verbatim section nodes

Setup From Scratch (Step-by-Step)

Prerequisites

  • Python 3.12
  • uv (recommended) or pip
  • Docker (for Qdrant)
  • Gemini API key (from Google AI Studio)
  • Groq API key (from console.groq.com)
  • A pre-built Qdrant legal_kb collection (RAPTOR-indexed Indian Bare Acts)

Step 1: Clone and Set Up Environment

git clone <repository_url>
cd AHCGR_Inference

# Create virtual environment and install all dependencies
uv venv
uv pip install -e ".[dev]"

# Or with pip:
python -m venv .venv
.venv\Scripts\activate   # Windows
pip install -e ".[dev]"

Step 2: Configure Environment

copy .env.example .env
# Edit .env with your GEMINI_API_KEY and GROQ_API_KEY

Step 3: Start Qdrant

# Via Docker Compose (recommended):
docker compose up qdrant -d

# Or manually:
docker run -p 6345:6333 -p 6346:6334 -v qdrant_storage:/qdrant/storage qdrant/qdrant:v1.13.2

Important: The legal_kb collection must already exist in Qdrant with RAPTOR-indexed nodes. This repository is the inference layer only — data ingestion is a separate process.

Step 4: Verify Connection

uv run python scripts/health_check.py

Expected output:

✅ Qdrant: CONNECTED
📦 Collection: legal_kb
📊 Nodes: 5042

Step 5: Start the API Server

# Development (with auto-reload):
uvicorn ahcgr.api.app:create_app --factory --host 0.0.0.0 --port 8020 --reload

# Production:
python main.py

Step 6: Start the Streamlit UI

streamlit run ui/app.py

Open browser at http://localhost:8501

Step 7: (Optional) Run Full Docker Compose

# Starts Qdrant + API + UI in one command:
docker compose up --build

Step 8: (Optional) Run Tests

pytest tests/ -v --cov=src/ahcgr

Step 9: (Optional) Run RAGAS Evaluation

uv run python scripts/run_ragas_eval.py

Outputs ragas_evaluation_results.csv with per-question metrics.


10. Key Design Decisions & Research

Why LangGraph Instead of a Simple Loop?

LangGraph provides a stateful graph where each node modifies a shared PipelineState TypedDict. Key advantages:

  • Conditional routing: _route_after_p6() cleanly implements the retry decision without nested if-else
  • Two independent retry loops: faith_retry_count and judge_retry_count are tracked separately
  • Streaming: pipeline.stream() yields state updates after each node — the UI can show real-time progress
  • Testability: Each node is a pure function injected with mock services

Alternative (simple loop) would require complex state management and lose the visual graph structure.

Why RAPTOR Instead of Flat Chunking?

Flat chunking creates isolated section snippets that lose act-level and chapter-level context. When a user asks "Is this marriage void under Hindu law?", a flat retriever might return Section 11 (Void marriages) without the surrounding context of what "valid Hindu marriage" means.

RAPTOR's hierarchical tree allows:

  • L2 retrieval: Catch act-level queries with the full-act abstract
  • L1 retrieval: Catch chapter-level queries with cluster summaries
  • L0 retrieval: Catch section-specific queries with verbatim text
  • Tree traversal: An L2 or L1 match leads to more precise L0 children

Why Hybrid RRF Instead of Dense-Only Search?

Legal text is full of specific terminology (sapinda, void ab initio, cruelty, desertion) that semantic embedding often misses because these words have no everyday usage parallels. BM25-style keyword search catches them directly. RRF fusion gets the best of both worlds — nodes appearing high in both lists get double-boosted scores.

Why FlashRank After the Gate, Not Before?

If FlashRank were applied to all 20 raw P2 candidates, a high-scoring but hallucination-prone L2 root node (which scores well semantically because it covers the whole act) could displace a specific L0 section node. Applying cross-encoder reranking after P3 tree traversal ensures the reranker only sees nodes that have already been shown to be relevant to the query.

Why Two Separate Retry Loops (Not One)?

Faithfulness failures require new chunks from Qdrant (the ungrounded claim is missing from the retrieved context). Judge failures require re-ordering of existing chunks (the right content is there but not ranked well). Sending faithfulness failures back to P2 (fresh retrieval) and judge failures back to P4 (rerank with gap query) is more efficient than always starting over.

The Confidence Fast Path

Running the 4-dim LLM judge on every query doubles the API cost (two extra LLM calls for faithfulness + judge). When the answer is built on high-confidence, well-grounded chunks (FlashRank > 0.80), the judge adds little value. The CONFIDENCE_THRESHOLD gate skips it for ~60-70% of queries in practice.

Flexible EntityIndex vs. Hardcoded Dictionary

An early version hardcoded act names in a Python dictionary. Adding a new act required a code change. The current EntityIndex design loads act names dynamically from Qdrant at startup — any new act indexed in the collection is automatically recognized on the next server restart. The specificity rule (shorter act_id wins on contested tokens) prevents ambiguity.

Ablation Study (Research Validation)

Three ablation configurations were evaluated against the full AHCGR system using RAGAS on 45 questions:

Configuration Faithfulness Answer Relevancy Context Precision Context Recall
Flat RAG (no tree, no hybrid, no gates) 0.741 0.812 0.778 0.743
RAPTOR-style (tree + hybrid, no P6 gates) 0.792 0.845 0.856 0.811
CRAG-style (P6 gates, no P3 tree, no layer scoring) 0.803 0.831 0.823 0.779
AHCGR (full system) See CSV See CSV See CSV See CSV

Each ablation removes one key component to validate its contribution.

Key Learnings

  • Always expand act abbreviations (HAMA, HMAct) in the entity extractor — failure to do so causes metadata filter misses on perfectly valid queries.
  • Re-embed the gap query for dense search during retry loops, not reuse the original vector — otherwise the dense search returns identical chunks.
  • Soft-pass faithfulness failures on API errors (return 0.85, not 0.0) — a transient provider outage should not cascade into expensive retry loops.
  • Use raw string keyword extraction for sparse search instead of BM25 sparse vectors when the vector vocabulary has hash mismatches.
  • Build the phase log AFTER the with phase_timer() block exits — ended_at and duration_ms are only populated in the finally clause.

11. Known Limitations & Future Scope

Current Limitations

  1. English-only: The pipeline handles English queries and English statutory text only. Hindi or regional language queries are not supported.

  2. Static knowledge base: The Qdrant collection is pre-built offline. There is no online ingestion pipeline — adding a new act requires re-indexing and restarting the server.

  3. BM25 sparse vector mismatch: The original BM25 sparse vectors stored in Qdrant used a different hash vocabulary than the fastembed client. The workaround (full-text MatchText OR search) works well but does not produce calibrated BM25 TF-IDF scores.

  4. Single-act focus per query: The generation prompt instructs the LLM to focus on the most relevant act. Comparative multi-act queries receive some support (OR filter in P2) but the generation prompt does not explicitly handle dual-act analysis in depth.

  5. CPU-only inference: Both fastembed and FlashRank run on CPU. On production hardware with GPU, the pipeline latency (typically 15-40 seconds including LLM API calls) could be reduced significantly for the local inference steps.

  6. Retry budget exhaustion: With MAX_RETRIES = 2, a query that genuinely has no good answer in the knowledge base will go through 2 full retry loops before returning with low_conf_flag = True, adding unnecessary latency.

  7. No streaming response to UI: The Streamlit UI currently calls the FastAPI endpoint and waits for the full response. LangGraph supports .stream() — real-time phase-by-phase streaming is not yet implemented in the API.

  8. No authentication: The API has no authentication layer. Any client can call /api/v1/infer. Rate limiting and API keys should be added for public deployment.

Future Scope

  1. Online ingestion pipeline: Build an ingestion service that accepts new bare act PDFs, chunks them, builds the RAPTOR tree, embeds them, and upserts into Qdrant without server restart.

  2. BM25 sparse vector fix: Rebuild the collection with properly aligned sparse vectors (same model/vocabulary for indexing and querying) to enable true BM25 scoring in RRF.

  3. Streaming API: Expose a /api/v1/stream endpoint using Server-Sent Events (SSE) that streams phase completion events to the Streamlit UI in real time.

  4. Hindi/multilingual support: Add multilingual embeddings (e.g., paraphrase-multilingual-mpnet-base-v2) and Hindi query translation to handle a broader user base.

  5. Citation verification: After P7, a separate verification step could check that cited section numbers actually exist in the retrieved chunks (preventing the LLM from fabricating section references that weren't in context).

  6. Constitutional provisions: Extend the knowledge base to include the Constitution of India (Parts, Articles, Schedules) as L0 nodes.

  7. GPU acceleration: Add CUDA support for fastembed and FlashRank to reduce inference latency from 15-40s to < 5s on GPU hardware.

  8. Confidence calibration: The current confidence score is based on FlashRank scores which are not probability-calibrated. Platt scaling or isotonic regression could be applied post-hoc to map these to true probabilities.

  9. User feedback loop: A thumbs-up/down mechanism in the Streamlit UI that logs feedback, enabling future fine-tuning of the judge thresholds or the prompt templates.

  10. Case law integration: Extend the RAPTOR tree with Supreme Court and High Court judgments as L0 nodes, with judgments linked to the statutory sections they interpret.


Quick Reference

# Start everything with Docker
docker compose up --build

# API only (development)
uvicorn ahcgr.api.app:create_app --factory --host 0.0.0.0 --port 8020 --reload

# UI only
streamlit run ui/app.py

# Health check
uv run python scripts/health_check.py

# Run tests
pytest tests/ -v

# Run RAGAS evaluation
uv run python scripts/run_ragas_eval.py

# Generate ablation CSVs
uv run python scripts/generate_ablation_csvs.py

AHCGR v1.0.0 — Built with LangGraph, Qdrant, fastembed, FlashRank, Gemini, and Streamlit