from __future__ import annotations import asyncio import base64 import json import time from collections.abc import AsyncIterator from dataclasses import dataclass from typing import Any, cast from openai import AsyncOpenAI from ... import _debug from ...exceptions import AgentsException from ...logger import logger from ...tracing import Span, SpanError, TranscriptionSpanData, transcription_span from ..exceptions import STTWebsocketConnectionError from ..imports import np, npt, websockets from ..input import AudioInput, StreamedAudioInput from ..model import StreamedTranscriptionSession, STTModel, STTModelSettings EVENT_INACTIVITY_TIMEOUT = 1000 # Timeout for inactivity in event processing SESSION_CREATION_TIMEOUT = 10 # Timeout waiting for session.created event SESSION_UPDATE_TIMEOUT = 10 # Timeout waiting for session.updated event DEFAULT_TURN_DETECTION = {"type": "semantic_vad"} @dataclass class ErrorSentinel: error: Exception class SessionCompleteSentinel: pass class WebsocketDoneSentinel: pass def _audio_to_base64(audio_data: list[npt.NDArray[np.int16 | np.float32]]) -> str: concatenated_audio = np.concatenate(audio_data) if concatenated_audio.dtype == np.float32: # convert to int16 concatenated_audio = np.clip(concatenated_audio, -1.0, 1.0) concatenated_audio = (concatenated_audio * 32767).astype(np.int16) audio_bytes = concatenated_audio.tobytes() return base64.b64encode(audio_bytes).decode("utf-8") async def _wait_for_event( event_queue: asyncio.Queue[dict[str, Any]], expected_types: list[str], timeout: float ): """ Wait for an event from event_queue whose type is in expected_types within the specified timeout. """ start_time = time.time() while True: remaining = timeout - (time.time() - start_time) if remaining <= 0: raise TimeoutError(f"Timeout waiting for event(s): {expected_types}") evt = await asyncio.wait_for(event_queue.get(), timeout=remaining) evt_type = evt.get("type", "") if evt_type in expected_types: return evt elif evt_type == "error": raise Exception(f"Error event: {evt.get('error')}") class OpenAISTTTranscriptionSession(StreamedTranscriptionSession): """A transcription session for OpenAI's STT model.""" def __init__( self, input: StreamedAudioInput, client: AsyncOpenAI, model: str, settings: STTModelSettings, trace_include_sensitive_data: bool, trace_include_sensitive_audio_data: bool, ): self.connected: bool = False self._client = client self._model = model self._settings = settings self._turn_detection = settings.turn_detection or DEFAULT_TURN_DETECTION self._trace_include_sensitive_data = trace_include_sensitive_data self._trace_include_sensitive_audio_data = trace_include_sensitive_audio_data self._input_queue: asyncio.Queue[npt.NDArray[np.int16 | np.float32]] = input.queue self._output_queue: asyncio.Queue[str | ErrorSentinel | SessionCompleteSentinel] = ( asyncio.Queue() ) self._websocket: websockets.ClientConnection | None = None self._event_queue: asyncio.Queue[dict[str, Any] | WebsocketDoneSentinel] = asyncio.Queue() self._state_queue: asyncio.Queue[dict[str, Any]] = asyncio.Queue() self._turn_audio_buffer: list[npt.NDArray[np.int16 | np.float32]] = [] self._tracing_span: Span[TranscriptionSpanData] | None = None # tasks self._listener_task: asyncio.Task[Any] | None = None self._process_events_task: asyncio.Task[Any] | None = None self._stream_audio_task: asyncio.Task[Any] | None = None self._connection_task: asyncio.Task[Any] | None = None self._stored_exception: Exception | None = None def _start_turn(self) -> None: self._tracing_span = transcription_span( model=self._model, model_config={ "temperature": self._settings.temperature, "language": self._settings.language, "prompt": self._settings.prompt, "turn_detection": self._turn_detection, }, ) self._tracing_span.start() def _end_turn(self, _transcript: str) -> None: if len(_transcript) < 1: return if self._tracing_span: if self._trace_include_sensitive_audio_data: self._tracing_span.span_data.input = _audio_to_base64(self._turn_audio_buffer) self._tracing_span.span_data.input_format = "pcm" if self._trace_include_sensitive_data: self._tracing_span.span_data.output = _transcript self._tracing_span.finish() self._turn_audio_buffer = [] self._tracing_span = None async def _event_listener(self) -> None: assert self._websocket is not None, "Websocket not initialized" async for message in self._websocket: try: event = json.loads(message) if event.get("type") == "error": raise STTWebsocketConnectionError(f"Error event: {event.get('error')}") if event.get("type") in [ "session.updated", "transcription_session.updated", "session.created", "transcription_session.created", ]: await self._state_queue.put(event) await self._event_queue.put(event) except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise STTWebsocketConnectionError("Error parsing events") from e await self._event_queue.put(WebsocketDoneSentinel()) async def _configure_session(self) -> None: assert self._websocket is not None, "Websocket not initialized" await self._websocket.send( json.dumps( { "type": "transcription_session.update", "session": { "input_audio_format": "pcm16", "input_audio_transcription": {"model": self._model}, "turn_detection": self._turn_detection, }, } ) ) async def _setup_connection(self, ws: websockets.ClientConnection) -> None: self._websocket = ws self._listener_task = asyncio.create_task(self._event_listener()) try: event = await _wait_for_event( self._state_queue, ["session.created", "transcription_session.created"], SESSION_CREATION_TIMEOUT, ) except TimeoutError as e: wrapped_err = STTWebsocketConnectionError( "Timeout waiting for transcription_session.created event" ) await self._output_queue.put(ErrorSentinel(wrapped_err)) raise wrapped_err from e except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise e await self._configure_session() try: event = await _wait_for_event( self._state_queue, ["session.updated", "transcription_session.updated"], SESSION_UPDATE_TIMEOUT, ) if _debug.DONT_LOG_MODEL_DATA: logger.debug("Session updated") else: logger.debug(f"Session updated: {event}") except TimeoutError as e: wrapped_err = STTWebsocketConnectionError( "Timeout waiting for transcription_session.updated event" ) await self._output_queue.put(ErrorSentinel(wrapped_err)) raise wrapped_err from e except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise async def _handle_events(self) -> None: while True: try: event = await asyncio.wait_for( self._event_queue.get(), timeout=EVENT_INACTIVITY_TIMEOUT ) if isinstance(event, WebsocketDoneSentinel): # processed all events and websocket is done break event_type = event.get("type", "unknown") if event_type == "conversation.item.input_audio_transcription.completed": transcript = cast(str, event.get("transcript", "")) if len(transcript) > 0: self._end_turn(transcript) self._start_turn() await self._output_queue.put(transcript) await asyncio.sleep(0) # yield control except asyncio.TimeoutError: # No new events for a while. Assume the session is done. break except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise e await self._output_queue.put(SessionCompleteSentinel()) async def _stream_audio( self, audio_queue: asyncio.Queue[npt.NDArray[np.int16 | np.float32]] ) -> None: assert self._websocket is not None, "Websocket not initialized" self._start_turn() while True: buffer = await audio_queue.get() if buffer is None: break self._turn_audio_buffer.append(buffer) try: await self._websocket.send( json.dumps( { "type": "input_audio_buffer.append", "audio": base64.b64encode(buffer.tobytes()).decode("utf-8"), } ) ) except websockets.ConnectionClosed: break except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise e await asyncio.sleep(0) # yield control async def _process_websocket_connection(self) -> None: try: async with websockets.connect( "wss://api.openai.com/v1/realtime?intent=transcription", additional_headers={ "Authorization": f"Bearer {self._client.api_key}", "OpenAI-Beta": "realtime=v1", "OpenAI-Log-Session": "1", }, ) as ws: await self._setup_connection(ws) self._process_events_task = asyncio.create_task(self._handle_events()) self._stream_audio_task = asyncio.create_task(self._stream_audio(self._input_queue)) self.connected = True if self._listener_task: await self._listener_task else: logger.error("Listener task not initialized") raise AgentsException("Listener task not initialized") except Exception as e: await self._output_queue.put(ErrorSentinel(e)) raise e def _check_errors(self) -> None: if self._connection_task and self._connection_task.done(): exc = self._connection_task.exception() if exc and isinstance(exc, Exception): self._stored_exception = exc if self._process_events_task and self._process_events_task.done(): exc = self._process_events_task.exception() if exc and isinstance(exc, Exception): self._stored_exception = exc if self._stream_audio_task and self._stream_audio_task.done(): exc = self._stream_audio_task.exception() if exc and isinstance(exc, Exception): self._stored_exception = exc if self._listener_task and self._listener_task.done(): exc = self._listener_task.exception() if exc and isinstance(exc, Exception): self._stored_exception = exc def _cleanup_tasks(self) -> None: if self._listener_task and not self._listener_task.done(): self._listener_task.cancel() if self._process_events_task and not self._process_events_task.done(): self._process_events_task.cancel() if self._stream_audio_task and not self._stream_audio_task.done(): self._stream_audio_task.cancel() if self._connection_task and not self._connection_task.done(): self._connection_task.cancel() async def transcribe_turns(self) -> AsyncIterator[str]: self._connection_task = asyncio.create_task(self._process_websocket_connection()) while True: try: turn = await self._output_queue.get() except asyncio.CancelledError: break if ( turn is None or isinstance(turn, ErrorSentinel) or isinstance(turn, SessionCompleteSentinel) ): self._output_queue.task_done() break yield turn self._output_queue.task_done() if self._tracing_span: self._end_turn("") if self._websocket: await self._websocket.close() self._check_errors() if self._stored_exception: raise self._stored_exception async def close(self) -> None: if self._websocket: await self._websocket.close() self._cleanup_tasks() class OpenAISTTModel(STTModel): """A speech-to-text model for OpenAI.""" def __init__( self, model: str, openai_client: AsyncOpenAI, ): """Create a new OpenAI speech-to-text model. Args: model: The name of the model to use. openai_client: The OpenAI client to use. """ self.model = model self._client = openai_client @property def model_name(self) -> str: return self.model def _non_null_or_not_given(self, value: Any) -> Any: return value if value is not None else None # NOT_GIVEN async def transcribe( self, input: AudioInput, settings: STTModelSettings, trace_include_sensitive_data: bool, trace_include_sensitive_audio_data: bool, ) -> str: """Transcribe an audio input. Args: input: The audio input to transcribe. settings: The settings to use for the transcription. Returns: The transcribed text. """ with transcription_span( model=self.model, input=input.to_base64() if trace_include_sensitive_audio_data else "", input_format="pcm", model_config={ "temperature": self._non_null_or_not_given(settings.temperature), "language": self._non_null_or_not_given(settings.language), "prompt": self._non_null_or_not_given(settings.prompt), }, ) as span: try: response = await self._client.audio.transcriptions.create( model=self.model, file=input.to_audio_file(), prompt=self._non_null_or_not_given(settings.prompt), language=self._non_null_or_not_given(settings.language), temperature=self._non_null_or_not_given(settings.temperature), ) if trace_include_sensitive_data: span.span_data.output = response.text return response.text except Exception as e: span.span_data.output = "" span.set_error(SpanError(message=str(e), data={})) raise e async def create_session( self, input: StreamedAudioInput, settings: STTModelSettings, trace_include_sensitive_data: bool, trace_include_sensitive_audio_data: bool, ) -> StreamedTranscriptionSession: """Create a new transcription session. Args: input: The audio input to transcribe. settings: The settings to use for the transcription. trace_include_sensitive_data: Whether to include sensitive data in traces. trace_include_sensitive_audio_data: Whether to include sensitive audio data in traces. Returns: A new transcription session. """ return OpenAISTTTranscriptionSession( input, self._client, self.model, settings, trace_include_sensitive_data, trace_include_sensitive_audio_data, )