refactor(rtzr): FlushSentinel-based segment control and type safety improvements

CONTRIBUTING.md

@@ -81,6 +81,7 @@ uv pip install pip && uv run mypy --install-types --non-interactive \
     -p livekit.plugins.nltk \
     -p livekit.plugins.resemble \
     -p livekit.plugins.rime \
+    -p livekit.plugins.rtzr \
     -p livekit.plugins.silero \
     -p livekit.plugins.speechify \
     -p livekit.plugins.speechmatics \

livekit-plugins/livekit-plugins-rtzr/livekit/plugins/rtzr/__init__.py

@@ -16,7 +16,7 @@
 
 
 class RTZRPlugin(Plugin):
-    def __init__(self):
+    def __init__(self) -> None:
         super().__init__(__name__, __version__, __package__, logger)
 
 

livekit-plugins/livekit-plugins-rtzr/livekit/plugins/rtzr/py.typed

@@ -0,0 +1 @@
+

livekit-plugins/livekit-plugins-rtzr/livekit/plugins/rtzr/rtzrapi.py

@@ -4,7 +4,8 @@
 import os
 import time
 from collections.abc import Iterable
-from typing import Any
+from types import TracebackType
+from typing import TypedDict
 
 import aiohttp
 
@@ -41,11 +42,14 @@ class RTZRTimeoutError(RTZRAPIError):
 DEFAULT_SAMPLE_RATE = 8000
 
 
-def _format_keywords(
-    keywords: Iterable[str] | Iterable[tuple[str, float]],
-) -> str:
+class _Token(TypedDict):
+    access_token: str
+    expire_at: float
+
+
+def _format_keywords(keywords: Iterable[str | tuple[str, float]]) -> str:
     formatted: list[str] = []
-    keyword_list = list(keywords)
+    keyword_list: list[str | tuple[str, float]] = list(keywords)
     if len(keyword_list) > 100:
         raise ValueError("RTZR keyword boosting supports up to 100 keywords")
 
@@ -54,13 +58,18 @@ def _format_keywords(
             if len(item) != 2:
                 raise ValueError("RTZR keyword boosting tuples must be (keyword, boost)")
             word, boost = item
+            if not isinstance(word, str):
+                raise ValueError("RTZR keyword boosting keywords must be strings")
+            if not isinstance(boost, (int, float)):
+                raise ValueError("RTZR keyword boost must be a number")
             if not word:
                 raise ValueError("RTZR keyword boosting keywords must be non-empty")
             if len(word) > 20:
                 raise ValueError("RTZR keyword boosting keywords must be <= 20 chars")
-            if boost < -5.0 or boost > 5.0:
+            boost_value = float(boost)
+            if boost_value < -5.0 or boost_value > 5.0:
                 raise ValueError("RTZR keyword boost must be between -5.0 and 5.0")
-            formatted.append(f"{word}:{boost}")
+            formatted.append(f"{word}:{boost_value}")
             continue
 
         if not isinstance(item, str):
@@ -117,7 +126,6 @@ def __init__(
         client_secret: str | None = None,
         http_session: aiohttp.ClientSession | None = None,
     ) -> None:
-        self._logger = logging.getLogger(__name__)
         self.client_id = client_id or os.environ.get("RTZR_CLIENT_ID")
         self.client_secret = client_secret or os.environ.get("RTZR_CLIENT_SECRET")
 
@@ -126,23 +134,32 @@ def __init__(
 
         self._http_session = http_session
         self._owns_session = http_session is None  # Track if we own the session
-        self._token: dict[str, Any] | None = None
+        self._token: _Token | None = None
         self._api_base = "https://openapi.vito.ai"
         self._ws_base = "wss://" + self._api_base.split("://", 1)[1]
 
-    async def __aenter__(self):
+    async def __aenter__(self) -> RTZROpenAPIClient:
         """Async context manager entry."""
         return self
 
-    async def __aexit__(self, exc_type, exc_val, exc_tb):
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc_val: BaseException | None,
+        exc_tb: TracebackType | None,
+    ) -> None:
         """Async context manager exit."""
         await self.close()
 
     async def get_token(self) -> str:
         """Get a valid access token, refreshing if necessary."""
-        if self._token is None or self._token["expire_at"] < time.time() - 3600:
+        token = self._token
+        if token is None or token["expire_at"] < time.time() - 3600:
             await self._refresh_token()
-        return self._token["access_token"]
+            token = self._token
+        if token is None:
+            raise RTZRAPIError("Failed to obtain RTZR access token")
+        return token["access_token"]
 
     async def _refresh_token(self) -> None:
         """Refresh the access token."""
@@ -155,7 +172,13 @@ async def _refresh_token(self) -> None:
             ) as resp:
                 resp.raise_for_status()
                 data = await resp.json()
-                self._token = data
+                if not isinstance(data, dict):
+                    raise RTZRStatusError("Invalid token response payload")
+                access_token = data.get("access_token")
+                expire_at = data.get("expire_at")
+                if not isinstance(access_token, str) or not isinstance(expire_at, (int, float)):
+                    raise RTZRStatusError("Invalid token response payload")
+                self._token = {"access_token": access_token, "expire_at": float(expire_at)}
                 logger.debug("Successfully refreshed RTZR access token")
         except aiohttp.ClientResponseError as e:
             logger.error("RTZR authentication failed: %s %s", e.status, e.message)
@@ -219,7 +242,7 @@ def build_config(
         noise_threshold: float = 0.60,
         active_threshold: float = 0.80,
         use_punctuation: bool = False,
-        keywords: Iterable[str] | Iterable[tuple[str, float]] | None = None,
+        keywords: Iterable[str | tuple[str, float]] | None = None,
     ) -> dict[str, str]:
         """Build configuration dictionary for WebSocket connection."""
         config = {