Metadata-Version: 2.4
Name: pgrag
Version: 1.0.0
Summary: Self-hosted document ingestion and hybrid (vector + keyword) retrieval server backed by Postgres/pgvector and Google embeddings.
Author: Sagar Hedaoo
License-Expression: Apache-2.0
Project-URL: Homepage, https://github.com/sagarhedaoo/vectorrag
Project-URL: Repository, https://github.com/sagarhedaoo/vectorrag
Project-URL: Issues, https://github.com/sagarhedaoo/vectorrag/issues
Keywords: rag,embeddings,pgvector,retrieval,semantic-search,fastapi
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Programming Language :: Python :: 3.12
Classifier: Framework :: FastAPI
Classifier: Topic :: Database
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Operating System :: OS Independent
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: fastapi>=0.115
Requires-Dist: uvicorn[standard]>=0.30
Requires-Dist: pydantic>=2.7
Requires-Dist: pydantic-settings>=2.3
Requires-Dist: psycopg[binary,pool]>=3.2
Requires-Dist: pgvector>=0.3.6
Requires-Dist: arq>=0.26
Requires-Dist: redis>=5.0
Requires-Dist: google-genai>=1.0
Requires-Dist: alembic>=1.13
Requires-Dist: prometheus-client>=0.20
Provides-Extra: dev
Requires-Dist: pytest>=8.2; extra == "dev"
Requires-Dist: pytest-asyncio>=1.0; extra == "dev"
Requires-Dist: testcontainers[postgres]>=4.5; extra == "dev"
Requires-Dist: httpx>=0.27; extra == "dev"
Requires-Dist: ruff>=0.5; extra == "dev"
Requires-Dist: pyright>=1.1.380; extra == "dev"
Requires-Dist: pytest-cov>=5.0; extra == "dev"
Dynamic: license-file

# VectorRAG

> A self-hosted document ingestion and **hybrid retrieval** server for Retrieval-Augmented Generation pipelines. Bring your own embedding key and database — no SaaS dependency.

[![License: Apache 2.0](https://img.shields.io/badge/License-Apache_2.0-blue.svg)](LICENSE)
[![Python: 3.12+](https://img.shields.io/badge/Python-3.12%2B-3776AB.svg?logo=python&logoColor=white)](https://www.python.org/downloads/)
[![CI](https://github.com/sagarhedaoo/VectorRAG/actions/workflows/ci.yml/badge.svg)](https://github.com/sagarhedaoo/VectorRAG/actions/workflows/ci.yml)
[![Coverage](https://img.shields.io/badge/coverage-99%25-brightgreen.svg)](#tests)
[![Type checked: pyright (strict)](https://img.shields.io/badge/types-pyright%20strict-2D6FBC.svg)](https://github.com/microsoft/pyright)

VectorRAG ingests text documents, embeds them with **Google `gemini-embedding-001`**, stores the vectors in **Postgres + pgvector**, and serves **hybrid (vector + full-text)** retrieval over HTTP — with every retrieval knob (top-k, top-p, similarity threshold, MMR, hybrid fusion, metadata filters) adjustable per request. Async ingestion via an Arq worker, API-key auth, per-key rate limiting, Prometheus metrics, structured logs, and health probes are wired in by default.

---

## Why VectorRAG

There are excellent hosted vector databases. There are also good ingestion frameworks. VectorRAG sits in the gap between them:

- **It's a service, not a library.** A real FastAPI server with auth, rate limiting, observability, and health probes — not a notebook helper.
- **You self-host it.** Your data stays in your Postgres. Your embeddings cost what _you_ pay Google, not a markup. No outbound calls beyond the embedding API.
- **The retrieval layer is configurable per request.** Want pure semantic for one query and hybrid + MMR for another? Just change the request body — no redeploy, no separate index.
- **It's built to be operated.** API-key auth, per-key rate limits, Prometheus metrics with request id correlation, `/healthz` + `/readyz` probes, graceful worker drain, content-hash dedup, idempotent migrations. The boring stuff is done.

---

## Features

- **Hybrid retrieval** — HNSW vector ANN (pgvector `halfvec` + `halfvec_cosine_ops`) fused with Postgres full-text search (`ts_rank_cd` over `tsvector`). RRF or weighted-α fusion.
- **Every knob adjustable per request** — `top_k`, `candidate_k`, `ef_search`, `similarity_threshold`, `top_p` (softmax nucleus), `mmr` + `mmr_lambda`, `hybrid` toggle, `fusion`/`alpha`/`rrf_k`, JSONB metadata `filter`, `include` field selection, and a **reranker seam** ready for a cross-encoder drop-in.
- **Embeddings via Google Gemini** — `gemini-embedding-001` at 1536 dims through `google-genai`, with **task-aware** embeddings (`RETRIEVAL_DOCUMENT` for chunks, `RETRIEVAL_QUERY` for queries) and a **content-hash cache** that never re-embeds the same chunk twice.
- **Resilient composition** — `CachingEmbedder(RetryingEmbedder(GeminiEmbedder(...)))` with exponential backoff on transient failures.
- **Async ingestion** — Submit a document over HTTP, get a `job_id` back, watch the Arq worker chunk → embed → bulk-insert with status visible via the API. Failed jobs mark the document `failed`, increment attempts, and re-raise for Arq retries.
- **Production hardening built in** — API-key auth (SHA-256 hashed at rest, Bearer or `X-API-Key`), per-key fixed-window rate limiting, JSON logs with `X-Request-ID` correlation, Prometheus metrics (`http_requests_total`, `http_request_duration_seconds`), `/healthz` and `/readyz`.
- **Operator tooling** — `pg_dump` / `pg_restore` scripts + runbook, recall@k evaluation harness, `BatchingEmbedder` for cheap bulk loads.
- **Strict quality bar** — 120 tests, **99% coverage**, ruff clean, **pyright strict 0 errors**, CI runs on every push.

---

## Quick start — Docker Compose

The fastest path to a running stack:

```bash
git clone https://github.com/sagarhedaoo/VectorRAG.git
cd VectorRAG
cp .env.example .env                              # fill in GCP_API_KEY
docker compose up -d postgres redis               # bring up infrastructure
docker compose up -d --build api worker           # build and start the app
```

That's it. The API is on `http://localhost:8000`. Run `vectorrag create-key my-app` (see [Operator commands](#operator-commands)) to mint a key, then jump to [Usage](#usage).

---

## Installation

VectorRAG supports two install paths.

### Via `pip`

```bash
pip install pgrag
```

The PyPI distribution name is `pgrag` (the `vectorrag` name was already taken by an unrelated project); the import package is still `vectorrag`. This gives you the `vectorrag` console script (`vectorrag migrate / serve / worker / create-key`) and the importable package (`import vectorrag`).

### From source (today)

```bash
git clone https://github.com/sagarhedaoo/VectorRAG.git
cd VectorRAG
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"
```

Then bring up Postgres + Redis (`docker compose up -d postgres redis`) and apply migrations:

```bash
DATABASE_URL=postgresql://rag:rag@localhost:5432/rag alembic upgrade head
```

Run the API and worker in separate shells:

```bash
# Terminal 1 — HTTP API
uvicorn vectorrag.api.app:create_app --factory --reload --port 8000

# Terminal 2 — ingestion worker
arq vectorrag.ingestion.worker.WorkerSettings
```

---

## Usage

All examples assume `BASE=http://localhost:8000` and `KEY=<your-api-key>`.

### 1. Create a collection

```bash
curl -s -X POST $BASE/collections \
  -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"name":"my-docs"}'
```

Returns:

```json
{
  "id": "8c2a...",
  "name": "my-docs",
  "embedding_model": "gemini-embedding-001",
  "embedding_dim": 1536,
  "distance": "cosine",
  ...
}
```

### 2. Submit a document for ingestion

```bash
curl -s -X POST $BASE/collections/<collection_id>/documents \
  -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{
    "text": "VectorRAG stores embeddings in Postgres and serves hybrid retrieval...",
    "title": "intro",
    "metadata": {"source": "docs", "lang": "en"},
    "chunk_size": 512,
    "chunk_overlap": 64
  }'
```

Returns `202 Accepted` with `{ "document_id": ..., "job_id": ..., "deduplicated": false }`. The worker picks the job up, chunks the text, embeds each chunk (using cached results when available), and bulk-inserts. Poll `GET /documents/{id}` until `"status": "done"`.

Dedup is per-collection on `sha256(text)`. Submitting the same text twice in the same collection returns `deduplicated: true` and no new job.

### 3. Query the collection

```bash
curl -s -X POST $BASE/collections/<collection_id>/search \
  -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{
    "query": "hybrid retrieval",
    "top_k": 5,
    "hybrid": true,
    "fusion": "rrf",
    "filter": {"lang": "en"},
    "include": ["content", "metadata", "score"]
  }'
```

Returns:

```json
{
  "results": [
    {
      "chunk_id": "...",
      "score": 0.84,
      "content": "VectorRAG stores embeddings in Postgres...",
      "metadata": { "source": "docs", "lang": "en" }
    }
  ]
}
```

### Postman collection

A ready-to-use Postman collection lives in the repo (currently gitignored — copy or import locally). Set `baseUrl` and `apiKey` collection variables, then run **Collections → Create Collection** → **Documents → Submit Document** → **Search → Search**. Auto-captured `collectionId`, `documentId`, and `jobId` flow between requests.

---

## Configuration

Set via environment variables (or a local `.env`; see `.env.example`):

| Var                      | Required | Default                | Purpose                                                                                              |
| ------------------------ | -------- | ---------------------- | ---------------------------------------------------------------------------------------------------- |
| `GCP_API_KEY`            | ✓        | —                      | Google Gemini embedding API key                                                                      |
| `DATABASE_URL`           | ✓        | —                      | Postgres connection string (must have pgvector)                                                      |
| `REDIS_URL`              | ✓        | —                      | Redis URL for the Arq queue                                                                          |
| `EMBEDDING_MODEL`        |          | `gemini-embedding-001` | Embedding model name                                                                                 |
| `EMBEDDING_DIM`          |          | `1536`                 | Embedding dimensionality (Matryoshka-truncated)                                                      |
| `EMBED_BATCH_SIZE`       |          | `100`                  | Documents per embedding call                                                                         |
| `HNSW_EF_SEARCH_DEFAULT` |          | `80`                   | Default HNSW recall/speed knob                                                                       |
| `API_KEYS_BOOTSTRAP`     |          | `""`                   | If set, this raw key is inserted as an API key on first boot (idempotent). Convenient for local dev. |

---

## Search parameters

Every request to `POST /collections/{id}/search` accepts the following:

| Parameter              | Type                  | Default                          | Notes                                                                                 |
| ---------------------- | --------------------- | -------------------------------- | ------------------------------------------------------------------------------------- |
| `query`                | string\|null          | —                                | Required unless `query_vector` is provided.                                           |
| `query_vector`         | float[]\|null         | —                                | Bypass embedding by sending a vector. Length should match `embedding_dim`.            |
| `top_k`                | int (1..1000)         | `10`                             | Final results returned. Must be ≤ `candidate_k`.                                      |
| `candidate_k`          | int (1..10000)        | `100`                            | Candidates pulled from each index before fusion.                                      |
| `ef_search`            | int\|null (1..1000)   | `HNSW_EF_SEARCH_DEFAULT`         | HNSW recall/speed dial. Set higher for better recall.                                 |
| `similarity_threshold` | float\|null (-1..1)   | —                                | Drop candidates with `vector_similarity` below this. Keyword-only hits are preserved. |
| `top_p`                | float\|null (0..1]    | —                                | Nucleus cutoff over softmaxed scores.                                                 |
| `mmr`                  | bool                  | `false`                          | Apply Maximal Marginal Relevance after threshold/top_p.                               |
| `mmr_lambda`           | float (0..1)          | `0.5`                            | 0 = max diversity, 1 = max relevance.                                                 |
| `hybrid`               | bool                  | `true`                           | Combine vector ANN with full-text search.                                             |
| `fusion`               | `"rrf"`\|`"weighted"` | `"rrf"`                          | Reciprocal Rank Fusion or min-max-normalized weighted blend.                          |
| `alpha`                | float (0..1)          | `0.5`                            | Weighted fusion only — vector contribution weight.                                    |
| `rrf_k`                | int ≥1                | `60`                             | RRF constant.                                                                         |
| `filter`               | object                | `{}`                             | JSONB containment filter on `chunks.metadata` (e.g. `{"lang": "en"}`).                |
| `distance`             | `"cosine"`            | `"cosine"`                       | Only cosine supported in v1; others return 400.                                       |
| `rerank`               | bool                  | `false`                          | Enable the reranker seam (no-op until a cross-encoder is wired).                      |
| `include`              | string[]              | `["content","metadata","score"]` | Subset of `content`, `metadata`, `score`, `embedding`.                                |

Two common recipes:

- **Pure semantic:** `{"query": "...", "top_k": 5, "hybrid": false}`
- **Hybrid + diverse + filtered:** `{"query": "...", "top_k": 5, "fusion": "rrf", "mmr": true, "filter": {"lang": "en"}}`

---

## API reference

All data endpoints require `Authorization: Bearer <key>` (or `X-API-Key: <key>`). Health and `/metrics` are open.

| Method   | Path                          | Notes                                                                                       |
| -------- | ----------------------------- | ------------------------------------------------------------------------------------------- |
| `POST`   | `/collections`                | Create a collection. **409** on duplicate name.                                             |
| `GET`    | `/collections`                | List (`limit`/`offset` query params).                                                       |
| `GET`    | `/collections/{id}`           | Get one. **404** if missing.                                                                |
| `DELETE` | `/collections/{id}`           | Cascades to documents and chunks. **204** on success.                                       |
| `POST`   | `/collections/{id}/documents` | Submit a document. **202** on accept; **404** on missing collection; **422** on empty text. |
| `GET`    | `/documents/{id}`             | Document status + metadata.                                                                 |
| `GET`    | `/jobs/{id}`                  | Job status, progress, attempts, error.                                                      |
| `POST`   | `/collections/{id}/search`    | Hybrid search. **400** on non-cosine `distance`; **404** on missing collection.             |
| `POST`   | `/embeddings`                 | Debug: returns `{dim, embedding}` for arbitrary text.                                       |
| `GET`    | `/healthz`                    | Liveness — always `200 {"status":"ok"}`.                                                    |
| `GET`    | `/readyz`                     | Readiness — **200** if DB reachable, **503** otherwise.                                     |
| `GET`    | `/metrics`                    | Prometheus text format.                                                                     |

---

## Architecture

```
                  ┌──────────────┐   enqueue job    ┌─────────┐    pull job    ┌──────────┐
   submit    ───▶│  FastAPI API │ ───────────────▶ │  Redis  │ ─────────────▶│  worker  │
                  │  (auth +     │                  │  (Arq)  │                │  (Arq)   │
   search    ───▶│  rate limit) │                  └─────────┘                └────┬─────┘
                  └──────┬───────┘                                                  │
                         │ vector search                            chunk → embed   │
                         │ + full-text                              → bulk insert   │
                         ▼                                                          ▼
                  ┌────────────────────────────────────────────────────────────────────┐
                  │  Postgres + pgvector                                                │
                  │   chunks(embedding halfvec(1536) HNSW cosine + content_tsv GIN)     │
                  │   collections / documents / jobs / embedding_cache / api_keys       │
                  └────────────────────────────────────────────────────────────────────┘
```

Four moving pieces: the **FastAPI API**, an **Arq worker**, **Postgres + pgvector**, and **Redis** for the queue. Embedding calls fan out to **Google Gemini** through a `CachingEmbedder(RetryingEmbedder(GeminiEmbedder(...)))` chain so identical chunks never re-embed and transient failures retry with exponential backoff.

The retrieval pipeline runs in this order: embed query → vector ANN + FTS in parallel → fuse (RRF or weighted) → similarity threshold → top-p nucleus → optional reranker → optional MMR → truncate to top-k.

---

## Project structure

```
src/vectorrag/
  api/             # FastAPI app, deps, routes (collections, documents, jobs, search, health)
  embeddings/      # Embedder protocol + Stub / Gemini / Retrying / Caching / Batching
  ingestion/       # chunker, pipeline, Arq worker
  observability/   # request-id middleware + JSON logs + Prometheus metrics
  repositories/    # async SQL access (collections, documents, chunks, jobs, api_keys)
  retrieval/       # Candidate, vector_search, keyword_search, fusion, postprocess, mmr, rerank, run_search
  auth.py          # Bearer / X-API-Key extraction + require_api_key dep + bootstrap
  ratelimit.py     # InMemoryRateLimiter + enforce_rate_limit dep
  config.py        # pydantic-settings Settings
  db.py            # async psycopg3 pool + pgvector type registration
  eval.py          # recall@k harness core
  hashing.py       # sha256_hex
  schemas.py       # all request/response Pydantic models
migrations/        # Alembic
scripts/
  backup.sh        # pg_dump
  restore.sh       # pg_restore
  eval.py          # recall@k CLI
tests/             # 120 tests, 99% coverage
docs/
  backup-and-restore.md
```

---

## Development

```bash
git clone https://github.com/sagarhedaoo/VectorRAG.git
cd VectorRAG
python -m venv .venv
source .venv/bin/activate
pip install -e ".[dev]"
docker compose up -d postgres redis    # tests use their own throwaway containers,
                                       # but you'll want this stack for manual testing
```

### Quality gate

The same gate CI runs:

```bash
ruff check .
pyright            # strict, 0 errors expected
pytest --cov=vectorrag --cov-report=term-missing --cov-fail-under=90
```

### Tests

Integration tests spin up a `pgvector/pgvector:pg17` container via [testcontainers](https://testcontainers.com/) — you need Docker running. No API keys or external services needed; `GeminiEmbedder` is faked in tests.

### Conventions

- **TDD-first.** Write the failing test, watch it fail, then implement.
- **Type-strict.** Public functions declare parameter and return types. `pyright` strict mode must stay at 0 errors.
- **Conventional commit messages** (`feat:`, `fix:`, `refactor:`, `chore:`, `docs:`, `test:`).
- See [CONTRIBUTING.md](CONTRIBUTING.md) for the contributor flow.

---

## Operator commands

Bringing up a server day-to-day:

```bash
docker compose up -d              # bring everything up (postgres, redis, api, worker)
vectorrag migrate                 # apply schema migrations
vectorrag create-key my-app       # mint a server API key
vectorrag serve                   # run the HTTP API
vectorrag worker                  # run the ingestion worker
```

Backup / restore:

```bash
DATABASE_URL=postgresql://rag:rag@localhost:5432/rag bash scripts/backup.sh backup.dump
DATABASE_URL=postgresql://rag:rag@localhost:5432/rag bash scripts/restore.sh backup.dump
```

Evaluation:

```bash
python scripts/eval.py --base-url http://localhost:8000 --collection <id> \
    --api-key $KEY --golden golden.jsonl --k 10
```

See [docs/backup-and-restore.md](docs/backup-and-restore.md) for the full runbook.

---

## Roadmap

- **v0.2 — Production hardening for scale.** Optional binary-quantization + rerank path for ≥10M-chunk deployments, partitioned `chunks` table, OpenTelemetry tracing.
- **v0.3 — Beyond the basics.** Cross-encoder reranker implementation behind the existing seam, batch embedding API path (50% off via Gemini batch endpoint), HyDE / query expansion experiments.

---

## License

Released under the [Apache License 2.0](LICENSE).

---

## Acknowledgements

VectorRAG stands on the shoulders of:

- **[pgvector](https://github.com/pgvector/pgvector)** — Postgres vector search.
- **[FastAPI](https://fastapi.tiangolo.com/)** + **[Pydantic](https://docs.pydantic.dev/)** — the HTTP and validation layer.
- **[psycopg3](https://www.psycopg.org/)** — async Postgres driver.
- **[Arq](https://arq-docs.helpmanual.io/)** — Redis-backed async task queue.
- **[Google AI Studio](https://ai.google.dev/) / `google-genai`** — embeddings.
- **[Alembic](https://alembic.sqlalchemy.org/)** — schema migrations.
- **[testcontainers-python](https://testcontainers-python.readthedocs.io/)** — test infrastructure.
