# Building Custom Models

RoohAI's registry system makes it straightforward to add new STT, TTS, LLM, or VAD implementations. This guide walks through implementing and registering a custom model.

## Model Interface

Every model must subclass one of the four base classes in `roohai/base.py` and implement all abstract methods.

### STTModel

```python
import numpy as np
from roohai.base import STTModel

class MySTT(STTModel):
    META = {
        "display_name": "My STT",
        "description": "Description shown in the UI wizard.",
        "kind": "local",  # or "cloud"
        "pip_dependencies": ["my-stt-library>=1.0"],
        "hf_models": ["my-org/my-model"],       # HuggingFace models to download
        "estimated_download_mb": 500,
        "streaming": False,                       # Set True if real-time streaming is supported
        "config_fields": [
            {
                "name": "model_id",
                "label": "Model ID",
                "type": "text",               # text | secret | select
                "required": False,
                "default": "my-org/my-model",
            },
        ],
    }

    def __init__(self, model_id: str = "my-org/my-model", **kwargs) -> None:
        self.model_id = model_id
        self._model = None

    def load(self) -> None:
        """Load model into memory. Called once at startup."""
        from my_stt_library import load_model
        self._model = load_model(self.model_id)

    def transcribe(self, audio: np.ndarray, sample_rate: int = 16000) -> str:
        """Convert audio to text.

        Args:
            audio: Float32 numpy array, mono, normalized to [-1, 1].
            sample_rate: Sample rate of the audio (typically 16000).

        Returns:
            Transcribed text string.
        """
        return self._model.recognize(audio, sr=sample_rate)

    def unload(self) -> None:
        """Release model from memory."""
        self._model = None

    @property
    def is_loaded(self) -> bool:
        return self._model is not None
```

### TTSModel

```python
import numpy as np
from roohai.base import TTSModel

class MyTTS(TTSModel):
    META = {
        "display_name": "My TTS",
        "description": "Custom text-to-speech engine.",
        "kind": "cloud",
        "pip_dependencies": ["my-tts-sdk"],
        "hf_models": [],
        "estimated_download_mb": 0,
        "config_fields": [
            {
                "name": "api_key",
                "label": "API Key",
                "type": "secret",
                "required": True,
                "env_var": "MY_TTS_API_KEY",
                "placeholder": "sk-xxx",
            },
            {
                "name": "voice",
                "label": "Voice",
                "type": "select",
                "required": False,
                "default": "default",
                "options": ["default", "female-1", "male-1"],
            },
        ],
    }

    def __init__(self, api_key: str | None = None, voice: str = "default", **kwargs) -> None:
        self._api_key = api_key
        self.voice = voice
        self._client = None

    def load(self) -> None:
        import os
        from my_tts_sdk import Client
        key = self._api_key or os.environ.get("MY_TTS_API_KEY")
        if not key:
            raise RuntimeError("API key required")
        self._client = Client(api_key=key)

    def synthesize(self, text: str) -> tuple[np.ndarray, int]:
        """Convert text to speech audio.

        Args:
            text: The text to synthesize.

        Returns:
            Tuple of (audio_array, sample_rate).
            audio_array: Float32 numpy array, mono.
            sample_rate: Sample rate as integer (e.g. 24000).
        """
        audio_bytes = self._client.synthesize(text, voice=self.voice)
        # Parse WAV bytes into numpy array
        import io, soundfile as sf
        data, sr = sf.read(io.BytesIO(audio_bytes), dtype="float32")
        return data, sr

    def unload(self) -> None:
        self._client = None

    @property
    def is_loaded(self) -> bool:
        return self._client is not None
```

### LLMModel

```python
from roohai.base import LLMModel

class MyLLM(LLMModel):
    META = {
        "display_name": "My LLM",
        "description": "Custom language model.",
        "kind": "cloud",
        "pip_dependencies": ["openai>=1.0"],
        "hf_models": [],
        "estimated_download_mb": 0,
        "config_fields": [
            {
                "name": "api_key",
                "label": "API Key",
                "type": "secret",
                "required": True,
                "env_var": "OPENAI_API_KEY",
            },
            {
                "name": "model",
                "label": "Model",
                "type": "select",
                "required": False,
                "default": "gpt-4o-mini",
                "options": ["gpt-4o-mini", "gpt-4o"],
            },
        ],
    }

    def __init__(self, api_key: str | None = None, model: str = "gpt-4o-mini", **kwargs) -> None:
        self._api_key = api_key
        self.model = model
        self._client = None

    def load(self) -> None:
        import os
        from openai import OpenAI
        key = self._api_key or os.environ.get("OPENAI_API_KEY")
        self._client = OpenAI(api_key=key)

    async def chat(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> str:
        """Generate a response to a message.

        Args:
            message: The user's message.
            history: Optional list of {"role": "user"|"assistant", "content": "..."} dicts.
            system_prompt: Optional system prompt override.
            session_id: Unique connection identifier (WebSocket UUID or WebRTC session ID).
                Use this to maintain per-session state like conversation history.

        Returns:
            The assistant's response text.
        """
        import asyncio

        def _call():
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            if history:
                messages.extend(history)
            messages.append({"role": "user", "content": message})

            response = self._client.chat.completions.create(
                model=self.model,
                messages=messages,
            )
            return response.choices[0].message.content

        return await asyncio.to_thread(_call)

    def unload(self) -> None:
        self._client = None

    @property
    def is_loaded(self) -> bool:
        return self._client is not None
```

#### Adding Streaming LLM Support

The base `LLMModel` provides a default `chat_stream()` that yields the full `chat()` result as a single chunk. To enable true token-level streaming (lower latency — TTS starts on the first completed sentence while the LLM is still generating), override `chat_stream()`:

```python
from collections.abc import AsyncIterator

class MyLLM(LLMModel):
    # ... chat() and other methods ...

    async def chat_stream(
        self,
        message: str,
        history: list[dict] | None = None,
        system_prompt: str | None = None,
        session_id: str | None = None,
    ) -> AsyncIterator[str]:
        """Yield text tokens as they arrive from the API."""
        # Example with OpenAI-style streaming:
        import asyncio

        def _stream():
            messages = []
            if system_prompt:
                messages.append({"role": "system", "content": system_prompt})
            if history:
                messages.extend(history)
            messages.append({"role": "user", "content": message})

            for chunk in self._client.chat.completions.create(
                model=self.model,
                messages=messages,
                stream=True,
            ):
                delta = chunk.choices[0].delta.content
                if delta:
                    yield delta

        # Bridge sync generator to async (same pattern as BedrockLLM)
        loop = asyncio.get_event_loop()
        queue: asyncio.Queue[str | None] = asyncio.Queue()

        def _run():
            try:
                for token in _stream():
                    loop.call_soon_threadsafe(queue.put_nowait, token)
            finally:
                loop.call_soon_threadsafe(queue.put_nowait, None)

        fut = loop.run_in_executor(None, _run)
        while True:
            item = await queue.get()
            if item is None:
                break
            yield item
        await fut
```

If you don't override `chat_stream()`, the pipeline will fall back to the non-streaming `chat()` method — everything still works, just without the progressive text display and overlapped TTS.

### LLM Hooks

Instead of implementing a full `LLMModel` subclass, you can replace the LLM via a hook function. Hooks receive the same arguments as `LLMModel.chat()` including `session_id`:

```python
# Batch hook — returns the full response
async def my_hook(
    message: str,
    *,
    history: list[dict] | None = None,
    system_prompt: str | None = None,
    session_id: str | None = None,
) -> str:
    """Use session_id for per-connection state (conversation memory, etc.)."""
    return f"You said: {message}"

# Streaming hook — yields tokens progressively
async def my_stream_hook(
    message: str,
    *,
    history: list[dict] | None = None,
    system_prompt: str | None = None,
    session_id: str | None = None,
):
    for word in f"You said: {message}".split():
        yield word + " "

# Register hooks on the pipeline
pipeline.set_llm_hook(my_hook, stream_hook=my_stream_hook)
```

The `session_id` is a UUID unique to each WebSocket connection or WebRTC session. Use it to maintain per-session state like conversation history in an external store. See the example agents (`examples/voice-latest-news-agent/hooks.py`) for a real-world hook implementation.

### VADModel

```python
import numpy as np
from roohai.base import VADModel

class MyVAD(VADModel):
    META = {
        "display_name": "My VAD",
        "description": "Custom voice activity detector.",
        "kind": "local",
        "pip_dependencies": ["webrtcvad"],
        "hf_models": [],
        "estimated_download_mb": 1,
        "config_fields": [],
    }

    def __init__(self, **kwargs) -> None:
        self._vad = None

    def load(self) -> None:
        import webrtcvad
        self._vad = webrtcvad.Vad(3)  # aggressiveness 0-3

    def is_speech(self, audio_chunk: np.ndarray, sample_rate: int = 16000) -> float:
        """Determine if an audio chunk contains speech.

        Args:
            audio_chunk: Float32 numpy array, mono, normalized to [-1, 1].
            sample_rate: Sample rate (16000 expected).

        Returns:
            Speech probability from 0.0 (silence) to 1.0 (speech).
        """
        pcm = (audio_chunk * 32768).astype(np.int16).tobytes()
        return 1.0 if self._vad.is_speech(pcm, sample_rate) else 0.0

    def unload(self) -> None:
        self._vad = None

    @property
    def is_loaded(self) -> bool:
        return self._vad is not None
```

## Registering Your Model

### Option 1: RoohBuilder (Recommended)

The simplest way to use a custom model — no framework modifications needed. Use the
`custom_stt()`, `custom_tts()`, `custom_llm()`, or `custom_vad()` methods on the builder:

```python
from roohai import Rooh
from my_module import ElevenLabsTTS

pipeline = (
    Rooh.builder()
    .stt("deepgram", model="nova-3")
    .llm("bedrock", model_id="anthropic.claude-3-haiku-20240307-v1:0")
    .custom_tts("elevenlabs", ElevenLabsTTS, api_key="sk-xxx", voice_id="abc123")
    .vad()
    .system_prompt("You are a helpful voice assistant.")
    .build()
)
pipeline.load()
```

You can mix built-in providers (passed as strings) with custom classes freely:

```python
pipeline = (
    Rooh.builder()
    .custom_stt("my-stt", MyStreamingSTT, api_key="key-xxx")
    .custom_llm("my-llm", MyRAGModel, index_name="products")
    .custom_tts("my-tts", ElevenLabsTTS, voice_id="abc")
    .custom_vad("webrtc-vad", WebRTCVAD, aggressiveness=3)
    .build()
)
```

Each `custom_*` method:
1. Accepts a **name** (string identifier), the **class** itself, and any **kwargs** for the constructor
2. Registers the class in the global model registry
3. Sets it as the active model for that pipeline stage

The custom model then participates in the full pipeline (hot-swap, catalog API, metrics, etc.)
just like any built-in model.

### Option 2: Direct Registry API

Register a class directly with the registry for use with YAML config or `swap_model()`:

```python
from roohai.registry import registry
from my_module import MyTTS

registry.register_tts("my-tts", MyTTS, api_key="sk-xxx", voice="female-1")

# Now usable in config or via hot-swap:
pipeline.swap_model("tts", "my-tts")
```

### Option 3: Add to the Framework

For models you want permanently available as built-in providers, add them to the framework source.

Place your model file in the appropriate subpackage:

```
roohai/stt/my_stt.py
roohai/tts/my_tts.py
roohai/llm/my_llm.py
roohai/vad/my_vad.py
```

Update the subpackage `__init__.py`:

```python
# roohai/tts/__init__.py
from .my_tts import MyTTS
__all__ = [..., "MyTTS"]
```

Register in `roohai/pipeline.py`:

```python
# In _register_defaults()
_tts_classes = {
    ...,
    "my-tts": MyTTS,
}

# In _CLASS_MAP
"tts": {
    ...,
    "my-tts": "my-tts",
},
```

## Adding Streaming STT Support

If your STT model supports real-time streaming, set `"streaming": True` in META and implement these additional methods:

```python
class MyStreamingSTT(STTModel):
    META = {
        ...,
        "streaming": True,
    }

    async def start_stream(
        self,
        on_transcript: Callable[[str, bool], None],
        on_utterance_end: Callable[[str], None],
    ) -> None:
        """Open a streaming connection.

        Args:
            on_transcript: Called with (text, is_final) on each transcript update.
            on_utterance_end: Called with (full_text) when a complete utterance is detected.
        """
        ...

    async def send_chunk(self, audio_bytes: bytes) -> None:
        """Send raw PCM16 audio bytes to the streaming connection."""
        ...

    async def stop_stream(self) -> None:
        """Close the streaming connection and clean up."""
        ...
```

The `on_utterance_end` callback **must not block**. The server wraps it in `asyncio.create_task()` to keep the event loop responsive for barge-in detection.

## META Dictionary Reference

| Field | Type | Description |
|-------|------|-------------|
| `display_name` | str | Human-readable name shown in the UI |
| `description` | str | Short description for the wizard |
| `kind` | str | `"local"` or `"cloud"` |
| `pip_dependencies` | list[str] | Python packages to install |
| `hf_models` | list[str] | HuggingFace model IDs to pre-download |
| `estimated_download_mb` | int | Approximate total download size |
| `streaming` | bool | (STT only) Whether real-time streaming is supported |
| `config_fields` | list[dict] | User-configurable parameters (see below) |

### Config Field Schema

```python
{
    "name": "api_key",            # Parameter name (passed to constructor)
    "label": "API Key",           # Display label in UI
    "type": "text",               # "text", "secret", or "select"
    "required": True,             # Whether the field must be provided
    "default": "value",           # Default value (optional)
    "env_var": "MY_API_KEY",      # Environment variable fallback (optional)
    "placeholder": "sk-xxx",      # Input placeholder text (optional)
    "options": ["a", "b"],        # Choices for "select" type (optional)
    "depends_on": {"field": "v"}, # Show only when another field matches (optional)
}
```
