Metadata-Version: 2.4
Name: lex-router
Version: 0.1.0
Summary: A domain-agnostic adaptive query router for RAG systems. Routes queries using multi-signal retrieval score analysis, not vocabularies.
Home-page: https://github.com/Vedant-ghadi/LexLegal
Author: Vedant Ghadi
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Text Processing :: Linguistic
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Requires-Python: >=3.7
Description-Content-Type: text/markdown
License-File: LICENSE
Dynamic: author
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: license-file
Dynamic: requires-python
Dynamic: summary

# Lex-Router

**An adaptive, domain-agnostic query router for RAG systems.**

Most RAG systems run the same heavy retrieval pipeline for every query. Lex-Router fixes this by analyzing the **statistical distribution of pilot retrieval scores** to dynamically route queries to the optimal strategy. It works across any domain (legal, medical, financial, etc.) because routing decisions come from the retrieval landscape itself, not hardcoded vocabularies.

Built on the multi-signal pilot framework from the [lexLegal RAG Pipeline](https://github.com/Vedant-ghadi/LexLegal).

---

## Installation

```bash
pip install lex-router
```

**Zero dependencies.** Written entirely using the Python Standard Library (`math`, `re`, `dataclasses`).

---

## Core Features

- **Domain Agnostic**: Doesn't rely on keyword lists or LLM calls. Routes based on score variance, entropy, and statistical confidence.
- **Score Normalization**: Built-in support for multiple embedding score types (`COSINE`, `L2`, `DOT_PRODUCT`, `LOGITS`, `COSINE_DISTANCE`, `RAW`).
- **Auto-Calibration**: Automatically tunes internal thresholds to your specific embedding model to prevent "threshold drift".
- **Architecture Aware**: Gracefully handles Hybrid (Dense+Sparse), Dense-Only (e.g., Pinecone), and Sparse-Only (e.g., Elasticsearch) systems.
- **Ultra-Low Latency**: Sub-millisecond routing overhead, plus a configurable fast-path bypass for trivial queries.
- **Production Safe**: Robust sanitization against `NaN`/`Inf` inputs, mathematical overflows, and database timeouts.

---

## Use Cases

Lex-Router is designed to handle edge cases across diverse RAG architectures:

1. **Complex Domain Search (Legal/Medical)**: 
   If a user asks about "the MAE clause regarding pandemic events", the router detects high term rarity and high score variance, routing it to a `broad_expand` strategy (deep vector search + graph traversal).
2. **Ultra-Low Latency Voice Assistants**: 
   If a user asks "What is this?", the router hits the `fast_path_max_tokens` bypass, skipping the database pilot search entirely and routing instantly to save latency.
3. **Legacy Enterprise Search**: 
   If your company only uses a keyword-based Elasticsearch database, Lex-Router's `is_sparse_only` mode will dynamically disable dense variance checks and route based purely on BM25 score margins and entropy.
4. **Machine Learning Transition**: 
   If you want to train a custom ML model to route queries, use the `log_file` parameter. The router will log a rich 11-signal feature vector for every query, giving you perfect training data for an XGBoost or Random Forest model later.

---

## How It Works

The router runs a fast **pilot search** (e.g., top 10 results) through your retrieval backend and extracts 11 statistical signals:

| Signal | What It Measures |
|:---|:---|
| `max_dense`, `bm25_max` | Maximum retrieval strength for vector and keyword searches |
| `mean_dense`, `std_dense` | Vector score distribution and variance |
| `top1_top5_margin` | Confidence gap between the top result and the rest |
| `dense_bm25_overlap` | Agreement between dense and sparse retrievers |
| `entropy` | Score distribution chaos |
| `unique_doc_count` | Document diversity |
| `query_rarity`, `max_term_rarity` | IDF-based term specificity |

These signals map to **4 retrieval strategies**:

| Route | When It Triggers | Pool | Features Enabled |
|:---|:---|:---:|:---|
| `narrow_precise` | Rare anchor term or very confident top-1 | 50 | BM25 only |
| `normal` | Balanced signals (Standard behavior) | 100 | HyDE + BM25 + Graph |
| `broad_expand` | Generic terms, high variance, many docs | 200 | Full pipeline |
| `reject_or_fallback`| Both retrievers returned garbage scores | 0 | Trigger fallback |

---

## Quick Start

### Option A: Connect Your Pipeline via Adapter

Implement the `RetrievalAdapter` interface. You only need to override the search methods your system supports.

```python
from lex_router import AdaptiveRouter, RetrievalAdapter

class MyVectorAdapter(RetrievalAdapter):
    def __init__(self, vector_db):
        self.db = vector_db

    # Only implement dense_search if you have a dense-only system!
    def dense_search(self, query, k=10):
        results = self.db.search(query, top_k=k)
        scores = [r.score for r in results]
        doc_ids = [r.id for r in results]
        return scores, doc_ids

# Plug it in
router = AdaptiveRouter(adapter=MyVectorAdapter(my_db))

# Route the query
decision = router.route("What are the termination provisions?")
print(decision.route)  # e.g., 'normal'
print(decision.pool)   # e.g., 100
```

### Option B: Pass Raw Scores (No Adapter)

If you already have retrieval scores from your pipeline, pass them directly:

```python
from lex_router import AdaptiveRouter, ScoreType

router = AdaptiveRouter(
    dense_score_type=ScoreType.COSINE,
    sparse_score_type=ScoreType.RAW
)

decision = router.route_from_scores(
    dense_scores=[0.92, 0.87, 0.61, 0.45, 0.32],
    sparse_scores=[12.4, 8.1, 5.3],
    dense_docs=["doc_A", "doc_A", "doc_B"],
    sparse_docs=["doc_A", "doc_C", "doc_B"],
    query="Does the MAE clause exclude pandemic events?",
)

print(decision.route)               # 'narrow_precise'
print(decision.metadata['reason'])  # 'confident_top1'
```

---

## Advanced Usage

### Handling Different Embedders (`ScoreType`)

Different vector databases and embedding models return scores on completely different mathematical scales. Lex-Router normalizes these to a standard `[0, 1]` scale internally.

When initializing the router, declare your score types:

```python
from lex_router import ScoreType

router = AdaptiveRouter(
    dense_score_type=ScoreType.L2,              # Converts [0, inf] -> [1, 0]
    sparse_score_type=ScoreType.RAW             # Leaves BM25 scores unbounded
)
```
Supported types: `COSINE`, `COSINE_DISTANCE`, `L2`, `DOT_PRODUCT`, `LOGITS` (for cross-encoders), `RAW`, and `CUSTOM`.

### Auto-Calibrating Thresholds

Default routing thresholds (`t_high=0.75`, `t_low=0.35`) are optimized for specific models (like BGE-M3). If you use a different embedder (like OpenAI `text-embedding-3`), your score distribution will shift. Lex-Router can calibrate itself perfectly to your model:

```python
# Provide a baseline sample of typical retrieval scores (e.g., from 50 queries)
router = AdaptiveRouter.auto_calibrate(
    baseline_dense_scores=[
        [0.82, 0.71, 0.60], 
        [0.44, 0.42, 0.41], 
        # ... more score batches
    ],
    dense_score_type=ScoreType.COSINE
)
# The router automatically calculates optimal percentiles for thresholding.
```

### Ultra-Low Latency Bypass

For real-time voice applications or chat, you can bypass pilot retrieval entirely for very short queries:

```python
from lex_router import RouterConfig

# Any query <= 3 words skips pilot search and goes straight to narrow_precise
cfg = RouterConfig(fast_path_max_tokens=3)
router = AdaptiveRouter(config=cfg, adapter=my_adapter)
```

### Safe Logging for ML Training

Lex-Router is designed to transition from heuristics to machine learning. You can log every decision, along with its 11-feature signal vector, to train an ML model later. Use the router as a context manager to ensure safe resource handling:

```python
with AdaptiveRouter(adapter=my_adapter, log_file="routing_decisions.jsonl") as router:
    decision = router.route("query one")
    decision = router.route("query two")
    
    router.print_summary() # Prints route distribution percentages
```

---

## API & Parameters Reference

### `AdaptiveRouter`
The main class that orchestrates routing.
*   **`adapter`** (`RetrievalAdapter`): Your custom DB connection.
*   **`config`** (`RouterConfig`): Tuning parameters and thresholds.
*   **`pilot_k`** (`int`): Number of pilot results to retrieve. Default is `10`.
*   **`dense_score_type`** (`ScoreType`): Mathematical scale of vector scores. Default is `COSINE`.
*   **`sparse_score_type`** (`ScoreType`): Mathematical scale of keyword scores. Default is `RAW`.
*   **`log_file`** (`str`): Path to output a JSONL log of all decisions.

### `RouterConfig`
Data class containing all thresholds and boundaries. Pass this to `AdaptiveRouter(config=...)` to override defaults.
*   **`t_high` / `t_low`**: Confidence thresholds for determining strong or weak matches.
*   **`std_high`**: Threshold for variance (entropy).
*   **`narrow_pool` / `normal_pool` / `broad_pool`**: The sizes of the document pools returned in the `RouteDecision`.
*   **`fast_path_max_tokens`**: If > 0, queries with length equal to or less than this will bypass the database entirely for speed.
*   **`min_pilot_results`**: Minimum number of DB results required; if fewer are returned, degrades gracefully to `normal` route.

### `RouteDecision`
The output of `router.route()`.
*   **`route`** (`str`): The strategy chosen (e.g., `'narrow_precise'`).
*   **`pool`** (`int`): Suggested number of documents to retrieve.
*   **`use_hyde` / `use_bm25` / `use_xref`** (`bool`): Feature toggles.
*   **`metadata`** (`dict`): Contains the exact `reason` for the route, and the raw 11-signal vector.

---

## License

MIT License. See [LICENSE](LICENSE).
