From ed387e876a259f7f3e855830fcbc51d91f434844 Mon Sep 17 00:00:00 2001
From: milo157 <43028253+milo157@users.noreply.github.com>
Date: Tue, 1 Apr 2025 00:03:26 +0200
Subject: [PATCH] Merge pull request #1486 from CerebriumAI/feature/ultravox

Feature/ultravox - bug fixes
---
 .../foundational/07u-interruptible-ultravox.py |  2 +-
 src/pipecat/services/ultravox/stt.py           | 18 ++++++++++++------
 tests/test_piper_tts.py                        |  6 +++---
 tests/test_user_idle_processor.py              |  6 +++---
 4 files changed, 19 insertions(+), 13 deletions(-)

diff --git a/examples/foundational/07u-interruptible-ultravox.py b/examples/foundational/07u-interruptible-ultravox.py
index 1b3f853e0..0c9821196 100644
--- a/examples/foundational/07u-interruptible-ultravox.py
+++ b/examples/foundational/07u-interruptible-ultravox.py
@@ -34,7 +34,7 @@ logger.add(sys.stderr, level="DEBUG")
 # Want to initialize the ultravox processor since it takes time to load the model and dont
 # want to load it every time the pipeline is run
 ultravox_processor = UltravoxSTTService(
-    model_size="fixie-ai/ultravox-v0_4_1-llama-3_1-8b",
+    model_name="fixie-ai/ultravox-v0_5-llama-3_1-8b",
     hf_token=os.getenv("HF_TOKEN"),
 )
 
diff --git a/src/pipecat/services/ultravox/stt.py b/src/pipecat/services/ultravox/stt.py
index ef5cafb3f..9f1cab832 100644
--- a/src/pipecat/services/ultravox/stt.py
+++ b/src/pipecat/services/ultravox/stt.py
@@ -17,6 +17,13 @@ from loguru import logger
 
 from pipecat.frames.frames import (
     AudioRawFrame,
+    LLMFullResponseEndFrame,
+    LLMFullResponseStartFrame,
+    LLMTextFrame,
+    TranscriptionFrame,
+    TextFrame,
+    StartFrame,
+    EndFrame,
     CancelFrame,
     EndFrame,
     ErrorFrame,
@@ -177,7 +184,7 @@ class UltravoxSTTService(AIService):
     to generate text transcriptions.
 
     Args:
-        model_size: The Ultravox model to use (ModelSize enum or string)
+        model_name: The Ultravox model to use (ModelSize enum or string)
         hf_token: Hugging Face token for model access
         temperature: Sampling temperature for generation
         max_tokens: Maximum tokens to generate
@@ -194,7 +201,7 @@ class UltravoxSTTService(AIService):
     def __init__(
         self,
         *,
-        model_size: str = "fixie-ai/ultravox-v0_4_1-llama-3_1-8b",
+        model_name: str = "fixie-ai/ultravox-v0_5-llama-3_1-8b",
         hf_token: Optional[str] = None,
         temperature: float = 0.7,
         max_tokens: int = 100,
@@ -211,7 +218,6 @@ class UltravoxSTTService(AIService):
             logger.warning("No Hugging Face token provided. Model may not load correctly.")
 
         # Initialize model
-        model_name = model_size if isinstance(model_size, str) else model_size.value
         self._model = UltravoxModel(model_name=model_name)
 
         # Initialize service state
@@ -356,10 +362,10 @@ class UltravoxSTTService(AIService):
                     await self.start_ttfb_metrics()
                     await self.start_processing_metrics()
 
-                    async for response in self.model.generate(
+                    async for response in self._model.generate(
                         messages=[{"role": "user", "content": "<|audio|>\n"}],
-                        temperature=self.temperature,
-                        max_tokens=self.max_tokens,
+                        temperature=self._temperature,
+                        max_tokens=self._max_tokens,
                         audio=audio_float32,
                     ):
                         # Stop TTFB metrics after first response
diff --git a/tests/test_piper_tts.py b/tests/test_piper_tts.py
index 673ea087a..f433fea19 100644
--- a/tests/test_piper_tts.py
+++ b/tests/test_piper_tts.py
@@ -133,6 +133,6 @@ async def test_run_piper_tts_error(aiohttp_client):
         up_frames = frames_received[1]
 
         assert isinstance(up_frames[0], ErrorFrame), "Must receive an ErrorFrame for 404"
-        assert "status: 404" in up_frames[0].error, (
-            "ErrorFrame should contain details about the 404"
-        )
+        assert (
+            "status: 404" in up_frames[0].error
+        ), "ErrorFrame should contain details about the 404"
diff --git a/tests/test_user_idle_processor.py b/tests/test_user_idle_processor.py
index a2f2fd386..7ea6f8744 100644
--- a/tests/test_user_idle_processor.py
+++ b/tests/test_user_idle_processor.py
@@ -86,9 +86,9 @@ class TestUserIdleProcessor(unittest.IsolatedAsyncioTestCase):
             expected_down_frames=expected_down_frames,
         )
 
-        assert not callback_called.is_set(), (
-            "Idle callback was called even though bot speaking frames reset the timer"
-        )
+        assert (
+            not callback_called.is_set()
+        ), "Idle callback was called even though bot speaking frames reset the timer"
 
     async def test_idle_retry_callback(self):
         """Test that retry count increases until user activity resets it."""