livekit-plugins-aws 1.0.0rc6__py3-none-any.whl → 1.3.9__py3-none-any.whl

livekit/plugins/aws/experimental/realtime/turn_tracker.py

@@ -0,0 +1,171 @@
+from __future__ import annotations
+
+import datetime
+import enum
+import uuid
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+from livekit.agents import llm
+
+from ...log import logger
+
+# Nova Sonic's barge-in detection signal (raw content without newline)
+BARGE_IN_CONTENT = '{ "interrupted" : true }'
+
+
+class _Phase(enum.Enum):
+    IDLE = 0  # waiting for the USER to begin speaking
+    USER_SPEAKING = 1  # still receiving USER text+audio blocks
+    USER_FINISHED = 2  # first ASSISTANT speculative block observed
+    ASSISTANT_RESPONDING = 3  # ASSISTANT audio/text streaming
+    DONE = 4  # assistant audio ended (END_TURN) or barge-in (INTERRUPTED)
+
+
+# note: b/c user ASR text is transcribed server-side, a single turn constitutes
+# both the user and agent's speech
+@dataclass
+class _Turn:
+    turn_id: int
+    input_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+    created: datetime.datetime = field(default_factory=datetime.datetime.utcnow)
+    transcript: list[str] = field(default_factory=list)
+
+    phase: _Phase = _Phase.IDLE
+    ev_input_started: bool = False
+    ev_input_stopped: bool = False
+    ev_trans_completed: bool = False
+    ev_generation_sent: bool = False
+
+    def add_partial_text(self, text: str) -> None:
+        self.transcript.append(text)
+
+    @property
+    def curr_transcript(self) -> str:
+        return " ".join(self.transcript)
+
+
+class _TurnTracker:
+    def __init__(
+        self,
+        emit_fn: Callable[[str, Any], None],
+        emit_generation_fn: Callable[[], None],
+    ):
+        self._emit = emit_fn
+        self._turn_idx = 0
+        self._curr_turn: _Turn | None = None
+        self._emit_generation_fn = emit_generation_fn
+
+    # --------------------------------------------------------
+    #  PUBLIC ENTRY POINT
+    # --------------------------------------------------------
+    def feed(self, event: dict) -> None:
+        turn = self._ensure_turn()
+        kind = _classify(event)
+
+        if kind == "USER_TEXT_PARTIAL":
+            turn.add_partial_text(event["event"]["textOutput"]["content"])
+            self._maybe_emit_input_started(turn)
+            self._emit_transcript_updated(turn)
+            # note: cannot invoke self._maybe_input_stopped() here
+            # b/c there is no way to know if the user is done speaking
+
+        # will always be correlated b/c generate_reply() is a stub
+        # user ASR text ends when agent's ASR speculative text begins
+        # corresponds to beginning of agent's turn
+        elif kind == "TOOL_OUTPUT_CONTENT_START" or kind == "ASSISTANT_SPEC_START":
+            # must be a maybe methods b/c agent can chain multiple tool calls
+            self._maybe_emit_input_stopped(turn)
+            self._maybe_emit_transcript_completed(turn)
+            self._maybe_emit_generation_created(turn)
+
+        elif kind == "BARGE_IN":
+            logger.debug(f"BARGE-IN DETECTED IN TURN TRACKER: {turn}")
+            # start new turn immediately to make interruptions snappier
+            self._emit("input_speech_started", llm.InputSpeechStartedEvent())
+            turn.phase = _Phase.DONE
+
+        elif kind == "ASSISTANT_AUDIO_END":
+            if event["event"]["contentEnd"]["stopReason"] == "END_TURN":
+                turn.phase = _Phase.DONE
+
+        if turn.phase is _Phase.DONE:
+            self._curr_turn = None
+
+    def _ensure_turn(self) -> _Turn:
+        if self._curr_turn is None:
+            self._turn_idx += 1
+            self._curr_turn = _Turn(turn_id=self._turn_idx)
+        return self._curr_turn
+
+    def _maybe_emit_input_started(self, turn: _Turn) -> None:
+        if not turn.ev_input_started:
+            turn.ev_input_started = True
+            self._emit("input_speech_started", llm.InputSpeechStartedEvent())
+            turn.phase = _Phase.USER_SPEAKING
+
+    def _maybe_emit_input_stopped(self, turn: _Turn) -> None:
+        if not turn.ev_input_stopped:
+            turn.ev_input_stopped = True
+            self._emit(
+                "input_speech_stopped", llm.InputSpeechStoppedEvent(user_transcription_enabled=True)
+            )
+            turn.phase = _Phase.USER_FINISHED
+
+    def _emit_transcript_updated(self, turn: _Turn) -> None:
+        self._emit(
+            "input_audio_transcription_completed",
+            llm.InputTranscriptionCompleted(
+                item_id=turn.input_id,
+                transcript=turn.curr_transcript,
+                is_final=False,
+            ),
+        )
+
+    def _maybe_emit_transcript_completed(self, turn: _Turn) -> None:
+        if not turn.ev_trans_completed:
+            turn.ev_trans_completed = True
+            self._emit(
+                "input_audio_transcription_completed",
+                # Q: does input_id need to match /w the _ResponseGeneration.input_id?
+                llm.InputTranscriptionCompleted(
+                    item_id=turn.input_id,
+                    transcript=turn.curr_transcript,
+                    is_final=True,
+                ),
+            )
+
+    def _maybe_emit_generation_created(self, turn: _Turn) -> None:
+        if not turn.ev_generation_sent:
+            turn.ev_generation_sent = True
+            logger.debug(
+                f"[GEN] TurnTracker calling emit_generation_fn() for turn_id={turn.turn_id}"
+            )
+            self._emit_generation_fn()
+            turn.phase = _Phase.ASSISTANT_RESPONDING
+        else:
+            logger.debug(f"[GEN] TurnTracker SKIPPED - already sent for turn_id={turn.turn_id}")
+
+
+def _classify(ev: dict) -> str:
+    e = ev.get("event", {})
+    if "textOutput" in e and e["textOutput"]["role"] == "USER":
+        return "USER_TEXT_PARTIAL"
+
+    if "contentStart" in e and e["contentStart"]["type"] == "TOOL":
+        return "TOOL_OUTPUT_CONTENT_START"
+
+    if "contentStart" in e and e["contentStart"]["role"] == "ASSISTANT":
+        add = e["contentStart"].get("additionalModelFields", "")
+        if "SPECULATIVE" in add:
+            return "ASSISTANT_SPEC_START"
+
+    if "textOutput" in e and e["textOutput"]["content"] == BARGE_IN_CONTENT:
+        return "BARGE_IN"
+
+    # note: there cannot be any audio events for the user in the output event loop
+    # therefore, we know that the audio event must be for the assistant
+    if "contentEnd" in e and e["contentEnd"]["type"] == "AUDIO":
+        return "ASSISTANT_AUDIO_END"
+
+    return ""

livekit/plugins/aws/experimental/realtime/types.py

@@ -0,0 +1,38 @@
+from typing import Literal
+
+TURN_DETECTION = Literal["HIGH", "MEDIUM", "LOW"]
+MODALITIES = Literal["audio", "mixed"]
+REALTIME_MODELS = Literal["amazon.nova-sonic-v1:0", "amazon.nova-2-sonic-v1:0"]
+
+SONIC1_VOICES = Literal[
+    "matthew",  # English (US) - Masculine
+    "tiffany",  # English (US) - Feminine
+    "amy",  # English (GB) - Feminine
+    "lupe",  # Spanish - Feminine
+    "carlos",  # Spanish - Masculine
+    "ambre",  # French - Feminine
+    "florian",  # French - Masculine
+    "greta",  # German - Feminine
+    "lennart",  # German - Masculine
+    "beatrice",  # Italian - Feminine
+    "lorenzo",  # Italian - Masculine
+]
+
+SONIC2_VOICES = Literal[
+    "matthew",  # English (US) - Masculine - Polyglot
+    "tiffany",  # English (US) - Feminine - Polyglot
+    "amy",  # English (GB) - Feminine
+    "olivia",  # English (US) - Feminine
+    "lupe",  # Spanish - Feminine
+    "carlos",  # Spanish - Masculine
+    "ambre",  # French - Feminine
+    "florian",  # French - Masculine
+    "tina",  # German - Feminine
+    "lennart",  # German - Masculine
+    "beatrice",  # Italian - Feminine
+    "lorenzo",  # Italian - Masculine
+    "carolina",  # Portuguese (Brazilian) - Feminine
+    "leo",  # Portuguese (Brazilian) - Masculine
+    "arjun",  # Hindi - Masculine
+    "kiara",  # Hindi - Feminine
+]

livekit/plugins/aws/llm.py

@@ -14,15 +14,21 @@
 # limitations under the License.
 from __future__ import annotations
 
-import asyncio
 import os
 from dataclasses import dataclass
-from typing import Any, Literal
+from typing import Any, cast
 
-import boto3
+import aioboto3  # type: ignore
+from botocore.config import Config  # type: ignore
 
 from livekit.agents import APIConnectionError, APIStatusError, llm
-from livekit.agents.llm import ChatContext, FunctionTool, FunctionToolCall, ToolChoice
+from livekit.agents.llm import (
+    ChatContext,
+    FunctionTool,
+    FunctionToolCall,
+    RawFunctionTool,
+    ToolChoice,
+)
 from livekit.agents.types import (
     DEFAULT_API_CONNECT_OPTIONS,
     NOT_GIVEN,
@@ -32,34 +38,39 @@ from livekit.agents.types import (
 from livekit.agents.utils import is_given
 
 from .log import logger
-from .utils import get_aws_credentials, to_chat_ctx, to_fnc_ctx
+from .utils import to_fnc_ctx
 
-TEXT_MODEL = Literal["anthropic.claude-3-5-sonnet-20241022-v2:0"]
+DEFAULT_TEXT_MODEL = "amazon.nova-2-lite-v1:0"
 
 
 @dataclass
 class _LLMOptions:
-    model: str | TEXT_MODEL
+    model: str
     temperature: NotGivenOr[float]
     tool_choice: NotGivenOr[ToolChoice]
     max_output_tokens: NotGivenOr[int]
     top_p: NotGivenOr[float]
     additional_request_fields: NotGivenOr[dict[str, Any]]
+    cache_system: bool
+    cache_tools: bool
 
 
 class LLM(llm.LLM):
     def __init__(
         self,
         *,
-        model: NotGivenOr[str | TEXT_MODEL] = NOT_GIVEN,
+        model: NotGivenOr[str] = DEFAULT_TEXT_MODEL,
         api_key: NotGivenOr[str] = NOT_GIVEN,
         api_secret: NotGivenOr[str] = NOT_GIVEN,
-        region: NotGivenOr[str] = NOT_GIVEN,
+        region: NotGivenOr[str] = "us-east-1",
         temperature: NotGivenOr[float] = NOT_GIVEN,
         max_output_tokens: NotGivenOr[int] = NOT_GIVEN,
         top_p: NotGivenOr[float] = NOT_GIVEN,
         tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
         additional_request_fields: NotGivenOr[dict[str, Any]] = NOT_GIVEN,
+        cache_system: bool = False,
+        cache_tools: bool = False,
+        session: aioboto3.Session | None = None,
     ) -> None:
         """
         Create a new instance of AWS Bedrock LLM.
@@ -70,7 +81,8 @@ class LLM(llm.LLM):
         See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse_stream.html for more details on the AWS Bedrock Runtime API.
 
         Args:
-            model (TEXT_MODEL, optional): model or inference profile arn to use(https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-use.html). Defaults to 'anthropic.claude-3-5-sonnet-20240620-v1:0'.
+            model (str, optional): model or inference profile arn to use(https://docs.aws.amazon.com/bedrock/latest/userguide/inference-profiles-use.html).
+                Defaults to 'amazon.nova-2-lite-v1:0'.
             api_key(str, optional): AWS access key id.
             api_secret(str, optional): AWS secret access key
             region (str, optional): The region to use for AWS API requests. Defaults value is "us-east-1".
@@ -79,36 +91,57 @@ class LLM(llm.LLM):
             top_p (float, optional): The nucleus sampling probability for response generation. Defaults to None.
             tool_choice (ToolChoice, optional): Specifies whether to use tools during response generation. Defaults to "auto".
             additional_request_fields (dict[str, Any], optional): Additional request fields to send to the AWS Bedrock Converse API. Defaults to None.
+            cache_system (bool, optional): Caches system messages to reduce token usage. Defaults to False.
+            cache_tools (bool, optional): Caches tool definitions to reduce token usage. Defaults to False.
+            session (aioboto3.Session, optional): Optional aioboto3 session to use.
         """  # noqa: E501
         super().__init__()
-        self._api_key, self._api_secret, self._region = get_aws_credentials(
-            api_key, api_secret, region
+
+        self._session = session or aioboto3.Session(
+            aws_access_key_id=api_key if is_given(api_key) else None,
+            aws_secret_access_key=api_secret if is_given(api_secret) else None,
+            region_name=region if is_given(region) else None,
         )
 
-        model = model if is_given(model) else os.environ.get("BEDROCK_INFERENCE_PROFILE_ARN")
-        if not model:
+        bedrock_model = (
+            model if is_given(model) else os.environ.get("BEDROCK_INFERENCE_PROFILE_ARN")
+        )
+        if not bedrock_model:
             raise ValueError(
                 "model or inference profile arn must be set using the argument or by setting the BEDROCK_INFERENCE_PROFILE_ARN environment variable."  # noqa: E501
             )
         self._opts = _LLMOptions(
-            model=model,
+            model=bedrock_model,
             temperature=temperature,
             tool_choice=tool_choice,
             max_output_tokens=max_output_tokens,
             top_p=top_p,
             additional_request_fields=additional_request_fields,
+            cache_system=cache_system,
+            cache_tools=cache_tools,
         )
 
+    @property
+    def model(self) -> str:
+        return self._opts.model
+
+    @property
+    def provider(self) -> str:
+        return "AWS Bedrock"
+
     def chat(
         self,
         *,
         chat_ctx: ChatContext,
-        tools: list[FunctionTool] | None = None,
+        tools: list[FunctionTool | RawFunctionTool] | None = None,
+        parallel_tool_calls: NotGivenOr[bool] = NOT_GIVEN,
         conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
-        temperature: NotGivenOr[float] = NOT_GIVEN,
         tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
+        temperature: NotGivenOr[float] = NOT_GIVEN,
+        extra_kwargs: NotGivenOr[dict[str, Any]] = NOT_GIVEN,
     ) -> LLMStream:
-        opts = {}
+        opts: dict[str, Any] = {}
+        extra_kwargs = extra_kwargs if is_given(extra_kwargs) else {}
 
         if is_given(self._opts.model):
             opts["modelId"] = self._opts.model
@@ -119,8 +152,14 @@ class LLM(llm.LLM):
             if not tools:
                 return None
 
-            tool_config: dict[str, Any] = {"tools": to_fnc_ctx(tools)}
-            tool_choice = tool_choice if is_given(tool_choice) else self._opts.tool_choice
+            tools_list = to_fnc_ctx(tools)
+            if self._opts.cache_tools:
+                tools_list.append({"cachePoint": {"type": "default"}})
+
+            tool_config: dict[str, Any] = {"tools": tools_list}
+            tool_choice = (
+                cast(ToolChoice, tool_choice) if is_given(tool_choice) else self._opts.tool_choice
+            )
             if is_given(tool_choice):
                 if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
                     tool_config["toolChoice"] = {"tool": {"name": tool_choice["function"]["name"]}}
@@ -136,12 +175,17 @@ class LLM(llm.LLM):
         tool_config = _get_tool_config()
         if tool_config:
             opts["toolConfig"] = tool_config
-        messages, system_message = to_chat_ctx(chat_ctx, id(self))
+        messages, extra_data = chat_ctx.to_provider_format(format="aws")
         opts["messages"] = messages
-        if system_message:
-            opts["system"] = [system_message]
-
-        inference_config = {}
+        if extra_data.system_messages:
+            system_messages: list[dict[str, str | dict]] = [
+                {"text": content} for content in extra_data.system_messages
+            ]
+            if self._opts.cache_system:
+                system_messages.append({"cachePoint": {"type": "default"}})
+            opts["system"] = system_messages
+
+        inference_config: dict[str, Any] = {}
         if is_given(self._opts.max_output_tokens):
             inference_config["maxTokens"] = self._opts.max_output_tokens
         temperature = temperature if is_given(temperature) else self._opts.temperature
@@ -156,11 +200,9 @@ class LLM(llm.LLM):
 
         return LLMStream(
             self,
-            aws_access_key_id=self._api_key,
-            aws_secret_access_key=self._api_secret,
-            region_name=self._region,
             chat_ctx=chat_ctx,
-            tools=tools,
+            tools=tools or [],
+            session=self._session,
             conn_options=conn_options,
             extra_kwargs=opts,
         )
@@ -171,24 +213,16 @@ class LLMStream(llm.LLMStream):
         self,
         llm: LLM,
         *,
-        aws_access_key_id: str,
-        aws_secret_access_key: str,
-        region_name: str,
         chat_ctx: ChatContext,
+        session: aioboto3.Session,
         conn_options: APIConnectOptions,
-        tools: list[FunctionTool] | None,
+        tools: list[FunctionTool | RawFunctionTool],
         extra_kwargs: dict[str, Any],
     ) -> None:
         super().__init__(llm, chat_ctx=chat_ctx, tools=tools, conn_options=conn_options)
-        self._client = boto3.client(
-            "bedrock-runtime",
-            region_name=region_name,
-            aws_access_key_id=aws_access_key_id,
-            aws_secret_access_key=aws_secret_access_key,
-        )
         self._llm: LLM = llm
         self._opts = extra_kwargs
-
+        self._session = session
         self._tool_call_id: str | None = None
         self._fnc_name: str | None = None
         self._fnc_raw_arguments: str | None = None
@@ -197,23 +231,22 @@ class LLMStream(llm.LLMStream):
     async def _run(self) -> None:
         retryable = True
         try:
-            response = self._client.converse_stream(**self._opts)  # type: ignore
-            request_id = response["ResponseMetadata"]["RequestId"]
-            if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
-                raise APIStatusError(
-                    f"aws bedrock llm: error generating content: {response}",
-                    retryable=False,
-                    request_id=request_id,
-                )
-
-            for chunk in response["stream"]:
-                chat_chunk = self._parse_chunk(request_id, chunk)
-                if chat_chunk is not None:
-                    retryable = False
-                    self._event_ch.send_nowait(chat_chunk)
-
-                # Let other coroutines run
-                await asyncio.sleep(0)
+            config = Config(user_agent_extra="x-client-framework:livekit-plugins-aws")
+            async with self._session.client("bedrock-runtime", config=config) as client:
+                response = await client.converse_stream(**self._opts)
+                request_id = response["ResponseMetadata"]["RequestId"]
+                if response["ResponseMetadata"]["HTTPStatusCode"] != 200:
+                    raise APIStatusError(
+                        f"aws bedrock llm: error generating content: {response}",
+                        retryable=False,
+                        request_id=request_id,
+                    )
+
+                async for chunk in response["stream"]:
+                    chat_chunk = self._parse_chunk(request_id, chunk)
+                    if chat_chunk is not None:
+                        retryable = False
+                        self._event_ch.send_nowait(chat_chunk)
 
         except Exception as e:
             raise APIConnectionError(
@@ -223,37 +256,42 @@ class LLMStream(llm.LLMStream):
 
     def _parse_chunk(self, request_id: str, chunk: dict) -> llm.ChatChunk | None:
         if "contentBlockStart" in chunk:
-            tool_use = chunk["contentBlockStart"]["start"]["toolUse"]
-            self._tool_call_id = tool_use["toolUseId"]
-            self._fnc_name = tool_use["name"]
-            self._fnc_raw_arguments = ""
+            start = chunk["contentBlockStart"]["start"]
+            if "toolUse" in start:
+                tool_use = start["toolUse"]
+                self._tool_call_id = tool_use["toolUseId"]
+                self._fnc_name = tool_use["name"]
+                self._fnc_raw_arguments = ""
 
         elif "contentBlockDelta" in chunk:
             delta = chunk["contentBlockDelta"]["delta"]
             if "toolUse" in delta:
                 self._fnc_raw_arguments += delta["toolUse"]["input"]
             elif "text" in delta:
-                self._text += delta["text"]
+                return llm.ChatChunk(
+                    id=request_id,
+                    delta=llm.ChoiceDelta(content=delta["text"], role="assistant"),
+                )
+            else:
+                logger.warning(f"aws bedrock llm: unknown chunk type: {chunk}")
 
         elif "metadata" in chunk:
             metadata = chunk["metadata"]
             return llm.ChatChunk(
-                request_id=request_id,
+                id=request_id,
                 usage=llm.CompletionUsage(
                     completion_tokens=metadata["usage"]["outputTokens"],
                     prompt_tokens=metadata["usage"]["inputTokens"],
                     total_tokens=metadata["usage"]["totalTokens"],
+                    prompt_cached_tokens=(
+                        metadata["usage"]["cacheReadInputTokens"]
+                        if "cacheReadInputTokens" in metadata["usage"]
+                        else 0
+                    ),
                 ),
             )
         elif "contentBlockStop" in chunk:
-            if self._text:
-                chat_chunk = llm.ChatChunk(
-                    id=request_id,
-                    delta=llm.ChoiceDelta(content=self._text, role="assistant"),
-                )
-                self._text = ""
-                return chat_chunk
-            elif self._tool_call_id:
+            if self._tool_call_id:
                 if self._tool_call_id is None:
                     logger.warning("aws bedrock llm: no tool call id in the response")
                     return None

livekit/plugins/aws/log.py

@@ -1,3 +1,7 @@
 import logging
 
 logger = logging.getLogger("livekit.plugins.aws")
+smithy_logger = logging.getLogger("smithy_aws_event_stream.aio")
+smithy_logger.setLevel(logging.INFO)
+bedrock_client_logger = logging.getLogger("aws_sdk_bedrock_runtime.client")
+bedrock_client_logger.setLevel(logging.INFO)

livekit/plugins/aws/models.py

@@ -1,7 +1,7 @@
 from typing import Literal
 
-TTS_SPEECH_ENGINE = Literal["standard", "neural", "long-form", "generative"]
-TTS_LANGUAGE = Literal[
+TTSSpeechEngine = Literal["standard", "neural", "long-form", "generative"]
+TTSLanguages = Literal[
     "arb",
     "cmn-CN",
     "cy-GB",
@@ -45,4 +45,5 @@ TTS_LANGUAGE = Literal[
     "de-CH",
 ]
 
-TTS_OUTPUT_FORMAT = Literal["mp3"]
+TTSEncoding = Literal["mp3"]
+TTSTextType = Literal["text", "ssml"]