datarobot-fastrag

FastRAG is an async-native rewrite of DRUM 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
• LLM safety guardrails via datarobot-moderations

Not implemented

• Artifact guessing
• Transform and custom tasks API
• directAccess routes

Installation

FastRAG requires Python 3.12 or later.

pip install datarobot-fastrag

Or using uv:

uv add datarobot-fastrag

For local development from source:

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:

name: My LLM model
type: inference
targetType: textgeneration

Step 2 — implement custom.py

The minimal interface for a chat model:

# 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

fastrag server --code-dir ./my_model

Test it:

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:

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

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:

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.

Configuration

Configuration can be provided via CLI arguments or environment variables:

Development

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

To run the tests:

make test

To run with coverage:

make cov

To format the code:

make fmt

To clean up build and test artifacts:

make clean