Metadata-Version: 2.4
Name: datarobot-fastrag
Version: 0.2.10
Summary: Async FastAPI runner for DataRobot custom LLM and RAG models
Author-email: Nickolai Novik <nickolai@datarobot.com>, Mikael Baath <mikael.baath@datarobot.com>, Alex Shoop <alexander.shoop@datarobot.com>
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: AUTHORS
Requires-Dist: fastapi>=0.128.7
Requires-Dist: httpx>=0.28.1
Requires-Dist: openai>=2.20.0
Requires-Dist: pydantic>=2.0.0
Requires-Dist: pydantic-settings>=2.7.1
Requires-Dist: pyyaml>=6.0.3
Requires-Dist: uvicorn[standard]>=0.27.0
Requires-Dist: pandas>=2.2.0
Requires-Dist: python-multipart>=0.0.9
Requires-Dist: opentelemetry-api>=1.23.0
Requires-Dist: opentelemetry-sdk>=1.23.0
Requires-Dist: opentelemetry-instrumentation-fastapi>=0.44b0
Requires-Dist: opentelemetry-exporter-otlp>=1.23.0
Requires-Dist: opentelemetry-instrumentation-openai>=0.52.3
Requires-Dist: datarobot-moderations>=11.2.47
Provides-Extra: build
Requires-Dist: uv~=0.9.6; extra == "build"
Dynamic: license-file

# datarobot-fastrag

FastRAG is an async-native rewrite of [DRUM](https://github.com/datarobot/datarobot-user-models) built on FastAPI. It serves DataRobot custom models using the same `custom.py` hook interface as DRUM, but handles each request on an async event loop instead of blocking a thread. For I/O-bound LLM workloads (vector DB + LLM calls), this typically gives 3–5× higher throughput at the same concurrency level.

Existing DRUM `custom.py` files work without modification — sync hooks are run in a thread pool automatically.

## Features

- `async def chat()` / `async def score()` run natively on the event loop
- Sync hooks offloaded to a `ThreadPoolExecutor` (drop-in compatible with existing models)
- OpenAI-compatible chat completions API (`/v1/chat/completions`, streaming included)
- Predict/score API (`/predict/`)
- OpenTelemetry instrumentation built in
- Deployment prediction stats (Total Predictions and errors) reported to DataRobot
- Chat prompt/response rows reported to DataRobot for Data exploration (Data quality, Data export)
- LLM safety guardrails via [datarobot-moderations](https://pypi.org/project/datarobot-moderations/)

### Not implemented
- Artifact guessing
- Transform and custom tasks API
- `directAccess` routes

## Installation

FastRAG requires Python 3.12 or later.

```bash
pip install datarobot-fastrag
```

Or using uv:

```bash
uv add datarobot-fastrag
```

For local development from source:

```bash
git clone https://github.com/datarobot-oss/datarobot-fastrag
cd datarobot-fastrag
uv pip install -e .
```

## Writing a custom LLM model

FastRAG uses the same `custom.py` hook convention as DRUM. For LLM models, you implement `chat()`. Making it `async` is what unlocks the throughput benefit.

### Step 1 — create your model directory

```
my_model/
├── custom.py
└── model-metadata.yaml
```

`model-metadata.yaml`:

```yaml
name: My LLM model
type: inference
targetType: textgeneration
```

### Step 2 — implement `custom.py`

The minimal interface for a chat model:

```python
# custom.py
import httpx

async def load_model(code_dir: str):
    # Return anything — it is passed as `model` to every hook.
    # Initialise your clients here (LLM, vector DB, etc).
    return httpx.AsyncClient()

async def chat(completion_create_params: dict, model, **kwargs):
    messages = completion_create_params["messages"]
    user_prompt = next(m["content"] for m in reversed(messages) if m["role"] == "user")

    # Both awaits release the event loop, so other requests run concurrently.
    # db_results = await model.post("http://vector-db/search", json={"q": user_prompt})
    # answer = await model.post("http://llm/generate", json={"prompt": ...})

    return {
        "id": "chatcmpl-1",
        "object": "chat.completion",
        "created": 0,
        "model": completion_create_params["model"],
        "choices": [{
            "index": 0,
            "message": {"role": "assistant", "content": f"Echo: {user_prompt}"},
            "finish_reason": "stop",
        }],
        "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
    }
```

### Step 3 — run locally

```bash
fastrag server --code-dir ./my_model
```

Test it:

```bash
curl -X POST http://localhost:8080/v1/chat/completions/ \
  -H "Content-Type: application/json" \
  -d '{"model": "my-model", "messages": [{"role": "user", "content": "hello"}]}'
```

Or with the OpenAI Python client:

```python
from openai import AsyncOpenAI

client = AsyncOpenAI(base_url="http://localhost:8080/v1", api_key="unused")
response = await client.chat.completions.create(
    model="my-model",
    messages=[{"role": "user", "content": "hello"}],
)
print(response.choices[0].message.content)
```

### Step 4 — deploy on DataRobot

Upload `custom.py` and `model-metadata.yaml` as a custom model in the DataRobot UI, and select the **[GenAI] Python 3.12 with Moderations** execution environment. When the `GENAI_RAG_FRAG_RUNNER` platform flag is enabled for your org, the environment starts FastRAG automatically; otherwise it falls back to DRUM.

### All supported hooks

| Hook | Signature | Notes |
|---|---|---|
| `load_model` | `(code_dir: str) -> Any` | Return value becomes `model` in all other hooks. Runs once at startup. |
| `init` | `(code_dir: str) -> None` | Side-effecting setup (logging, connections). Runs before `load_model`. |
| `chat` | `(completion_create_params: dict, model: Any, **kwargs) -> dict \| Iterator` | OpenAI chat completions. Return a dict or a (sync/async) generator for streaming. |
| `score` | `(data: pd.DataFrame, model: Any, **kwargs) -> pd.DataFrame` | Tabular predictions. |
| `score_unstructured` | `(data: Any, model: Any, **kwargs) -> Any` | Raw bytes in/out. |
| `get_supported_llm_models` | `(model: Any) -> list[Model]` | Populates `/v1/models`. |

All hooks can be `async def` or plain `def`. Sync hooks run in a thread pool.

### Streaming

Return a generator from `chat()` that yields `ChatCompletionChunk` objects (or plain dicts). Both sync generators and `async` generators work:

```python
async def chat(completion_create_params, model, **kwargs):
    if completion_create_params.get("stream"):
        async def gen():
            for token in ["Hello", " world"]:
                yield {"choices": [{"delta": {"content": token}, "finish_reason": None, "index": 0}]}
            yield {"choices": [{"delta": {}, "finish_reason": "stop", "index": 0}]}
        return gen()
    # non-streaming fallback ...
```

A complete working example (including streaming) is at [`tests/models/python3_dummy_chat/custom.py`](tests/models/python3_dummy_chat/custom.py).

## Configuration

Configuration can be provided via CLI arguments or environment variables:

| CLI Argument | Environment Variable | Description |
|--------------|----------------------|-------------|
| `--code-dir` | `CODE_DIR` | Directory containing `custom.py` and model files. |
| `--address` | `ADDRESS` | Host and port to bind to (e.g., `0.0.0.0:8080`). |
| `--max-workers` | `MAX_WORKERS` | Number of worker processes. |
| `--verbose` | `VERBOSE` | Enable verbose logging. |

### Prediction stats reporting

FastRAG reports chat-completion prediction counts to DataRobot using the same env
vars as DRUM (`EXTERNAL_WEB_SERVER_URL` / `API_TOKEN` / `DEPLOYMENT_ID` /
`MODEL_ID`; `MLOPS_*` aliases are accepted). Each chat request counts as 1, and
4xx/5xx are reported as `userError` / `systemError` with 0 predictions. Records
are queued and POSTed in batches; reporting never blocks a request. Reporting is
off only when those credentials are missing (typical for local runs).

Responses carry an `X-Drum-Version` header (`1.17.12`, overridable with
`FASTRAG_DRUM_VERSION`) whenever reporting is on. DataRobot's predictions gateway
reads that header to decide whether the model reports its own chat monitoring; with
no header it assumes any 5xx went unreported and files a record of its own that
carries one prediction, so a failed chat request would land twice in `Total
Requests` and once in `Total Predictions`. The header is omitted when reporting is
off, which keeps the gateway's coverage for deployments that report nothing.

### Chat prediction data reporting

Like DRUM, FastRAG also reports one prompt/response row per chat completion to
`/api/v2/deployments/<id>/predictionInputs/fromJSON/`, using the same credentials.
These rows back the deployment's Data exploration (Data quality, Data export) and
join to custom metrics by association ID. The ID is the request's
`datarobot_association_id` or a generated one, and is returned on the response
and on every stream chunk. The prompt column name is read from the deployment at
startup, and DataRobot keeps the rows only while the deployment's prediction row
storage is enabled. Streamed responses are reported once the stream finishes;
failed requests are not reported.

Structured `/predict/` requests are deliberately **not** reported. Those go
through DataRobot's prediction API, which counts the scored rows itself and feeds
them into deployment stats.

## Development

The project uses `uv` for dependency management and a `Makefile` for common tasks.

To run the tests:

```bash
make test
```

To run with coverage:

```bash
make cov
```

To format the code:

```bash
make fmt
```

To clean up build and test artifacts:

```bash
make clean
```
