Metadata-Version: 2.4
Name: vs-rag
Version: 0.1.0
Summary: Enterprise RAG framework with pluggable layers
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: pydantic>=2.0
Requires-Dist: httpx>=0.25
Requires-Dist: vs-common>=0.1.3
Requires-Dist: qdrant-client>=1.7
Requires-Dist: sentence-transformers>=2.0
Requires-Dist: rank-bm25>=0.2
Requires-Dist: tiktoken>=0.5

# vs-rag

Enterprise RAG framework with pluggable layers. Drop in your embedder, vector store, LLM, and chunker — the pipeline handles retrieval, reranking, faithfulness checking, citation verification, caching, and Graph RAG automatically.

---

## Why vs-rag

Building a RAG system from scratch means wiring together 8–10 components and debugging failures across all of them. vs-rag gives you a production-ready pipeline out of the box:

- **Hybrid retrieval** — vector search + BM25 fused with configurable weights, so keyword-heavy and semantic queries both work well
- **Reranking** — cross-encoder reranker pushes the most relevant chunks to the top before generation
- **Faithfulness guard** — LLM-scored faithfulness check prevents hallucinated answers from reaching the user; triggers fallback or abstain automatically
- **Graph RAG** — entity extraction at ingest time + graph traversal at query time surfaces related chunks that vector search alone would miss (great for multi-hop questions)
- **Caching** — query-level and embedding-level caching to avoid redundant LLM and embedding calls
- **Citations** — every answer comes with chunk-level citations so you know exactly where each claim came from
- **Pluggable everything** — every layer (embedder, store, chunker, reranker, LLM, cache) is an abstract base class; swap implementations without touching the pipeline

---

## Installation

```bash
pip install vs-rag
```

**Infrastructure dependencies** (run locally or in Docker):

| Service | Used for | Default URL |
|---------|----------|-------------|
| Qdrant  | Vector store | `http://localhost:6333` |
| Ollama  | Embeddings + LLM (optional) | `http://localhost:11434` |

Start Qdrant:
```bash
docker run -p 6333:6333 qdrant/qdrant
```

Start Ollama with the default embedding model:
```bash
ollama pull nomic-embed-text:latest
```

---

## Quickstart

### 1. Config file (`config.ini`)

```ini
[embedder]
provider = litellm
model = ollama/nomic-embed-text:latest
url = http://localhost:11434
dimension = 768

[vector_store]
provider = qdrant
url = http://localhost:6333

[chunker]
strategy = recursive
chunk_size = 512
chunk_overlap = 50

[retriever]
top_k = 20
vector_weight = 0.7
bm25_weight = 0.3

[reranker]
provider = cross_encoder
top_n = 5

[faithfulness]
threshold = 0.75
fallback = abstain
```

### 2. Implement `VsLLMClient`

vs-rag doesn't ship with an LLM client — you bring your own. Implement two methods:

```python
from vs_rag import VsLLMClient

class MyLLMClient(VsLLMClient):

    async def rephrase_query(self, query: str, num_variations: int) -> list[str]:
        # Call your LLM to generate query variations for multi-query retrieval
        ...

    async def check_faithfulness(self, answer: str, context: str) -> float:
        # Call your LLM to score how grounded the answer is in the context
        # Return a float between 0.0 and 1.0
        ...
```

### 3. Build the pipeline

```python
import asyncio
from vs_common.config.vs_ini_config import VsIniConfig
from vs_rag import RagFactory, Document

config = VsIniConfig("config.ini")
pipeline = RagFactory.from_config(config, llm=MyLLMClient())

async def main():
    # Ingest a document
    doc = Document(filename="handbook.md", content=open("handbook.md").read(), content_type="markdown")
    result = await pipeline.ingest(doc, user_ref="user_123")
    print(f"Ingested {result.chunks_created} chunks into '{result.collection}'")

    # Query
    async def generate(prompt: str) -> str:
        # call your LLM with the prompt that already contains the retrieved context
        return my_llm_call(prompt)

    response = await pipeline.query(
        query="What is the parental leave policy?",
        user_ref="user_123",
        generate_fn=generate,
    )

    print(response.answer)
    print(f"Confidence: {response.confidence.overall:.2f}")
    for c in response.citations:
        print(f"  [{c.source}] {c.text[:60]}...")

asyncio.run(main())
```

### Response shape

```python
class RagResponse:
    answer: str
    confidence: ConfidenceScore   # retrieval, faithfulness, overall — all 0.0–1.0
    citations: list[Citation]     # per-sentence source attribution
    fallback_triggered: bool      # True if faithfulness check failed
    cached: bool                  # True if result came from query cache
```

---

## Features

### Chunking strategies

Set `chunker.strategy` in config:

| Strategy | Best for |
|----------|----------|
| `recursive` (default) | General text, tries paragraph → sentence splits |
| `fixed` | Uniform token budgets |
| `semantic` | Splits on semantic similarity boundaries (requires embedder) |
| `document_aware` | Markdown/structured docs — splits on headings, preserves section context |

```ini
[chunker]
strategy = document_aware
chunk_size = 1024
```

### Hybrid retrieval tuning

Vector search finds semantically similar chunks; BM25 finds keyword-matching chunks. Adjust the blend:

```ini
[retriever]
vector_weight = 0.7   # increase for semantic/conceptual queries
bm25_weight = 0.3     # increase for exact keyword / filter queries
top_k = 20
```

### Reranking

The cross-encoder reranker re-scores the top-k retrieved chunks using a more accurate (but slower) model before sending them to the generator. Reduce `top_n` to send fewer, higher-quality chunks:

```ini
[reranker]
provider = cross_encoder
top_n = 5   # chunks sent to the LLM after reranking
```

To skip reranking (faster, lower quality):
```ini
[reranker]
provider = passthrough
```

### Faithfulness guard

The pipeline scores the generated answer against the retrieved context. If the score is below `threshold`, it triggers the fallback:

```ini
[faithfulness]
threshold = 0.75

# abstain: immediately return "I don't know" — safest
# retry: re-run generation up to max_retries times before abstaining
fallback = abstain

# only relevant when fallback = retry
max_retries = 2
```

If no `VsLLMClient` is provided, faithfulness checking is skipped and all answers pass through.

### Caching

Two independent cache layers:

**Query cache** — caches the retrieved + reranked chunks for a `user_ref:query` key. Same query from the same user skips the entire retrieval pipeline.

**Embedding cache** — caches embedding vectors so repeated texts (e.g. chunks that appear in multiple documents) are not re-embedded.

```ini
[cache]
query_backend = vs
query_prefix = rag:query:
embedding_backend = vs
embedding_prefix = rag:embed:
```

TTLs are set in code when building the pipeline:

```python
pipeline = RagPipeline(
    ...
    query_cache_ttl=3600,      # 1 hour
    embedding_cache_ttl=86400, # 24 hours
)
```

The built-in `vs` backend delegates to `VsCacheManager` (backed by Redis or in-memory depending on your vs-common config). To use a different backend, implement `VsQueryCache` or `VsEmbeddingCache` and register it:

```python
from vs_rag.pipeline.rag_registry import RagComponentRegistry
RagComponentRegistry.register_query_cache("my_backend", lambda cfg: MyQueryCache(cfg))
```

### Graph RAG

Graph RAG extracts entities and relationships from each chunk at ingest time and builds a knowledge graph. At query time, entities mentioned in the question are used to traverse the graph and pull in related chunks that vector search alone would not find — essential for multi-hop questions.

#### Implement `VsEntityExtractor`

```python
from vs_rag import VsEntityExtractor
from vs_rag.schema.graph import Entity, Relationship

class MyEntityExtractor(VsEntityExtractor):

    async def extract(self, text: str, chunk_id: str, collection: str):
        # Call your LLM to extract entities and relationships from text
        # Return (List[Entity], List[Relationship])
        ...
```

#### Enable it

```python
pipeline = RagFactory.from_config(
    config,
    llm=MyLLMClient(),
    entity_extractor=MyEntityExtractor(),
    # graph_store defaults to InMemoryGraphStore if not provided
)
```

Graph traversal is automatic — the `HybridRetriever` enriches results with graph-connected chunks transparently. No changes to query calls.

#### Bring your own graph store

The default `InMemoryGraphStore` is lost on restart. For persistence, implement `VsGraphStore`:

```python
from vs_rag import VsGraphStore
from vs_rag.schema.graph import Entity, Relationship, GraphResult

class Neo4jGraphStore(VsGraphStore):
    async def add_entities(self, entities): ...
    async def add_relationships(self, relationships): ...
    async def traverse(self, entity_names, top_k): ...
    async def clear(self): ...
```

```python
pipeline = RagFactory.from_config(config, llm=..., entity_extractor=..., graph_store=Neo4jGraphStore())
```

---

## Customization

Every component is an abstract base class. Register your own implementation and vs-rag will use it automatically.

### Custom embedder

```python
from vs_rag import VsEmbedder
from vs_rag.pipeline.rag_registry import RagComponentRegistry

class MyEmbedder(VsEmbedder):
    async def embed(self, texts): ...
    async def embed_one(self, text): ...
    def dimension(self): return 1536

RagComponentRegistry.register_embedder("my_embedder", lambda cfg: MyEmbedder())
```

```ini
[embedder]
provider = my_embedder
```

### Custom vector store

```python
from vs_rag import VsVectorStore
from vs_rag.pipeline.rag_registry import RagComponentRegistry

class PineconeStore(VsVectorStore):
    ...

RagComponentRegistry.register_vector_store("pinecone", lambda cfg, embedder: PineconeStore(cfg))
```

```ini
[vector_store]
provider = pinecone
```

### Custom chunker

```python
from vs_rag import VsChunker
from vs_rag.pipeline.rag_registry import RagComponentRegistry

class MyChunker(VsChunker):
    async def chunk(self, text, metadata): ...

RagComponentRegistry.register_chunker("my_chunker", lambda cfg, embedder: MyChunker())
```

```ini
[chunker]
strategy = my_chunker
```

The same pattern works for rerankers (`register_reranker`), faithfulness checkers (`register_faithfulness_checker`), and fallback handlers (`register_fallback`).

---

## Multi-user / multi-collection

vs-rag is multi-tenant by design. Pass `user_ref` to `ingest` and `query` — the pipeline automatically routes each user to their own Qdrant collections and scopes cache keys by user.

```python
# User A's documents stay isolated from User B's
await pipeline.ingest(doc, user_ref="user_a")
await pipeline.ingest(doc, user_ref="user_b")

response = await pipeline.query("...", user_ref="user_a")  # only searches user_a's collections
```

The collection prefix is configurable:
```ini
[rag]
collection_prefix = myapp
```

---

## Accuracy evaluation

A built-in eval runner measures correctness, faithfulness, and retrieval quality against a QA dataset:

```bash
cd tests
python eval_runner.py --dataset datasets/company_handbook_qa.json

# Skip re-ingestion if already done
python eval_runner.py --dataset datasets/company_handbook_qa.json --skip-ingest

# Enable Graph RAG (requires re-ingest)
python eval_runner.py --dataset datasets/company_handbook_qa.json --graph
```

Results are saved to `tests/eval_report_<dataset_name>.json` with per-question breakdown and category-level aggregates.

---

## Built-in components reference

| Layer | Built-in options |
|-------|-----------------|
| Embedder | `litellm` (any model via LiteLLM) |
| Vector store | `qdrant` |
| Chunker | `recursive`, `fixed`, `semantic`, `document_aware` |
| Reranker | `cross_encoder`, `passthrough` |
| Fallback | `abstain`, `retry` |
| Query cache | `vs` (VsCacheManager) |
| Embedding cache | `vs` (VsCacheManager) |
| Graph store | `InMemoryGraphStore` |
