Merge branch 'pipecat-ai:main' into fix/smallwebrtcrequesthandler-return-type
This commit is contained in:
Kingston Kuan
@@ -1,13 +0,0 @@
|
||||
#
|
||||
# Copyright (c) 2024–2025, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
import sys
|
||||
|
||||
from pipecat.services import DeprecatedModuleProxy
|
||||
|
||||
from .stt import *
|
||||
|
||||
sys.modules[__name__] = DeprecatedModuleProxy(globals(), "ultravox", "ultravox.stt")
|
||||
|
||||
549
src/pipecat/services/ultravox/llm.py
Normal file
549
src/pipecat/services/ultravox/llm.py
Normal file
@@ -0,0 +1,549 @@
|
||||
#
|
||||
# Copyright (c) 2024–2025, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""Ultravox Realtime API service implementation.
|
||||
|
||||
This module provides real-time conversational AI capabilities using Ultravox's
|
||||
Realtime API, supporting both text and audio modalities with
|
||||
voice transcription, streaming responses, and tool usage.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import datetime
|
||||
import json
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Literal, Optional, Union
|
||||
|
||||
import aiohttp
|
||||
from loguru import logger
|
||||
from openai.types import chat as openai_chat_types
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
from pipecat.adapters.schemas.tools_schema import ToolsSchema
|
||||
from pipecat.audio.utils import create_stream_resampler
|
||||
from pipecat.frames.frames import (
|
||||
AggregationType,
|
||||
CancelFrame,
|
||||
EndFrame,
|
||||
Frame,
|
||||
InputAudioRawFrame,
|
||||
InputTextRawFrame,
|
||||
LLMContextFrame,
|
||||
LLMFullResponseEndFrame,
|
||||
LLMFullResponseStartFrame,
|
||||
LLMTextFrame,
|
||||
LLMUpdateSettingsFrame,
|
||||
StartFrame,
|
||||
TranscriptionFrame,
|
||||
TTSAudioRawFrame,
|
||||
TTSStartedFrame,
|
||||
TTSStoppedFrame,
|
||||
TTSTextFrame,
|
||||
UserAudioRawFrame,
|
||||
)
|
||||
from pipecat.processors.aggregators.llm_context import LLMContext
|
||||
from pipecat.processors.aggregators.llm_response import (
|
||||
LLMAssistantAggregatorParams,
|
||||
LLMUserAggregatorParams,
|
||||
)
|
||||
from pipecat.processors.aggregators.llm_response_universal import LLMContextAggregatorPair
|
||||
from pipecat.processors.aggregators.openai_llm_context import (
|
||||
OpenAILLMContext,
|
||||
OpenAILLMContextFrame,
|
||||
)
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
from pipecat.services.llm_service import FunctionCallFromLLM, LLMService
|
||||
from pipecat.utils.time import time_now_iso8601
|
||||
|
||||
try:
|
||||
from websockets.asyncio import client as websocket_client
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
logger.error("In order to use Ultravox, you need to `pip install pipecat-ai[ultravox]`.")
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
class AgentInputParams(BaseModel):
|
||||
"""Input parameters for Ultravox Realtime generation using a pre-defined Agent.
|
||||
|
||||
Parameters:
|
||||
api_key: Ultravox API key for authentication.
|
||||
agent_id: The ID of the Ultravox Realtime agent you'd like to use. Agents
|
||||
are pre-configured to handle calls consistently. You can create and edit
|
||||
agents in the Ultravox console (https://app.ultravox.ai/agents) or using
|
||||
the Ultravox API (https://docs.ultravox.ai/api-reference/agents/agents-post).
|
||||
template_context: Context variables to use when instantiating a call with the
|
||||
agent. Defaults to an empty dict.
|
||||
metadata: Metadata to attach to the call. Default to an empty dict.
|
||||
max_duration: The maximum duration of the call. Defaults to None, which will
|
||||
use the agent's default maximum duration.
|
||||
extra: Extra parameters to include in the agent call creation request. Defaults
|
||||
to an empty dict. See the Ultravox API documentation for valid arguments:
|
||||
https://docs.ultravox.ai/api-reference/agents/agents-calls-post
|
||||
"""
|
||||
|
||||
api_key: str
|
||||
agent_id: uuid.UUID
|
||||
template_context: Dict[str, Any] = Field(default_factory=dict)
|
||||
metadata: Dict[str, str] = Field(default_factory=dict)
|
||||
max_duration: Optional[datetime.timedelta] = Field(
|
||||
default=None, ge=datetime.timedelta(seconds=10), le=datetime.timedelta(hours=1)
|
||||
)
|
||||
extra: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class OneShotInputParams(BaseModel):
|
||||
"""Input parameters for Ultravox Realtime generation using a one-off call.
|
||||
|
||||
Parameters:
|
||||
api_key: Ultravox API key for authentication.
|
||||
system_prompt: System prompt to guide the model's behavior. Defaults to None.
|
||||
temperature: Sampling temperature for response generation. Defaults to 0.
|
||||
model: Model identifier to use. Defaults to "fixie-ai/ultravox".
|
||||
voice: Voice identifier for speech generation. Defaults to None.
|
||||
metadata: Metadata to attach to the call. Default to an empty dict.
|
||||
max_duration: The maximum duration of the call. Defaults to one hour.
|
||||
extra: Extra parameters to include in the call creation request. Defaults
|
||||
to an empty dict. See the Ultravox API documentation for valid arguments:
|
||||
https://docs.ultravox.ai/api-reference/calls/calls-post
|
||||
"""
|
||||
|
||||
api_key: str
|
||||
system_prompt: Optional[str] = None
|
||||
temperature: float = Field(default=0.0, ge=0.0, le=1.0)
|
||||
model: Optional[str] = None
|
||||
voice: Optional[uuid.UUID] = None
|
||||
metadata: Dict[str, str] = Field(default_factory=dict)
|
||||
max_duration: datetime.timedelta = Field(
|
||||
default=datetime.timedelta(hours=1),
|
||||
ge=datetime.timedelta(seconds=10),
|
||||
le=datetime.timedelta(hours=1),
|
||||
)
|
||||
extra: Dict[str, Any] = Field(default_factory=dict)
|
||||
|
||||
|
||||
class JoinUrlInputParams(BaseModel):
|
||||
"""Input parameters for joining an existing Ultravox Realtime call via join URL.
|
||||
|
||||
Parameters:
|
||||
join_url: The join URL for the existing Ultravox Realtime call.
|
||||
"""
|
||||
|
||||
join_url: str
|
||||
|
||||
|
||||
class UltravoxRealtimeLLMService(LLMService):
|
||||
"""Provides access to the Ultravox Realtime API.
|
||||
|
||||
This service enables real-time conversations with Ultravox, supporting both
|
||||
text and audio output. It handles voice transcription, streaming audio
|
||||
responses, and tool usage.
|
||||
|
||||
Note: Ultravox is an audio-native model, so voice transcriptions are not used
|
||||
by the model and may not always align with its understanding of user input.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
params: Union[AgentInputParams, OneShotInputParams, JoinUrlInputParams],
|
||||
one_shot_selected_tools: Optional[ToolsSchema] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the Ultravox Realtime LLM service.
|
||||
|
||||
Args:
|
||||
api_key: Ultravox API key for authentication.
|
||||
params: Configuration parameters for the model.
|
||||
one_shot_selected_tools: ToolsSchema for tools to use with this call.
|
||||
May only be set with OneShotInputParams.
|
||||
**kwargs: Additional arguments passed to parent LLMService.
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
self._params = params
|
||||
if one_shot_selected_tools:
|
||||
if not isinstance(self._params, OneShotInputParams):
|
||||
logger.warning(
|
||||
"one_shot_selected_tools may only be set when using OneShotInputParams; ignoring."
|
||||
)
|
||||
else:
|
||||
self._selected_tools = one_shot_selected_tools
|
||||
|
||||
self._socket: Optional[websocket_client.ClientConnection] = None
|
||||
self._receive_task: Optional[asyncio.Task] = None
|
||||
self._disconnecting = False
|
||||
self._bot_responding: Literal[None, "text", "voice"] = None
|
||||
self._last_user_id: Optional[str] = None
|
||||
|
||||
self._sample_rate = 48000
|
||||
self._resampler = create_stream_resampler()
|
||||
|
||||
#
|
||||
# standard AIService frame handling
|
||||
#
|
||||
|
||||
async def start(self, frame: StartFrame):
|
||||
"""Start the service and establish connection.
|
||||
|
||||
Args:
|
||||
frame: The start frame.
|
||||
"""
|
||||
await super().start(frame)
|
||||
|
||||
try:
|
||||
match self._params:
|
||||
case JoinUrlInputParams():
|
||||
join_url = self._params.join_url
|
||||
case AgentInputParams():
|
||||
join_url = await self._start_agent_call(self._params)
|
||||
case OneShotInputParams():
|
||||
join_url = await self._start_one_shot_call(self._params)
|
||||
|
||||
logger.info(f"Joining Ultravox Realtime call via URL: {join_url}")
|
||||
self._socket = await websocket_client.connect(join_url)
|
||||
self._receive_task = self.create_task(self._receive_messages())
|
||||
except Exception as e:
|
||||
await self.push_error("Failed to connect to Ultravox", e, fatal=True)
|
||||
|
||||
async def _start_agent_call(self, params: AgentInputParams) -> str:
|
||||
request_body = {
|
||||
"templateContext": params.template_context,
|
||||
"metadata": params.metadata,
|
||||
"medium": {
|
||||
"serverWebSocket": {
|
||||
"inputSampleRate": self._sample_rate,
|
||||
}
|
||||
},
|
||||
}
|
||||
if params.max_duration:
|
||||
request_body["maxDuration"] = f"{params.max_duration.total_seconds():3f}s"
|
||||
request_body = request_body | params.extra
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
f"https://api.ultravox.ai/api/agents/{params.agent_id}/calls",
|
||||
headers={"X-Api-Key": params.api_key},
|
||||
json=request_body,
|
||||
) as response:
|
||||
if response.status != 201:
|
||||
error_text = await response.text()
|
||||
raise Exception(f"Ultravox API error {response.status}: {error_text}")
|
||||
return (await response.json())["joinUrl"]
|
||||
|
||||
async def _start_one_shot_call(self, params: OneShotInputParams) -> str:
|
||||
request_body = {
|
||||
"systemPrompt": params.system_prompt,
|
||||
"temperature": params.temperature,
|
||||
"model": params.model,
|
||||
"voice": str(params.voice) if params.voice else None,
|
||||
"metadata": params.metadata,
|
||||
"maxDuration": f"{params.max_duration.total_seconds():3f}s",
|
||||
"selectedTools": self._to_selected_tools(self._selected_tools)
|
||||
if self._selected_tools
|
||||
else [],
|
||||
"medium": {
|
||||
"serverWebSocket": {
|
||||
"inputSampleRate": self._sample_rate,
|
||||
}
|
||||
},
|
||||
} | params.extra
|
||||
async with aiohttp.ClientSession() as session:
|
||||
async with session.post(
|
||||
"https://api.ultravox.ai/api/calls",
|
||||
headers={"X-Api-Key": params.api_key},
|
||||
json=request_body,
|
||||
) as response:
|
||||
if response.status != 201:
|
||||
error_text = await response.text()
|
||||
raise Exception(f"Ultravox API error {response.status}: {error_text}")
|
||||
return (await response.json())["joinUrl"]
|
||||
|
||||
def _to_selected_tools(self, tool: ToolsSchema) -> List[Dict[str, Any]]:
|
||||
result: List[Dict[str, Any]] = []
|
||||
for standard_tool in tool.standard_tools:
|
||||
result.append(
|
||||
{
|
||||
"temporaryTool": {
|
||||
"modelToolName": standard_tool.name,
|
||||
"description": standard_tool.description,
|
||||
"dynamicParameters": [
|
||||
{
|
||||
"name": k,
|
||||
"location": "PARAMETER_LOCATION_BODY",
|
||||
"schema": v,
|
||||
"required": k in standard_tool.required,
|
||||
}
|
||||
for k, v in standard_tool.properties.items()
|
||||
],
|
||||
"client": {},
|
||||
}
|
||||
}
|
||||
)
|
||||
return result
|
||||
|
||||
async def stop(self, frame: EndFrame):
|
||||
"""Stop the service and close connections.
|
||||
|
||||
Args:
|
||||
frame: The end frame.
|
||||
"""
|
||||
await super().stop(frame)
|
||||
await self._disconnect()
|
||||
|
||||
async def cancel(self, frame: CancelFrame):
|
||||
"""Cancel the service and close connections.
|
||||
|
||||
Args:
|
||||
frame: The cancel frame.
|
||||
"""
|
||||
await super().cancel(frame)
|
||||
await self._disconnect()
|
||||
|
||||
async def _disconnect(self):
|
||||
self._disconnecting = True
|
||||
if self._socket:
|
||||
await self._socket.close()
|
||||
self._socket = None
|
||||
if self._receive_task:
|
||||
await self.cancel_task(self._receive_task, timeout=1.0)
|
||||
self._receive_task = None
|
||||
|
||||
#
|
||||
# frame processing
|
||||
# StartFrame, StopFrame, CancelFrame implemented in base class
|
||||
#
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
"""Process incoming frames for the Ultravox Realtime service.
|
||||
|
||||
Args:
|
||||
frame: The frame to process.
|
||||
direction: The frame processing direction.
|
||||
"""
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
if isinstance(frame, (LLMContextFrame, OpenAILLMContextFrame)):
|
||||
context = (
|
||||
frame.context
|
||||
if isinstance(frame, LLMContextFrame)
|
||||
else LLMContext.from_openai_context(frame.context)
|
||||
)
|
||||
await self._handle_context(context)
|
||||
elif isinstance(frame, LLMUpdateSettingsFrame):
|
||||
if "output_medium" in frame.settings:
|
||||
await self._update_output_medium(frame.settings.get("output_medium"))
|
||||
elif isinstance(frame, InputTextRawFrame):
|
||||
await self._send_user_text(frame.text)
|
||||
await self.push_frame(frame, direction)
|
||||
elif isinstance(frame, InputAudioRawFrame):
|
||||
await self._send_user_audio(frame)
|
||||
await self.push_frame(frame, direction)
|
||||
else:
|
||||
await self.push_frame(frame, direction)
|
||||
|
||||
async def _handle_context(self, context: LLMContext):
|
||||
# Ultravox handles all context server-side, so the only context we may
|
||||
# need to handle here is new function call results.
|
||||
for message in reversed(context.messages):
|
||||
if message.get("role") != "tool":
|
||||
break
|
||||
content = message.get("content")
|
||||
socket_message = {
|
||||
"type": "client_tool_result",
|
||||
"invocationId": message.get("tool_call_id"),
|
||||
"result": content
|
||||
if isinstance(content, str)
|
||||
else "".join(t.get("text") for t in content),
|
||||
}
|
||||
await self._send(socket_message)
|
||||
|
||||
async def _send_user_audio(self, frame: InputAudioRawFrame):
|
||||
"""Send user audio frame to Ultravox Realtime."""
|
||||
if not self._socket:
|
||||
return
|
||||
self._last_user_id = frame.user_id if isinstance(frame, UserAudioRawFrame) else None
|
||||
audio = frame.audio
|
||||
if frame.sample_rate != self._sample_rate:
|
||||
audio = await self._resampler.resample(audio, frame.sample_rate, self._sample_rate)
|
||||
await self._send(audio)
|
||||
|
||||
async def _send_user_text(self, text: str):
|
||||
"""Send user text via Ultravox Realtime.
|
||||
|
||||
Args:
|
||||
text: The text to send as user input.
|
||||
"""
|
||||
if not self._socket:
|
||||
return
|
||||
await self._send({"type": "user_text_message", "text": text})
|
||||
|
||||
async def _update_output_medium(self, output_medium: str):
|
||||
output_medium = output_medium.lower()
|
||||
if output_medium == "audio":
|
||||
output_medium = "voice"
|
||||
if output_medium.lower() not in {"voice", "text"}:
|
||||
logger.warning(f"Unsupported Ultravox output medium: {output_medium}")
|
||||
return
|
||||
await self._send({"type": "set_output_medium", "medium": output_medium})
|
||||
|
||||
async def _send(self, content: Union[bytes, Dict[str, Any]]):
|
||||
"""Send content via the WebSocket connection.
|
||||
|
||||
Args:
|
||||
content: The content to send, either as bytes or a JSON-serializable dict.
|
||||
"""
|
||||
if self._disconnecting or not self._socket:
|
||||
return
|
||||
|
||||
try:
|
||||
if isinstance(content, bytes):
|
||||
await self._socket.send(content)
|
||||
else:
|
||||
await self._socket.send(json.dumps(content))
|
||||
except Exception as e:
|
||||
if self._disconnecting or not self._socket:
|
||||
return
|
||||
await self.push_error("Ultravox websocket send error", e, fatal=True)
|
||||
|
||||
#
|
||||
# response handling
|
||||
#
|
||||
|
||||
async def _receive_messages(self):
|
||||
"""Receive messages from the Ultravox Realtime WebSocket."""
|
||||
if not self._socket:
|
||||
return
|
||||
async for message in self._socket:
|
||||
try:
|
||||
if isinstance(message, bytes):
|
||||
await self._handle_audio(message)
|
||||
continue
|
||||
|
||||
data = json.loads(message)
|
||||
match data.get("type"):
|
||||
case "state":
|
||||
if self._bot_responding and data.get("state") != "speaking":
|
||||
await self._handle_response_end()
|
||||
case "client_tool_invocation":
|
||||
await self._handle_tool_invocation(
|
||||
data.get("toolName"), data.get("invocationId"), data.get("parameters")
|
||||
)
|
||||
case "transcript":
|
||||
match data.get("role"):
|
||||
case "user":
|
||||
if not data.get("final"):
|
||||
logger.warning(
|
||||
"Unexpected non-final user transcript from Ultravox Realtime; ignoring."
|
||||
)
|
||||
else:
|
||||
await self._handle_user_transcript(data.get("text"))
|
||||
case "agent":
|
||||
await self._handle_agent_transcript(
|
||||
data.get("medium"),
|
||||
data.get("text"),
|
||||
data.get("delta"),
|
||||
data.get("final", False),
|
||||
)
|
||||
case _:
|
||||
logger.debug(
|
||||
f"Received transcript with unknown role from Ultravox Realtime: {data}"
|
||||
)
|
||||
case _:
|
||||
logger.debug(f"Received unhandled Ultravox message: {data}")
|
||||
except Exception as e:
|
||||
if self._disconnecting or not self._socket:
|
||||
return
|
||||
await self.push_error("Ultravox websocket receive error", e, fatal=True)
|
||||
|
||||
async def _handle_audio(self, audio: bytes):
|
||||
"""Handle incoming audio bytes from Ultravox Realtime."""
|
||||
if not audio:
|
||||
return
|
||||
if not self._bot_responding:
|
||||
await self.push_frame(LLMFullResponseStartFrame())
|
||||
await self.push_frame(TTSStartedFrame())
|
||||
self._bot_responding = "voice"
|
||||
await self.push_frame(TTSAudioRawFrame(audio, self._sample_rate, 1))
|
||||
|
||||
async def _handle_response_end(self):
|
||||
if self._bot_responding == "voice":
|
||||
await self.push_frame(TTSStoppedFrame())
|
||||
await self.push_frame(LLMFullResponseEndFrame())
|
||||
self._bot_responding = None
|
||||
|
||||
async def _handle_tool_invocation(
|
||||
self, tool_name: str, invocation_id: str, parameters: Dict[str, Any]
|
||||
):
|
||||
await self.run_function_calls(
|
||||
[
|
||||
FunctionCallFromLLM(
|
||||
function_name=tool_name,
|
||||
tool_call_id=invocation_id,
|
||||
arguments=parameters,
|
||||
context=None,
|
||||
)
|
||||
]
|
||||
)
|
||||
|
||||
async def _handle_user_transcript(self, text: str):
|
||||
await self.push_frame(
|
||||
TranscriptionFrame(
|
||||
user_id=self._last_user_id or "",
|
||||
timestamp=time_now_iso8601(),
|
||||
result=text,
|
||||
text=text,
|
||||
),
|
||||
FrameDirection.UPSTREAM,
|
||||
)
|
||||
|
||||
async def _handle_agent_transcript(
|
||||
self, medium: str, text: Optional[str], delta: Optional[str], final: bool
|
||||
):
|
||||
if text or delta:
|
||||
frame = LLMTextFrame(text=text or delta)
|
||||
frame.skip_tts = medium == "voice"
|
||||
await self.push_frame(frame)
|
||||
if medium == "text":
|
||||
if text:
|
||||
await self.push_frame(LLMFullResponseStartFrame())
|
||||
await self.push_frame(TTSStartedFrame())
|
||||
await self.push_frame(TTSTextFrame(text=text, aggregated_by=AggregationType.WORD))
|
||||
self._bot_responding = "text"
|
||||
elif final:
|
||||
await self.push_frame(LLMFullResponseEndFrame())
|
||||
self._bot_responding = None
|
||||
elif delta:
|
||||
await self.push_frame(TTSTextFrame(text=delta, aggregated_by=AggregationType.WORD))
|
||||
|
||||
def create_context_aggregator(
|
||||
self,
|
||||
context: OpenAILLMContext,
|
||||
*,
|
||||
user_params: LLMUserAggregatorParams = LLMUserAggregatorParams(),
|
||||
assistant_params: LLMAssistantAggregatorParams = LLMAssistantAggregatorParams(),
|
||||
) -> LLMContextAggregatorPair:
|
||||
"""Create an instance of LLMContextAggregatorPair from an OpenAILLMContext.
|
||||
|
||||
Constructor keyword arguments for both the user and assistant aggregators can be provided.
|
||||
|
||||
NOTE: this method exists only for backward compatibility. New code
|
||||
should instead do::
|
||||
|
||||
context = LLMContext(...)
|
||||
context_aggregator = LLMContextAggregatorPair(context)
|
||||
|
||||
Args:
|
||||
context: The LLM context to use.
|
||||
user_params: User aggregator parameters. Defaults to LLMUserAggregatorParams().
|
||||
assistant_params: Assistant aggregator parameters. Defaults to LLMAssistantAggregatorParams().
|
||||
|
||||
Returns:
|
||||
A pair of user and assistant context aggregators.
|
||||
"""
|
||||
context = LLMContext.from_openai_context(context)
|
||||
assistant_params.expect_stripped_words = False
|
||||
return LLMContextAggregatorPair(
|
||||
context, user_params=user_params, assistant_params=assistant_params
|
||||
)
|
||||
@@ -1,448 +0,0 @@
|
||||
#
|
||||
# Copyright (c) 2024–2025, Daily
|
||||
#
|
||||
# SPDX-License-Identifier: BSD 2-Clause License
|
||||
#
|
||||
|
||||
"""This module implements Ultravox speech-to-text with a locally-loaded model."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import time
|
||||
from typing import AsyncGenerator, List, Optional
|
||||
|
||||
import numpy as np
|
||||
from huggingface_hub import login
|
||||
from loguru import logger
|
||||
|
||||
from pipecat.frames.frames import (
|
||||
AudioRawFrame,
|
||||
CancelFrame,
|
||||
EndFrame,
|
||||
ErrorFrame,
|
||||
Frame,
|
||||
LLMFullResponseEndFrame,
|
||||
LLMFullResponseStartFrame,
|
||||
LLMTextFrame,
|
||||
StartFrame,
|
||||
UserStartedSpeakingFrame,
|
||||
UserStoppedSpeakingFrame,
|
||||
)
|
||||
from pipecat.processors.frame_processor import FrameDirection
|
||||
from pipecat.services.ai_service import AIService
|
||||
|
||||
try:
|
||||
from transformers import AutoTokenizer
|
||||
from vllm import AsyncLLMEngine, SamplingParams
|
||||
from vllm.engine.arg_utils import AsyncEngineArgs
|
||||
except ModuleNotFoundError as e:
|
||||
logger.error(f"Exception: {e}")
|
||||
logger.error("In order to use Ultravox, you need to `pip install pipecat-ai[ultravox]`.")
|
||||
raise Exception(f"Missing module: {e}")
|
||||
|
||||
|
||||
class AudioBuffer:
|
||||
"""Buffer to collect audio frames before processing.
|
||||
|
||||
Manages the collection and state of audio frames during speech
|
||||
recording sessions, including timing and processing flags.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the audio buffer."""
|
||||
self.frames: List[AudioRawFrame] = []
|
||||
self.started_at: Optional[float] = None
|
||||
self.is_processing: bool = False
|
||||
|
||||
|
||||
class UltravoxModel:
|
||||
"""Model wrapper for the Ultravox multimodal model.
|
||||
|
||||
This class handles loading and running the Ultravox model for speech-to-text
|
||||
transcription using vLLM for efficient inference.
|
||||
"""
|
||||
|
||||
def __init__(self, model_name: str = "fixie-ai/ultravox-v0_5-llama-3_1-8b"):
|
||||
"""Initialize the Ultravox model.
|
||||
|
||||
Args:
|
||||
model_name: The name or path of the Ultravox model to load.
|
||||
Defaults to "fixie-ai/ultravox-v0_5-llama-3_1-8b".
|
||||
"""
|
||||
self.model_name = model_name
|
||||
self._initialize_engine()
|
||||
self._initialize_tokenizer()
|
||||
self.stop_token_ids = None
|
||||
|
||||
def _initialize_engine(self):
|
||||
"""Initialize the vLLM engine for inference."""
|
||||
engine_args = AsyncEngineArgs(
|
||||
model=self.model_name,
|
||||
gpu_memory_utilization=0.9,
|
||||
max_model_len=8192,
|
||||
trust_remote_code=True,
|
||||
)
|
||||
self.engine = AsyncLLMEngine.from_engine_args(engine_args)
|
||||
|
||||
def _initialize_tokenizer(self):
|
||||
"""Initialize the tokenizer for the model."""
|
||||
self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
|
||||
|
||||
def format_prompt(self, messages: list):
|
||||
"""Format chat messages into a prompt for the model.
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries with 'role' and 'content'.
|
||||
|
||||
Returns:
|
||||
str: Formatted prompt string ready for model input.
|
||||
"""
|
||||
return self.tokenizer.apply_chat_template(
|
||||
messages, tokenize=False, add_generation_prompt=True
|
||||
)
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
messages: list,
|
||||
temperature: float = 0.7,
|
||||
max_tokens: int = 100,
|
||||
audio: np.ndarray = None,
|
||||
):
|
||||
"""Generate text from audio input using the model.
|
||||
|
||||
Args:
|
||||
messages: List of message dictionaries for conversation context.
|
||||
temperature: Sampling temperature for generation randomness.
|
||||
max_tokens: Maximum number of tokens to generate.
|
||||
audio: Audio data as numpy array in float32 format.
|
||||
|
||||
Yields:
|
||||
str: JSON chunks of the generated response in OpenAI format.
|
||||
"""
|
||||
sampling_params = SamplingParams(
|
||||
temperature=temperature, max_tokens=max_tokens, stop_token_ids=self.stop_token_ids
|
||||
)
|
||||
|
||||
mm_data = {"audio": audio}
|
||||
inputs = {"prompt": self.format_prompt(messages), "multi_modal_data": mm_data}
|
||||
results_generator = self.engine.generate(inputs, sampling_params, str(time.time()))
|
||||
|
||||
previous_text = ""
|
||||
first_chunk = True
|
||||
|
||||
async for output in results_generator:
|
||||
prompt_output = output.outputs
|
||||
new_text = prompt_output[0].text[len(previous_text) :]
|
||||
previous_text = prompt_output[0].text
|
||||
|
||||
# Construct OpenAI-compatible chunk
|
||||
chunk = {
|
||||
"id": str(int(time.time() * 1000)),
|
||||
"object": "chat.completion.chunk",
|
||||
"created": int(time.time()),
|
||||
"model": self.model_name,
|
||||
"choices": [
|
||||
{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": None,
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
# Include the role in the first chunk
|
||||
if first_chunk:
|
||||
chunk["choices"][0]["delta"]["role"] = "assistant"
|
||||
first_chunk = False
|
||||
|
||||
# Add new text to the delta if any
|
||||
if new_text:
|
||||
chunk["choices"][0]["delta"]["content"] = new_text
|
||||
|
||||
# Capture a finish reason if it's provided
|
||||
finish_reason = prompt_output[0].finish_reason or None
|
||||
if finish_reason and finish_reason != "none":
|
||||
chunk["choices"][0]["finish_reason"] = finish_reason
|
||||
|
||||
yield json.dumps(chunk)
|
||||
|
||||
|
||||
class UltravoxSTTService(AIService):
|
||||
"""Service to transcribe audio using the Ultravox multimodal model.
|
||||
|
||||
This service collects audio frames during speech and processes them with
|
||||
Ultravox to generate text transcriptions. It handles real-time audio
|
||||
buffering, model warm-up, and streaming text generation.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
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,
|
||||
**kwargs,
|
||||
):
|
||||
"""Initialize the UltravoxSTTService.
|
||||
|
||||
Args:
|
||||
model_name: The Ultravox model to use. Defaults to
|
||||
"fixie-ai/ultravox-v0_5-llama-3_1-8b".
|
||||
hf_token: Hugging Face token for model access. If None, will try
|
||||
to use HF_TOKEN environment variable.
|
||||
temperature: Sampling temperature for generation. Defaults to 0.7.
|
||||
max_tokens: Maximum tokens to generate. Defaults to 100.
|
||||
**kwargs: Additional arguments passed to AIService.
|
||||
"""
|
||||
super().__init__(**kwargs)
|
||||
|
||||
# Authenticate with Hugging Face if token provided
|
||||
if hf_token:
|
||||
login(token=hf_token)
|
||||
elif os.environ.get("HF_TOKEN"):
|
||||
login(token=os.environ.get("HF_TOKEN"))
|
||||
else:
|
||||
logger.warning("No Hugging Face token provided. Model may not load correctly.")
|
||||
|
||||
# Initialize model
|
||||
self._model = UltravoxModel(model_name=model_name)
|
||||
|
||||
# Initialize service state
|
||||
self._buffer = AudioBuffer()
|
||||
self._temperature = temperature
|
||||
self._max_tokens = max_tokens
|
||||
self._connection_active = False
|
||||
self._warm_up_duration_sec = 1
|
||||
|
||||
logger.info(f"Initialized UltravoxSTTService with model: {model_name}")
|
||||
|
||||
async def _warm_up_model(self):
|
||||
"""Warm up the model with silent audio to improve first inference performance.
|
||||
|
||||
This method generates a short segment of silent audio and runs it through
|
||||
the model to ensure the model is fully loaded and optimized for the first
|
||||
real inference request.
|
||||
"""
|
||||
logger.info("Warming up Ultravox model with silent audio...")
|
||||
|
||||
# Generate silent audio at 16kHz sample rate
|
||||
sample_rate = 16000
|
||||
silent_audio = self._generate_silent_audio(sample_rate, self._warm_up_duration_sec)
|
||||
|
||||
try:
|
||||
# Process the silent audio with the model
|
||||
messages = [{"role": "user", "content": "<|audio|>\n"}]
|
||||
warmup_generator = self._model.generate(
|
||||
messages=messages,
|
||||
temperature=self._temperature,
|
||||
max_tokens=self._max_tokens,
|
||||
audio=silent_audio,
|
||||
)
|
||||
|
||||
# Consume the generator to actually run the inference
|
||||
async for _ in warmup_generator:
|
||||
pass
|
||||
|
||||
logger.info("Model warm-up completed successfully")
|
||||
except Exception as e:
|
||||
await self.push_error(error_msg=f"Unknown error occurred: {e}", exception=e)
|
||||
|
||||
def _generate_silent_audio(self, sample_rate=16000, duration_sec=1.0):
|
||||
"""Generate silent audio as a numpy array.
|
||||
|
||||
Args:
|
||||
sample_rate: Sample rate in Hz
|
||||
duration_sec: Duration of silence in seconds
|
||||
|
||||
Returns:
|
||||
np.ndarray: Float32 array of zeros representing silent audio
|
||||
"""
|
||||
# Calculate number of samples
|
||||
num_samples = int(sample_rate * duration_sec)
|
||||
|
||||
# Create silent audio as float32 in the [-1.0, 1.0] range
|
||||
silent_audio = np.zeros(num_samples, dtype=np.float32)
|
||||
|
||||
logger.info(f"Generated {duration_sec}s of silent audio ({num_samples} samples)")
|
||||
return silent_audio
|
||||
|
||||
def can_generate_metrics(self) -> bool:
|
||||
"""Indicates whether this service can generate metrics.
|
||||
|
||||
Returns:
|
||||
bool: True, as this service supports metric generation.
|
||||
"""
|
||||
return True
|
||||
|
||||
async def start(self, frame: StartFrame):
|
||||
"""Handle service start.
|
||||
|
||||
Starts the service, marks it as active, and performs model warm-up
|
||||
to ensure optimal performance for the first inference.
|
||||
|
||||
Args:
|
||||
frame: StartFrame that triggered this method.
|
||||
"""
|
||||
await super().start(frame)
|
||||
self._connection_active = True
|
||||
|
||||
await self._warm_up_model()
|
||||
|
||||
logger.info("UltravoxSTTService started")
|
||||
|
||||
async def stop(self, frame: EndFrame):
|
||||
"""Handle service stop.
|
||||
|
||||
Stops the service and marks it as inactive.
|
||||
|
||||
Args:
|
||||
frame: EndFrame that triggered this method.
|
||||
"""
|
||||
await super().stop(frame)
|
||||
self._connection_active = False
|
||||
logger.info("UltravoxSTTService stopped")
|
||||
|
||||
async def cancel(self, frame: CancelFrame):
|
||||
"""Handle service cancellation.
|
||||
|
||||
Cancels the service, clears any buffered audio, and marks it as inactive.
|
||||
|
||||
Args:
|
||||
frame: CancelFrame that triggered this method.
|
||||
"""
|
||||
await super().cancel(frame)
|
||||
self._connection_active = False
|
||||
self._buffer = AudioBuffer()
|
||||
logger.info("UltravoxSTTService cancelled")
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
"""Process incoming frames.
|
||||
|
||||
This method collects audio frames during speech and processes them
|
||||
when speech ends to generate text transcriptions.
|
||||
|
||||
Args:
|
||||
frame: The frame to process.
|
||||
direction: Direction of the frame (input/output).
|
||||
"""
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
if isinstance(frame, UserStartedSpeakingFrame):
|
||||
logger.info("Speech started")
|
||||
self._buffer = AudioBuffer()
|
||||
self._buffer.started_at = time.time()
|
||||
|
||||
elif isinstance(frame, AudioRawFrame) and self._buffer.started_at is not None:
|
||||
self._buffer.frames.append(frame)
|
||||
|
||||
elif isinstance(frame, UserStoppedSpeakingFrame):
|
||||
if self._buffer.frames and not self._buffer.is_processing:
|
||||
logger.info("Speech ended, processing buffer...")
|
||||
await self.process_generator(self._process_audio_buffer())
|
||||
return # Return early to avoid pushing None frame
|
||||
|
||||
# Only push the original frame if we haven't processed audio
|
||||
if frame is not None:
|
||||
await self.push_frame(frame, direction)
|
||||
|
||||
async def _process_audio_buffer(self) -> AsyncGenerator[Frame, None]:
|
||||
"""Process collected audio frames with Ultravox.
|
||||
|
||||
This method concatenates audio frames, processes them with the model,
|
||||
and yields the resulting text frames.
|
||||
|
||||
Yields:
|
||||
Frame: TextFrame containing the transcribed text
|
||||
"""
|
||||
try:
|
||||
self._buffer.is_processing = True
|
||||
|
||||
# Check if we have valid frames before processing
|
||||
if not self._buffer.frames:
|
||||
logger.warning("No audio frames to process")
|
||||
yield ErrorFrame("No audio frames to process")
|
||||
return
|
||||
|
||||
# Process audio frames
|
||||
audio_arrays = []
|
||||
for f in self._buffer.frames:
|
||||
if hasattr(f, "audio") and f.audio:
|
||||
# Handle bytes data - these are int16 PCM samples
|
||||
if isinstance(f.audio, bytes):
|
||||
try:
|
||||
# Convert bytes to int16 array
|
||||
arr = np.frombuffer(f.audio, dtype=np.int16)
|
||||
if arr.size > 0: # Check if array is not empty
|
||||
audio_arrays.append(arr)
|
||||
except Exception as e:
|
||||
yield ErrorFrame(error=f"Unknown error occurred: {e}")
|
||||
# Handle numpy array data
|
||||
elif isinstance(f.audio, np.ndarray):
|
||||
if f.audio.size > 0: # Check if array is not empty
|
||||
# Ensure it's int16 data
|
||||
if f.audio.dtype != np.int16:
|
||||
logger.info(f"Converting array from {f.audio.dtype} to int16")
|
||||
audio_arrays.append(f.audio.astype(np.int16))
|
||||
else:
|
||||
audio_arrays.append(f.audio)
|
||||
|
||||
# Only proceed if we have valid audio arrays
|
||||
if not audio_arrays:
|
||||
logger.warning("No valid audio data found in frames")
|
||||
yield ErrorFrame("No valid audio data found in frames")
|
||||
return
|
||||
|
||||
# Concatenate audio frames - all should be int16 now
|
||||
audio_data = np.concatenate(audio_arrays)
|
||||
|
||||
audio_int16 = audio_data # Already in int16 format
|
||||
# Save int16 audio
|
||||
|
||||
# Convert int16 to float32 and normalize for model input
|
||||
audio_float32 = audio_int16.astype(np.float32) / 32768.0
|
||||
|
||||
# Generate text using the model
|
||||
if self._model:
|
||||
try:
|
||||
logger.info("Generating text from audio using model...")
|
||||
|
||||
# Start metrics tracking
|
||||
await self.start_ttfb_metrics()
|
||||
await self.start_processing_metrics()
|
||||
|
||||
yield LLMFullResponseStartFrame()
|
||||
|
||||
async for response in self._model.generate(
|
||||
messages=[{"role": "user", "content": "<|audio|>\n"}],
|
||||
temperature=self._temperature,
|
||||
max_tokens=self._max_tokens,
|
||||
audio=audio_float32,
|
||||
):
|
||||
# Stop TTFB metrics after first response
|
||||
await self.stop_ttfb_metrics()
|
||||
|
||||
chunk = json.loads(response)
|
||||
if "choices" in chunk and len(chunk["choices"]) > 0:
|
||||
delta = chunk["choices"][0]["delta"]
|
||||
if "content" in delta:
|
||||
new_text = delta["content"]
|
||||
if new_text:
|
||||
yield LLMTextFrame(text=new_text)
|
||||
|
||||
# Stop processing metrics after completion
|
||||
await self.stop_processing_metrics()
|
||||
|
||||
yield LLMFullResponseEndFrame()
|
||||
|
||||
except Exception as e:
|
||||
yield ErrorFrame(error=f"Unknown error occurred: {e}")
|
||||
else:
|
||||
yield ErrorFrame("No model available for text generation")
|
||||
|
||||
except Exception as e:
|
||||
yield ErrorFrame(f"Error processing audio: {str(e)}")
|
||||
finally:
|
||||
self._buffer.is_processing = False
|
||||
self._buffer.frames = []
|
||||
self._buffer.started_at = None
|
||||
Reference in New Issue
Block a user