Metadata-Version: 2.4
Name: kolzchut-ragbot
Version: 2.0.0
Summary: A search engine using machine learning models and Elasticsearch for advanced document retrieval.
Home-page: https://github.com/shmuelrob/rag-bot
Author: Shmuel Robinov
Author-email: shmuel_robinov@webiks.com
Classifier: Programming Language :: Python :: 3
Classifier: License :: OSI Approved :: MIT License
Classifier: Operating System :: OS Independent
Requires-Python: >=3.10
Description-Content-Type: text/markdown
Requires-Dist: elasticsearch==8.17.1
Requires-Dist: sentence-transformers==3.4.1
Requires-Dist: torch==2.6.0
Requires-Dist: transformers==4.48.3
Dynamic: author
Dynamic: author-email
Dynamic: classifier
Dynamic: description
Dynamic: description-content-type
Dynamic: home-page
Dynamic: requires-dist
Dynamic: requires-python
Dynamic: summary

# kolzchut-ragbot

Retrieval engine for the Kol-Zchut RAG chatbot. It embeds documents with sentence-transformer
models, stores them in Elasticsearch, and returns the most relevant paragraphs for a user question.

- **PyPI package:** [`kolzchut-ragbot`](https://pypi.org/project/kolzchut-ragbot/)
- **Python import name:** `kolzchut_ragbot`
- **Repository:** https://github.com/kolzchut/ragbot

---

## 1. Project Overview

### What this repository is

`kolzchut-ragbot` is a **library**, not a service. It has no HTTP server, no CLI and no `main.py`.
You install it as a dependency and call its functions from your own application.

It covers the **"R" (Retrieval) part of RAG**:

| Responsibility | In this repo? |
| --- | --- |
| Turning text into embedding vectors | Yes |
| Creating and writing Elasticsearch indices | Yes |
| Vector search + result fusion + ranking | Yes |
| Talking to an LLM / building the final answer | **No** |
| Prompts, conversation history, HTTP API | **No** — lives in the consuming service |
| Rebuilding a full page from its paragraphs | **No** — the host reads its own `full_pages` index |

### Architectural role

```
                ┌──────────────────────────────┐
                │  chatbot (FastAPI service)   │
                │  routes, prompts, sessions   │
                └───────────────┬──────────────┘
                                │ imports
                                ▼
                ┌──────────────────────────────┐
                │      kolzchut-ragbot         │
                │   Engine  ·  Model  ·  Docs  │
                └───────┬──────────────┬───────┘
                        │              │
        embedding models│              │queries
        (local folder)  ▼              ▼
              SentenceTransformer   Elasticsearch
```

The main consumer in this monorepo is [`chatbot`](../chatbot). It creates the engine once at
startup in `chatbot/app/src/main.py`:

```python
from kolzchut_ragbot.engine import engine_factory
from kolzchut_ragbot.model import es_client_factory

es_model = es_client_factory(es_client)
engine = engine_factory(configs_class=configs, es_client=es_client)
```

and then calls `engine.retrieve_documents(...)` on every search request and
`engine.update_docs(...)` whenever a wiki page changes.

### Design decisions worth knowing up front

1. **The document schema is not hardcoded.** It is loaded at import time from a JSON file
   (see [The document definition file](#22-the-document-definition-file)). Change the JSON, and the
   Elasticsearch mapping, the embedded fields and the validation rules all change with it.
2. **Everything is a singleton.** `factory()`, `es_client_factory()` and `engine_factory()` each
   cache one global instance. Calling them again returns the same object and **ignores the new
   arguments**. This matters in tests — see [Running the tests](#7-running-the-tests).
3. **The library never reads a `.env` file itself.** It only reads `os.getenv(...)`. The hosting
   application is responsible for loading the environment.

---

## 2. How It Works

### 2.1 Key modules

| File | Purpose |
| --- | --- |
| `kolzchut_ragbot/Document.py` | Loads and validates the document definition JSON. Entry point: `factory()`, which returns the cached `DocumentDefinitions`. |
| `kolzchut_ragbot/model.py` | Everything that touches Elasticsearch directly: index naming, mapping creation, writing documents, vector search. |
| `kolzchut_ragbot/engine.py` | The orchestration layer: embed the query, search, fuse the per-model results, and trim to the requested paragraph budget. |
| `kolzchut_ragbot/config.py` | Reads `MODELS_LOCATION` from the environment. (`ES_EMBEDDING_INDEX` is read in `model.py`, where it is used.) |

### 2.2 The document definition file

This JSON file is the schema for your data. Its path comes from the
`DOCUMENT_DEFINITION_CONFIG` environment variable (default: `example-conf.json`).

```json
{
  "identifier_field": "page_id",
  "saved_fields": {
    "title": "text",
    "page_id": "integer",
    "url": "text",
    "content": "text",
    "summary": "text",
    "revision_date": "date"
  },
  "field_for_llm": "content",
  "models": {
    "me5_large-v10": "content",
    "me5_large-titles-retrieval-mistakes": "title",
    "me5_large_summary-retrieval-mistakes": "summary"
  }
}
```

| Key | Meaning |
| --- | --- |
| `saved_fields` | Field name to Elasticsearch field type. These become the data fields in the generated index mapping. |
| `identifier_field` | Which saved field identifies a page. Must exist in `saved_fields`. |
| `models` | **Embedding model folder name to the field that model embeds.** Note the direction: the key is the model, the value is the field. Each entry produces one vector field. |
| `field_for_llm` | Which field is passed to the LLM as the answer context. Optional. |

`initialize_definitions()` raises a `ValueError` at import time if `identifier_field` or any value
in `models` is not listed in `saved_fields`. A bad config fails fast, on import — not at query time.

### 2.3 Indexing flow (`engine.update_docs`)

```
list of paragraph dicts
        │
        ▼
for each (model, field) in definitions.models:
    encode doc[field]  ──►  doc["<field>_<model>_vectors"]
        │
        ▼
stamp doc["last_update"] = now
        │
        ▼
Model.create_or_update_documents()
        │
        ├─ index name = index_from_page_id(page_id)
        │        = "<ES_EMBEDDING_INDEX>_<round(page_id / 1000)>"
        │
        ├─ if update=True: delete_by_query on "<ES_EMBEDDING_INDEX>*"
        │                  for this identifier  (replace, not duplicate)
        │
        ├─ create the index if it does not exist (with the generated mapping)
        │
        └─ index every paragraph one by one
```

**Why the index is split.** `index_from_page_id` puts roughly 1,000 page IDs into each physical
index (`embedded_fusion_0`, `embedded_fusion_1`, and so on). Searches always run against the
wildcard `<ES_EMBEDDING_INDEX>*`, so the split is invisible to readers. It keeps individual indices
small and makes rebuilding a slice of the corpus cheaper.

**Vector field naming.** For model `me5_large-v10` embedding field `content`, the vector field is
`content_me5_large-v10_vectors`. Any field ending in `_vectors` is treated as an embedding and is
stripped from API responses by `strip_vector_fields()`.

**Mapping.** `create_mapping()` declares every vector field as `dense_vector` with **1024
dimensions**, hardcoded. All embedding models used must output 1024-dimensional vectors.

### 2.4 Retrieval flow (`engine.retrieve_documents`)

```
query string
    │
    ▼
1. Embed the query once per model            (engine.search_documents)
    │
    ▼
2. Model.search() — one Elasticsearch query per model.
   Uses script_score with cosineSimilarity(...) + 1.0
   (the +1.0 keeps scores positive, which Elasticsearch requires)
   Each model returns up to `retrieval_size` paragraphs.
    │
    ▼
3. Drop paragraphs scoring below `paragraph_relevance_threshold` (0 = keep all)
    │
    ▼
4. Reciprocal Rank Fusion across the per-model ranked lists:
       score(page) = sum of  weight / (k + rank + 1)
   Pages ranked high by several models rise to the top.
    │
    ▼
5. Keep the top `len(paragraphs_per_page)` pages.
   For the page at position i, keep at most `paragraphs_per_page[i]` paragraphs,
   skipping any paragraph already selected.
       e.g. [3, 2, 1] → 3 paragraphs from the best page,
                        2 from the second, 1 from the third.
    │
    ▼
returns (top_k_documents, all_docs_and_scores, retrieval_time)
```

**What Reciprocal Rank Fusion (RRF) does.** Different models produce scores on different scales, so
their raw scores cannot be compared. RRF ignores the scores and uses only the **position** in each
list. A page at rank 0 in one list and rank 2 in another gets `1/(k+1) + 1/(k+3)`. The constant `k`
dampens the advantage of the very top positions.

> **`k` here is not the textbook 60.** `search_documents` passes the page count as `k`
> (`reciprocal_rank_fusion(ids_for_fusion, k=top_k)`), so with a typical `paragraphs_per_page` of
> three to five, `k` is small and the top-ranked pages dominate far more strongly than standard RRF
> would make them. This is deliberate. Keep it in mind before tuning retrieval.

**Return values.**

- `top_k_documents` — a list of paragraph dicts ready to hand to an LLM, each with its `_id`.
- `all_docs_and_scores` — a debug/analytics view: per model, every candidate with its raw
  Elasticsearch score, title, page id and vector-free `_source`. The chatbot admin screens use this.

### 2.5 The configuration object

`Engine` takes a `configs_class` and simply stores it on the instance. The library itself never calls
anything on it; the hosting application reads it back through `engine.configs_class`. In the chatbot,
`saved_config.py` provides it and `updater_service.py` calls `get_config()` to read `banned_fields`.

There is no abstract base class to inherit from — the contract is duck-typed. Anything exposing
`get_config()` works.

---

## 3. Prerequisites

| Requirement | Version / note |
| --- | --- |
| Python | 3.10 or newer (`setup.py` enforces `>=3.10`) |
| Elasticsearch | 8.x, reachable from your machine |
| Embedding models | Downloaded locally as sentence-transformer folders |
| Disk | The models are large (roughly 2 GB each); make sure you have room |
| GPU | Optional. The code uses CUDA automatically if `torch.cuda.is_available()`, otherwise CPU. |

You do **not** need an LLM API key to work on this repository — the library does not call an LLM.

---

## 4. Getting Started

### Step 1 — Install

**As a dependency in another project** — published at
[pypi.org/project/kolzchut-ragbot](https://pypi.org/project/kolzchut-ragbot/):

```bash
pip install kolzchut-ragbot
```

**For development on this repository:**

```bash
git clone https://github.com/kolzchut/ragbot.git
```

```bash
python -m venv venv
```

```bash
source venv/bin/activate
```

On Windows use `venv\Scripts\activate` instead.

```bash
pip install -r requirements.txt
```

### Step 2 — Get the embedding models

The library loads models from a local folder — it never downloads them. Each model lives in its own
subfolder, and the **folder name must exactly match the key in the `models` section of your document
definition file**.

```
artifacts/
├── me5_large-v10/
├── me5_large-titles-retrieval-mistakes/
└── me5_large_summary-retrieval-mistakes/
```

The Kol-Zchut models are not in this repository. Ask a maintainer for the artifact storage location,
or copy them from a working `chatbot` checkout (`chatbot/app/artifacts/`). Access details are in the
internal handover document:
[מסמך חפיפה פנימי - כל זכות](https://docs.google.com/document/d/1Vw99KJXy1Cj2etQohB_2PE_puBszK8rSWG8kc9C_HBA).

### Step 3 — Write your document definition file

Copy `example-conf.json` and edit it to match your data. See
[The document definition file](#22-the-document-definition-file) for what each key means.

### Step 4 — Set the environment variables

```bash
cp .env.example .env
```

Then fill in the values:

| Variable | Default | What it does |
| --- | --- | --- |
| `MODELS_LOCATION` | `models` | Folder holding the embedding model subfolders. |
| `ES_EMBEDDING_INDEX` | `embeddings` | Index name prefix. Actual indices are `<prefix>_<n>`; searches use `<prefix>*`. |
| `DOCUMENT_DEFINITION_CONFIG` | `example-conf.json` | Path to the document definition JSON from Step 3. |

> The library reads these with `os.getenv` at **import time**. Set them before importing
> `kolzchut_ragbot`, and remember that a plain `.env` file is not loaded automatically — use
> `python-dotenv`, your shell, or Docker Compose.
>
> No secrets or credentials belong in this repository. Elasticsearch credentials are supplied by the
> hosting application through the `es_client` you pass in. For the real values, see the
> [internal handover document](https://docs.google.com/document/d/1Vw99KJXy1Cj2etQohB_2PE_puBszK8rSWG8kc9C_HBA).

### Step 5 — Run it

```python
import os

# Must be set before importing the library — it reads them on import.
os.environ["DOCUMENT_DEFINITION_CONFIG"] = "doc-config.json"
os.environ["MODELS_LOCATION"] = "artifacts"
os.environ["ES_EMBEDDING_INDEX"] = "embedded_fusion"

from elasticsearch import Elasticsearch
from kolzchut_ragbot.engine import engine_factory


class SimpleConfigs:
    """The configuration object is duck-typed - it only needs get_config()."""

    def get_config(self): return {"paragraph_identifier": "paragraph"}


es_client = Elasticsearch("http://localhost:9200")
engine = engine_factory(configs_class=SimpleConfigs(), es_client=es_client)

# Index a page (all paragraphs of one page in a single call).
engine.update_docs(
    list_of_docs=[
        {"page_id": 1, "title": "Some title", "content": "Paragraph text", "url": "https://..."},
    ],
    embed_only_fields=["content"],
    delete_existing=True,   # replace the existing copy of this page
)

# Search.
documents, scores, elapsed = engine.retrieve_documents(
    query="What are my rights?",
    retrieval_size=50,               # candidates fetched per model
    paragraphs_per_page=[3, 2, 1],   # 3 pages back: 3, 2 and 1 paragraphs
    paragraph_relevance_threshold=1.6,
)
```

> **The first call is slow.** Loading the sentence-transformer models into memory takes time. This is
> why `engine_factory` is a singleton — create the engine once at application startup.

---

## 5. Public API Reference

### `engine_factory(configs_class, es_client=None) -> Engine`

Returns the shared `Engine`. Only the first call actually builds one.

### `Engine.update_docs(list_of_docs, embed_only_fields=None, delete_existing=False)`

Embeds and indexes a list of paragraph dicts. All dicts must belong to the **same page** — the index
is chosen from the first dict's identifier, so a call mixing pages raises `ValueError`.
`embed_only_fields` limits which fields get vectors; `delete_existing=True` removes the previous copy
of the page before writing.

### `Engine.retrieve_documents(query, retrieval_size, paragraphs_per_page, paragraph_relevance_threshold=0)`

The main search entry point. Returns `(top_k_documents, all_docs_and_scores, retrieval_time)`. The
number of pages returned is `len(paragraphs_per_page)`.

### `Engine.search_documents(query, retrieval_size, paragraphs_per_page, paragraph_relevance_threshold=0)`

The same retrieval, without the timing wrapper.

### `Engine.configs_class` / `Engine.field_for_answer`

Two attributes the host reads back: the configuration object it passed in, and the field name from
`field_for_llm` in the definition file.

### `Engine.reciprocal_rank_fusion(ranking_lists, k=60, weights=None) -> list`

Merges several ranked ID lists into one. A pure function, easy to unit test.

### `strip_vector_fields(source) -> dict`

A copy of an Elasticsearch `_source` without any `*_vectors` key. Use it before returning documents
over an API — the vectors are large and useless to clients.

### `es_client_factory(es_client) -> Model`

Returns the shared `Model`, the low-level Elasticsearch wrapper. Like the other factories, only the
first call builds one.

### `Model.create_or_update_documents(paragraphs_dicts, update=False)`

Writes already-embedded paragraph dicts straight to Elasticsearch, creating the index and mapping if
needed. This is the layer `Engine.update_docs` sits on top of; call it directly only when the
documents already carry their vectors.

---

## 6. Repository Layout

```
ragbot/
├── kolzchut_ragbot/                    # the package
│   ├── engine.py                       # orchestration: embed → search → fuse → trim
│   ├── model.py                        # Elasticsearch: mapping, indexing, vector search
│   ├── Document.py                     # definition-file loading and validation
│   └── config.py                       # environment variables
├── test/                               # unit suite (mocked)
├── characterization/                   # integration suite against a real Elasticsearch
├── .github/workflows/                  # test, tag, publish
├── example-conf.json                   # sample document definition
├── requirements.txt                    # pinned runtime dependencies
├── setup.py                            # the file that actually builds the PyPI package
├── pyproject.toml                      # poetry metadata + bumpver settings
└── builder.sh                          # manual build + upload to TestPyPI
```

---

## 7. Running the Tests

The tests must be run **from inside the `test/` directory** — `conftest.py` sets
`DOCUMENT_DEFINITION_CONFIG` to `test/example-conf.json` and adds the repository root to `sys.path`.

```bash
cd test
```

```bash
pytest .
```

With coverage, exactly as CI runs it:

```bash
cd test && coverage run --rcfile=.coveragerc -m pytest . && coverage report
```

**Why the tests use mocks everywhere.** No Elasticsearch instance and no model files are needed —
`Engine` accepts `elastic_model` and `models` as constructor arguments, so the tests inject
`MagicMock()` objects instead. Keep that door open when you refactor: if you make `Engine` build its
own dependencies unconditionally, the whole suite breaks.

**Watch out for the singletons.** `engine_factory`, `es_client_factory` and `factory` all cache a
module-level global. A test that calls a factory will leak that instance into every later test in the
same process. Construct `Engine(...)` directly in tests, as `test_engine.py` does.

### The characterization suite

`characterization/` is a second, integration-level suite that runs against a **real Elasticsearch**
and pins down exactly what `chatbot` depends on. It needs a reachable cluster:

```bash
cd characterization && python -m pytest . -q
```

It defaults to `http://localhost:9199`; override with `CHARACTERIZATION_ES_URL`. It creates and drops
its own `ragbot_characterization*` indices and never touches application data. If no cluster is
reachable the tests skip rather than fail.

Only the embedding models are stubbed there (deterministic 1024-dimensional vectors) — everything
else is the real query path. Use this suite whenever you change retrieval or indexing behaviour.

CI (`.github/workflows/test.yaml`) runs the mocked unit suite on Python 3.12 for pushes and pull
requests targeting `master`. The characterization suite is not in CI, because it needs an
Elasticsearch instance — run it locally before changing retrieval.

---

## 8. Releasing a New Version

Releasing is **two steps, not one**. Publishing to PyPI does not update anything that consumes the
package — you have to bump the consumers yourself.

### Step 1 — Create the tag

Run the **Create Tag** workflow (`workflow_dispatch`) in GitHub Actions with the major, minor and
patch numbers. It rewrites the version in `setup.py` and `pyproject.toml`, commits, tags, and creates
a GitHub release.

Publishing that release then triggers **Publish to PyPI** (`publish.yaml`), which builds with
`setup.py` and uploads with `twine`. The result appears at
**https://pypi.org/project/kolzchut-ragbot/** — check the version is live there before moving on.

### Step 2 — Update every consumer that pins the version

> ### Do not forget this step
>
> A new version on PyPI changes nothing on its own. Until the consumers are bumped they keep
> installing the old release, and the change you just published simply will not run.

In `chatbot` the version is pinned in **two** files, and both must be updated to the version you just
released:

| File | Line | Format |
| --- | --- | --- |
| `chatbot/requirements.txt` | 24 | `kolzchut-ragbot==<version>` |
| `chatbot/pyproject.toml` | 16 | `kolzchut-ragbot = "<version>"` |

Then reinstall so the environment matches:

```bash
pip install -r requirements.txt
```

Before you start, confirm nothing else has begun pinning the package:

```bash
grep -rn "kolzchut-ragbot" --include="*.txt" --include="*.toml" .
```

> **PyPI is the source of truth for what is released**, not the version strings committed in this
> repository — the tag workflow rewrites them, so they can lag behind. Check the published version
> before pinning.

### Secrets

The PyPI token is stored as the `PYPI_API_TOKEN` GitHub secret, and `GH_ACTIONS_PAT` is used to push
the version commit and create the release. Never put these values in the repository — see the
[internal handover document](https://docs.google.com/document/d/1Vw99KJXy1Cj2etQohB_2PE_puBszK8rSWG8kc9C_HBA).

`builder.sh` is a manual alternative that uploads to **TestPyPI**, not PyPI.
