# 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: Add to the Framework

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",
},
```

### Option 2: Register at Runtime

```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:
pipeline.swap_model("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)
}
```
