cartesia 1.4.0__py3-none-any.whl → 2.0.0a2__py3-none-any.whl

cartesia/{_async_websocket.py → tts/_async_websocket.py}

@@ -1,4 +1,6 @@
 import asyncio
+import json
+import typing
 import uuid
 from collections import defaultdict
 from types import TracebackType
@@ -6,11 +8,26 @@ from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Union
 
 import aiohttp
 
-from cartesia._constants import DEFAULT_MODEL_ID, DEFAULT_OUTPUT_FORMAT, DEFAULT_VOICE_EMBEDDING
-from cartesia._types import OutputFormat, VoiceControls
-from cartesia._websocket import _WebSocket
-from cartesia.tts import TTS
-from cartesia.utils.tts import _construct_tts_request
+from cartesia.tts.requests import TtsRequestVoiceSpecifierParams
+from cartesia.tts.requests.output_format import OutputFormatParams
+from cartesia.tts.types import (
+    WebSocketResponse,
+    WebSocketResponse_Done,
+    WebSocketResponse_Error,
+    WebSocketResponse_FlushDone,
+    WebSocketTtsOutput,
+    WordTimestamps,
+)
+
+from ..core.pydantic_utilities import parse_obj_as
+from ._websocket import TtsWebsocket
+from .types.generation_request import GenerationRequest
+from .utils.constants import (
+    DEFAULT_MODEL_ID,
+    DEFAULT_OUTPUT_FORMAT,
+    DEFAULT_VOICE_EMBEDDING,
+)
+from .utils.tts import get_output_format
 
 
 class _AsyncTTSContext:
@@ -26,7 +43,9 @@ class _AsyncTTSContext:
 
     """
 
-    def __init__(self, context_id: str, websocket: "_AsyncWebSocket", timeout: float):
+    def __init__(
+        self, context_id: str, websocket: "AsyncTtsWebsocket", timeout: float = 30
+    ):
         self._context_id = context_id
         self._websocket = websocket
         self.timeout = timeout
@@ -38,60 +57,73 @@ class _AsyncTTSContext:
 
     async def send(
         self,
+        *,
         model_id: str,
         transcript: str,
-        output_format: OutputFormat,
-        voice_id: Optional[str] = None,
-        voice_embedding: Optional[List[float]] = None,
+        output_format: OutputFormatParams,
+        voice: TtsRequestVoiceSpecifierParams,
         context_id: Optional[str] = None,
-        continue_: bool = False,
-        flush: bool = False,
         duration: Optional[int] = None,
         language: Optional[str] = None,
+        stream: bool = True,
         add_timestamps: bool = False,
-        _experimental_voice_controls: Optional[VoiceControls] = None,
+        continue_: bool = False,
+        flush: bool = False,
     ) -> None:
         """Send audio generation requests to the WebSocket. The response can be received using the `receive` method.
 
         Args:
-            model_id: The ID of the model to use for generating audio.
-            transcript: The text to convert to speech.
-            output_format: A dictionary containing the details of the output format.
-            voice_id: The ID of the voice to use for generating audio.
-            voice_embedding: The embedding of the voice to use for generating audio.
-            context_id: The context ID to use for the request. If not specified, a random context ID will be generated.
-            continue_: Whether to continue the audio generation from the previous transcript or not.
-            flush: Whether to trigger a manual flush for the current context's generation.
-            duration: The duration of the audio in seconds.
-            language: The language code for the audio request. This can only be used with `model_id = sonic-multilingual`.
-            add_timestamps: Whether to return word-level timestamps.
-            _experimental_voice_controls: Experimental voice controls for controlling speed and emotion.
-                Note: This is an experimental feature and may change rapidly in future releases.
+            request: The request to generate audio.
 
         Returns:
             None.
         """
-        if context_id is not None and context_id != self._context_id:
-            raise ValueError("Context ID does not match the context ID of the current context.")
-        if continue_ and transcript == "" and not flush:
-            raise ValueError("Transcript cannot be empty when continue_ is True.")
-
         await self._websocket.connect()
-
-        request_body = _construct_tts_request(
-            model_id=model_id,
-            transcript=transcript,
-            output_format=output_format,
-            voice_id=voice_id,
-            voice_embedding=voice_embedding,
-            duration=duration,
-            language=language,
-            context_id=self._context_id,
-            add_timestamps=add_timestamps,
-            continue_=continue_,
-            flush=flush,
-            _experimental_voice_controls=_experimental_voice_controls,
-        )
+        assert self._websocket.websocket is not None, "WebSocket is not connected"
+
+        request_body = {
+            "model_id": model_id,
+            "transcript": transcript,
+            "output_format": (
+                output_format
+                if isinstance(output_format, dict)
+                else output_format.dict()
+            ),
+            "voice": (voice if isinstance(voice, dict) else voice.dict()),
+            "context_id": self._context_id,
+        }
+        if context_id is not None:
+            request_body["context_id"] = context_id
+        if duration is not None:
+            request_body["duration"] = duration
+        if language is not None:
+            request_body["language"] = language
+        if stream:
+            request_body["stream"] = stream
+        if add_timestamps:
+            request_body["add_timestamps"] = add_timestamps
+        if continue_:
+            request_body["continue"] = continue_
+        if flush:
+            request_body["flush"] = flush
+
+        if (
+            "context_id" in request_body
+            and request_body["context_id"] is not None
+            and request_body["context_id"] != self._context_id
+        ):
+            raise ValueError(
+                "Context ID does not match the context ID of the current context."
+            )
+        request_body["context_id"] = self._context_id
+
+        if (
+            "continue" in request_body
+            and request_body["continue"]
+            and request_body["transcript"] == ""
+            and ("flush" in request_body and not request_body["flush"])
+        ):
+            raise ValueError("Transcript cannot be empty when continue_ is True.")
 
         await self._websocket.websocket.send_json(request_body)
 
@@ -103,8 +135,11 @@ class _AsyncTTSContext:
         await self.send(
             model_id=DEFAULT_MODEL_ID,
             transcript="",
-            output_format=TTS.get_output_format(DEFAULT_OUTPUT_FORMAT),
-            voice_embedding=DEFAULT_VOICE_EMBEDDING,  # Default voice embedding since it's a required input for now.
+            output_format=get_output_format(DEFAULT_OUTPUT_FORMAT),
+            voice={
+                "mode": "embedding",
+                "embedding": DEFAULT_VOICE_EMBEDDING,
+            },
             context_id=self._context_id,
             continue_=False,
         )
@@ -114,8 +149,11 @@ class _AsyncTTSContext:
         await self.send(
             model_id=DEFAULT_MODEL_ID,
             transcript="",
-            output_format=TTS.get_output_format(DEFAULT_OUTPUT_FORMAT),
-            voice_embedding=DEFAULT_VOICE_EMBEDDING,  # Default voice embedding since it's a required input for now.
+            output_format=get_output_format(DEFAULT_OUTPUT_FORMAT),
+            voice={
+                "mode": "embedding",
+                "embedding": DEFAULT_VOICE_EMBEDDING,
+            },
             context_id=self._context_id,
             continue_=True,
             flush=True,
@@ -134,11 +172,23 @@ class _AsyncTTSContext:
                     response = await self._websocket._get_message(
                         self._context_id, timeout=self.timeout, flush_id=flush_id
                     )
-                    if "error" in response:
-                        raise RuntimeError(f"Error generating audio:\n{response['error']}")
-                    if response.get("flush_done") or response["done"]:
+                    response_obj = typing.cast(
+                        WebSocketResponse,
+                        parse_obj_as(
+                            type_=WebSocketResponse, object_=response  # type: ignore
+                        ),
+                    )
+                    if isinstance(response_obj, WebSocketResponse_Error):
+                        raise RuntimeError(
+                            f"Error generating audio:\n{response_obj.error}"
+                        )
+                    if isinstance(response_obj, WebSocketResponse_Done) or isinstance(
+                        response_obj, WebSocketResponse_FlushDone
+                    ):
                         break
-                    yield self._websocket._convert_response(response, include_context_id=True)
+                    yield self._websocket._convert_response(
+                        response_obj, include_context_id=True
+                    )
             except Exception as e:
                 if isinstance(e, asyncio.TimeoutError):
                     raise RuntimeError("Timeout while waiting for audio chunk")
@@ -146,7 +196,7 @@ class _AsyncTTSContext:
 
         return generator
 
-    async def receive(self) -> AsyncGenerator[Dict[str, Any], None]:
+    async def receive(self) -> AsyncGenerator[WebSocketTtsOutput, None]:
         """Receive the audio chunks from the WebSocket. This method is a generator that yields audio chunks.
 
         Returns:
@@ -157,11 +207,21 @@ class _AsyncTTSContext:
                 response = await self._websocket._get_message(
                     self._context_id, timeout=self.timeout
                 )
-                if "error" in response:
-                    raise RuntimeError(f"Error generating audio:\n{response['error']}")
-                if response["done"]:
+                response_obj = typing.cast(
+                    WebSocketResponse,
+                    parse_obj_as(
+                        type_=WebSocketResponse,  # type: ignore
+                        object_=response,
+                    ),
+                )
+
+                if isinstance(response_obj, WebSocketResponse_Error):
+                    raise RuntimeError(f"Error generating audio:\n{response_obj.error}")
+                if isinstance(response_obj, WebSocketResponse_Done):
                     break
-                yield self._websocket._convert_response(response, include_context_id=True)
+                yield self._websocket._convert_response(
+                    response_obj, include_context_id=True
+                )
         except Exception as e:
             if isinstance(e, asyncio.TimeoutError):
                 raise RuntimeError("Timeout while waiting for audio chunk")
@@ -192,7 +252,7 @@ class _AsyncTTSContext:
         self._close()
 
 
-class _AsyncWebSocket(_WebSocket):
+class AsyncTtsWebsocket(TtsWebsocket):
     """This class contains methods to generate audio using WebSocket asynchronously."""
 
     def __init__(
@@ -200,8 +260,8 @@ class _AsyncWebSocket(_WebSocket):
         ws_url: str,
         api_key: str,
         cartesia_version: str,
-        timeout: float,
         get_session: Callable[[], Optional[aiohttp.ClientSession]],
+        timeout: float = 30,
     ):
         """
         Args:
@@ -216,7 +276,7 @@ class _AsyncWebSocket(_WebSocket):
         self._get_session = get_session
         self.websocket = None
         self._context_queues: Dict[str, List[asyncio.Queue]] = {}
-        self._processing_task: asyncio.Task = None
+        self._processing_task: Optional[asyncio.Task] = None
 
     def __del__(self):
         try:
@@ -268,18 +328,17 @@ class _AsyncWebSocket(_WebSocket):
 
     async def send(
         self,
+        *,
         model_id: str,
         transcript: str,
-        output_format: OutputFormat,
-        voice_id: Optional[str] = None,
-        voice_embedding: Optional[List[float]] = None,
+        output_format: OutputFormatParams,
+        voice: TtsRequestVoiceSpecifierParams,
         context_id: Optional[str] = None,
         duration: Optional[int] = None,
         language: Optional[str] = None,
         stream: bool = True,
         add_timestamps: bool = False,
-        _experimental_voice_controls: Optional[VoiceControls] = None,
-    ) -> Union[bytes, AsyncGenerator[bytes, None]]:
+    ):
         """See :meth:`_WebSocket.send` for details."""
         if context_id is None:
             context_id = str(uuid.uuid4())
@@ -290,14 +349,12 @@ class _AsyncWebSocket(_WebSocket):
             model_id=model_id,
             transcript=transcript,
             output_format=output_format,
-            voice_id=voice_id,
-            voice_embedding=voice_embedding,
+            voice=voice,
             context_id=context_id,
             duration=duration,
             language=language,
             continue_=False,
             add_timestamps=add_timestamps,
-            _experimental_voice_controls=_experimental_voice_controls,
         )
 
         generator = ctx.receive()
@@ -305,18 +362,32 @@ class _AsyncWebSocket(_WebSocket):
         if stream:
             return generator
 
-        chunks = []
-        word_timestamps = defaultdict(list)
+        chunks: typing.List[str] = []
+        words: typing.List[str] = []
+        start: typing.List[float] = []
+        end: typing.List[float] = []
         async for chunk in generator:
-            if "audio" in chunk:
-                chunks.append(chunk["audio"])
-            if add_timestamps and "word_timestamps" in chunk:
-                for k, v in chunk["word_timestamps"].items():
-                    word_timestamps[k].extend(v)
-        out = {"audio": b"".join(chunks), "context_id": context_id}
-        if add_timestamps:
-            out["word_timestamps"] = word_timestamps
-        return out
+            if chunk.audio is not None:
+                chunks.append(chunk.audio)
+            if add_timestamps and chunk.word_timestamps is not None:
+                if chunk.word_timestamps is not None:
+                    words.extend(chunk.word_timestamps.words)
+                    start.extend(chunk.word_timestamps.start)
+                    end.extend(chunk.word_timestamps.end)
+
+        return WebSocketTtsOutput(
+            audio=b"".join(chunks),  # type: ignore
+            context_id=context_id,
+            word_timestamps=(
+                WordTimestamps(
+                    words=words,
+                    start=start,
+                    end=end,
+                )
+                if add_timestamps
+                else None
+            ),
+        )
 
     async def _process_responses(self):
         try:
@@ -332,12 +403,14 @@ class _AsyncWebSocket(_WebSocket):
             raise e
 
     async def _get_message(
-        self, context_id: str, timeout: float, flush_id: Optional[int] = -1
+        self, context_id: str, timeout: float, flush_id: int = -1
     ) -> Dict[str, Any]:
         if context_id not in self._context_queues:
             raise ValueError(f"Context ID {context_id} not found.")
         if len(self._context_queues[context_id]) <= flush_id:
-            raise ValueError(f"Flush ID {flush_id} not found for context ID {context_id}.")
+            raise ValueError(
+                f"Flush ID {flush_id} not found for context ID {context_id}."
+            )
         return await asyncio.wait_for(
             self._context_queues[context_id][flush_id].get(), timeout=timeout
         )
@@ -350,9 +423,11 @@ class _AsyncWebSocket(_WebSocket):
         if self._processing_task is None or self._processing_task.done():
             self._processing_task = asyncio.create_task(self._process_responses())
 
-    def context(self, context_id: Optional[str] = None) -> _AsyncTTSContext:
+    def context(self, context_id: Optional[str] = None):
         if context_id in self._context_queues:
-            raise ValueError(f"AsyncContext for context ID {context_id} already exists.")
+            raise ValueError(
+                f"AsyncContext for context ID {context_id} already exists."
+            )
         if context_id is None:
             context_id = str(uuid.uuid4())
         if context_id not in self._context_queues: