# Inference Middleware Plugins

Inference middleware plugins run in-process inside the Inference Gateway (IGW).
They intercept every inference request routed through a **VirtualModel** and can:

- **Mutate the request** before it reaches the backend model provider (e.g. keyword filtering, A/B routing, prompt augmentation)
- **Short-circuit the proxy** by returning a response directly, skipping the backend entirely (e.g. a cache hit, a rule-based refusal)
- **Mutate the response** before it reaches the caller (e.g. PII redaction, output sanitisation)
- **Fire-and-forget post-response work** without blocking the caller (e.g. logging, analytics)

---

## 1. How plugins are discovered

Register your middleware class under the `nemo.inference_middleware` entry-point group in `pyproject.toml`:

```toml
[project.entry-points."nemo.inference_middleware"]
"nemo-my-plugin" = "nemo_my_plugin.middleware:MyMiddleware"
```

The entry-point key (`"nemo-my-plugin"`) is the **plugin identity** — it is what
`MiddlewareCall.name` references in VirtualModel configs.  IGW discovers all
registered classes at startup, instantiates them, injects platform cache access,
and calls `on_startup()`.

---

## 2. The `NemoInferenceMiddleware` ABC

Subclass `NemoInferenceMiddleware` and override the hooks you need:

```python
from nemo_plugin.inference_middleware import (
    NemoInferenceMiddleware,
    RequestResult,
    ResponseResult,
    ImmediateResponse,
    InferenceMiddlewareError,
)

class MyMiddleware(NemoInferenceMiddleware):

    async def on_startup(self) -> None:
        """Load ML models, build HTTP clients, validate config at startup."""

    async def on_shutdown(self) -> None:
        """Release resources on graceful shutdown."""

    async def on_virtual_model_upserted(self, virtual_model) -> None:
        """Pre-warm per-VirtualModel resources (e.g. load a classifier for this VM)."""

    async def on_virtual_model_destroyed(self, virtual_model) -> None:
        """Release resources held for this VirtualModel."""

    async def get_middleware_config(self, config_type: str, config_id: str):
        """Fetch a stored config entity from your plugin's entity store."""

    async def validate_middleware_config(self, config_type: str, config) -> ...:
        """Validate and coerce a config object — called at cache-build time, never per-request."""

    async def process_request(self, request_body, request_headers, middleware_config) -> RequestResult:
        """Modify the request before proxying — or short-circuit with ImmediateResponse."""

    async def process_response(self, response_result, request_body, request_headers, response_headers, middleware_config) -> ResponseResult:
        """Modify the response before returning to the caller."""
```

All hooks default to no-ops.  Implement only what your plugin needs.

---

## 3. VirtualModels and MiddlewareCalls

A **VirtualModel** is an entity stored in IGW that maps a logical model name to a
backend model entity and ordered middleware pipelines.  Each pipeline entry is a
**MiddlewareCall**:

```json
{
  "name": "nemo-my-plugin",
  "config_type": "my_plugin_config",
  "config": { "threshold": 0.8 }
}
```

Fields:

| Field | Required | Description |
|---|---|---|
| `name` | ✓ | Entry-point key of the plugin — must match `pyproject.toml` |
| `config_type` | ✓ | Discriminator string — tells the plugin which schema applies |
| `config` | ✗ | Inline config dict. Mutually exclusive with `config_id` |
| `config_id` | ✗ | `"workspace/name"` reference to a stored entity. Mutually exclusive with `config` |

### Inline config vs. config_id

**Use `config` (inline)** when:
- The config is simple and specific to this VirtualModel
- No sharing or versioning is needed
- You don't want to manage a separate entity

**Use `config_id` (entity reference)** when:
- Multiple VirtualModels share the same config
- Operators need to update the config without editing every VirtualModel
- You want automatic propagation — IGW re-resolves on every polling cycle when the entity's `updated_at` changes, without any VirtualModel edit

---

## 4. Implementing `config_id` support

When `MiddlewareCall.config_id` is used, the referenced entity must exist somewhere in the entity store.  Your plugin must:

**Step 1 — Define a config entity** (extends `NemoEntity`, not `EntityBase` directly):

```python
from nemo_plugin.entity import NemoEntity

class MyPluginConfig(NemoEntity, entity_type="my_plugin_config"):
    """Stored config entity for MyMiddleware.

    entity_type must match MiddlewareCall.config_type.
    Use a plugin-scoped name to avoid cross-plugin collisions.
    """
    threshold: float = 0.8
    strong_model: str = "default/llama-70b"
    weak_model: str = "default/llama-8b"
```

**Step 2 — Expose a CRUD API** so operators can create and update configs:

```python
# In your NemoService — full CRUD for MyPluginConfig
# POST /v2/workspaces/{workspace}/my-plugin-configs
# GET  /v2/workspaces/{workspace}/my-plugin-configs
# GET  /v2/workspaces/{workspace}/my-plugin-configs/{name}
# PATCH /v2/workspaces/{workspace}/my-plugin-configs/{name}
# DELETE /v2/workspaces/{workspace}/my-plugin-configs/{name}
```

See `plugin-service` skill and `example-plugin/middleware_service.py` for the full pattern.

**Step 3 — Implement `get_middleware_config`** to fetch from the entity store:

```python
from nemo_plugin.entity_client import NemoEntitiesClient

async def get_middleware_config(self, config_type: str, config_id: str):
    if config_type != MyPluginConfig.__entity_type__:
        raise ValueError(f"Unknown config_type={config_type!r}")

    ws, name = config_id.split("/", 1)
    client = NemoEntitiesClient()
    return await client.get(MyPluginConfig, name=name, workspace=ws)
```

IGW calls `get_middleware_config` at VirtualModel cache-build time and on every
polling cycle.  It is **never called per-request**.

**Step 4 — Implement `validate_middleware_config`** to return the typed config:

```python
async def validate_middleware_config(self, config_type: str, config) -> MyPluginConfigData:
    if config_type != MyPluginConfig.__entity_type__:
        raise ValueError(f"Unknown config_type={config_type!r}")

    if isinstance(config, MyPluginConfig):
        # Convert entity (has name/workspace fields) → lightweight working type
        return MyPluginConfigData(threshold=config.threshold, ...)

    # Inline dict — validate against the working type directly
    return MyPluginConfigData.model_validate(config)
```

> **Tip:** Return a lightweight Pydantic `BaseModel` (not the `NemoEntity`) from
> `validate_middleware_config`.  The reason is practical: `NemoEntity` inherits
> `workspace: str` (required, no default) from `EntityBase`, so calling
> `MyPluginConfig.model_validate({"threshold": 0.8})` — an inline config dict
> that has no `workspace` — raises a validation error.  A separate
> `MyPluginConfigData(BaseModel)` with only the domain fields works for both
> the inline dict path and the entity store path.

---

## 5. Cache accessor methods

After `_inject_cache()` is called (before `on_startup()`), these methods are available:

```python
# All model entity IDs known to IGW: ["default/llama-3b", "default/llama-70b", ...]
self.list_model_entities_for_workspace()
self.list_model_entities_for_workspace("default")  # filter to workspace

# Model entity details (spec, finetuning_type, providers)
entity = self.get_model_entity("default/llama-3b")
if entity and entity.spec:
    print(entity.spec.is_chat, entity.spec.context_size)

# All active providers for a model — useful for latency-aware routing
providers = self.get_model_providers_for_model("default/llama-3b")

# Resolve a model entity to a backend URL and served model name
target = self.get_inference_url_and_model("default/llama-3b")
# target.model_provider_gateway_url → "http://nim-svc:8080/v1"
# target.served_model_name          → "meta/llama-3.2-3b-instruct"

# VirtualModel access — for meta-routing plugins
vm = self.get_virtual_model("default/my-alias")
self.list_virtual_models_for_workspace("default")
```

---

## 6. Request processing

```python
async def process_request(
    self,
    request_body: dict,      # OpenAI Chat Completions payload
    request_headers: dict,   # mutable — mutate in place to add/modify upstream headers
    middleware_config,        # validated config from validate_middleware_config
) -> RequestResult:
```

`request_body["model"]` contains the model entity name set by
`VirtualModel.default_model_entity` (if set).  Middleware may rewrite it freely.
After the request middleware chain completes, IGW reads `request_body["model"]`
and resolves it to a provider.

**Return a dict** to continue proxying (possibly with mutations):
```python
request_body["model"] = "default/strong-model"  # route to a different entity
return request_body
```

**Return `ImmediateResponse`** to skip the proxy entirely:
```python
return ImmediateResponse(data={
    "id": "cached",
    "object": "chat.completion",
    "created": int(time.time()),
    "model": request_body.get("model", ""),
    "choices": [{"index": 0, "message": {"role": "assistant", "content": cached_reply}, "finish_reason": "stop"}],
    "usage": {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0},
})
```

**Raise `InferenceMiddlewareError`** to return a specific HTTP error:
```python
raise InferenceMiddlewareError("Rate limit exceeded", status_code=429)
```

---

## 7. Response processing

```python
async def process_response(
    self,
    response_result: ResponseResult,  # dict (non-streaming) or AsyncIterator[dict] (SSE)
    request_body: dict,               # original request body (read-only by convention)
    request_headers: dict,            # read-only
    response_headers: dict,           # mutable — add/modify headers returned to caller
    middleware_config,
) -> ResponseResult:
```

For non-streaming responses, mutate and return the dict:
```python
if isinstance(response_result, dict):
    response_result["choices"][0]["message"]["content"] = redact(...)
return response_result
```

For streaming, wrap the iterator:
```python
async def _filtered(stream):
    async for chunk in stream:
        yield redact_chunk(chunk)
return _filtered(response_result)
```

---

## 8. Error handling

| Exception | HTTP status | When to use |
|---|---|---|
| `InferenceMiddlewareError(msg, status_code=N)` | N (default 500) | Any handled error with a specific status |
| `InferenceMiddlewareUnavailableError(msg)` | 503 | Upstream service unavailable |
| Unhandled exception | 500 | Let IGW catch and log |

---

## 9. Testing

Mock the cache accessor for unit tests:

```python
from unittest.mock import MagicMock
from nemo_plugin.inference_middleware import InferenceMiddlewareCacheAccessor

cache = MagicMock(spec=InferenceMiddlewareCacheAccessor)
cache.get_model_entity.return_value = None
plugin = MyMiddleware()
plugin._inject_cache(cache)
await plugin.on_startup()
```

See `plugin-testing` skill and `example-plugin/tests/test_inference_middleware.py` for complete examples.

---

## 10. Reference implementation

`plugins/example-plugin/` contains a working keyword content-filter plugin that demonstrates every concept in this document:

| File | What it shows |
|---|---|
| `middleware_config.py` | `NemoEntity` subclass for stored config |
| `middleware.py` | Full `NemoInferenceMiddleware` with both inline and `config_id` support, `ImmediateResponse`, response redaction |
| `middleware_service.py` | CRUD API for the config entity (required for `config_id` support) |
| `tests/test_inference_middleware.py` | Unit tests with mocked cache and entity client |
| `pyproject.toml` | `nemo.inference_middleware` entry-point registration |
