livekit-plugins-aws 0.1.0__py3-none-any.whl → 1.0.0__py3-none-any.whl

livekit/plugins/aws/llm.py

@@ -14,50 +14,52 @@
 # limitations under the License.
 from __future__ import annotations
 
-import asyncio
 import os
 from dataclasses import dataclass
-from typing import Any, Literal, MutableSet, Union
+from typing import Any, Literal
 
-import boto3
-from livekit.agents import (
-    APIConnectionError,
-    APIStatusError,
-    llm,
+import aioboto3
+
+from livekit.agents import APIConnectionError, APIStatusError, llm
+from livekit.agents.llm import ChatContext, FunctionTool, FunctionToolCall, ToolChoice
+from livekit.agents.types import (
+    DEFAULT_API_CONNECT_OPTIONS,
+    NOT_GIVEN,
+    APIConnectOptions,
+    NotGivenOr,
 )
-from livekit.agents.llm import LLMCapabilities, ToolChoice, _create_ai_function_info
-from livekit.agents.types import DEFAULT_API_CONNECT_OPTIONS, APIConnectOptions
+from livekit.agents.utils import is_given
 
-from ._utils import _build_aws_ctx, _build_tools, _get_aws_credentials
 from .log import logger
+from .utils import get_aws_async_session, to_chat_ctx, to_fnc_ctx
 
 TEXT_MODEL = Literal["anthropic.claude-3-5-sonnet-20241022-v2:0"]
-DEFAULT_REGION = "us-east-1"
 
 
 @dataclass
-class LLMOptions:
-    model: TEXT_MODEL | str
-    temperature: float | None
-    tool_choice: Union[ToolChoice, Literal["auto", "required", "none"]] = "auto"
-    max_output_tokens: int | None = None
-    top_p: float | None = None
-    additional_request_fields: dict[str, Any] | None = None
+class _LLMOptions:
+    model: str | TEXT_MODEL
+    temperature: NotGivenOr[float]
+    tool_choice: NotGivenOr[ToolChoice]
+    max_output_tokens: NotGivenOr[int]
+    top_p: NotGivenOr[float]
+    additional_request_fields: NotGivenOr[dict[str, Any]]
 
 
 class LLM(llm.LLM):
     def __init__(
         self,
         *,
-        model: TEXT_MODEL | str = "anthropic.claude-3-5-sonnet-20240620-v1:0",
-        api_key: str | None = None,
-        api_secret: str | None = None,
-        region: str = "us-east-1",
-        temperature: float = 0.8,
-        max_output_tokens: int | None = None,
-        top_p: float | None = None,
-        tool_choice: Union[ToolChoice, Literal["auto", "required", "none"]] = "auto",
-        additional_request_fields: dict[str, Any] | None = None,
+        model: NotGivenOr[str | TEXT_MODEL] = NOT_GIVEN,
+        api_key: NotGivenOr[str] = NOT_GIVEN,
+        api_secret: NotGivenOr[str] = NOT_GIVEN,
+        region: NotGivenOr[str] = NOT_GIVEN,
+        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,
+        session: aioboto3.Session | None = None,
     ) -> None:
         """
         Create a new instance of AWS Bedrock LLM.
@@ -65,7 +67,7 @@ class LLM(llm.LLM):
         ``api_key``  and ``api_secret`` must be set to your AWS Access key id and secret access key, either using the argument or by setting the
         ``AWS_ACCESS_KEY_ID`` and ``AWS_SECRET_ACCESS_KEY`` environmental variables.
 
-        See https://boto3.amazonaws.com/v1/documentation/api/latest/reference/services/bedrock-runtime/client/converse_stream.html for more details on the the AWS Bedrock Runtime API.
+        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'.
@@ -75,67 +77,94 @@ class LLM(llm.LLM):
             temperature (float, optional): Sampling temperature for response generation. Defaults to 0.8.
             max_output_tokens (int, optional): Maximum number of tokens to generate in the output. Defaults to None.
             top_p (float, optional): The nucleus sampling probability for response generation. Defaults to None.
-            tool_choice (ToolChoice or Literal["auto", "required", "none"], optional): Specifies whether to use tools during response generation. Defaults to "auto".
+            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.
-        """
-        super().__init__(
-            capabilities=LLMCapabilities(
-                supports_choices_on_int=True,
-                requires_persistent_functions=True,
-            )
-        )
-        self._api_key, self._api_secret = _get_aws_credentials(
-            api_key, api_secret, region
+            session (aioboto3.Session, optional): Optional aioboto3 session to use.
+        """  # noqa: E501
+        super().__init__()
+
+        self._session = session or get_aws_async_session(
+            api_key=api_key if is_given(api_key) else None,
+            api_secret=api_secret if is_given(api_secret) else None,
+            region=region if is_given(region) else None,
         )
 
-        self._model = model or os.environ.get("BEDROCK_INFERENCE_PROFILE_ARN")
-        if not self._model:
+        model = model if is_given(model) else os.environ.get("BEDROCK_INFERENCE_PROFILE_ARN")
+        if not model:
             raise ValueError(
-                "model or inference profile arn must be set using the argument or by setting the BEDROCK_INFERENCE_PROFILE_ARN environment variable."
+                "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=self._model,
+        self._opts = _LLMOptions(
+            model=model,
             temperature=temperature,
             tool_choice=tool_choice,
             max_output_tokens=max_output_tokens,
             top_p=top_p,
             additional_request_fields=additional_request_fields,
         )
-        self._region = region
-        self._running_fncs: MutableSet[asyncio.Task[Any]] = set()
 
     def chat(
         self,
         *,
-        chat_ctx: llm.ChatContext,
+        chat_ctx: ChatContext,
+        tools: list[FunctionTool] | None = None,
         conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
-        fnc_ctx: llm.FunctionContext | None = None,
-        temperature: float | None = None,
-        n: int | None = 1,
-        parallel_tool_calls: bool | None = None,
-        tool_choice: Union[ToolChoice, Literal["auto", "required", "none"]]
-        | None = None,
-    ) -> "LLMStream":
-        if tool_choice is None:
-            tool_choice = self._opts.tool_choice
-
-        if temperature is None:
-            temperature = self._opts.temperature
+        temperature: NotGivenOr[float] = NOT_GIVEN,
+        tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
+    ) -> LLMStream:
+        opts = {}
+
+        if is_given(self._opts.model):
+            opts["modelId"] = self._opts.model
+
+        def _get_tool_config() -> dict[str, Any] | None:
+            nonlocal tool_choice
+
+            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
+            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"]}}
+                elif tool_choice == "required":
+                    tool_config["toolChoice"] = {"any": {}}
+                elif tool_choice == "auto":
+                    tool_config["toolChoice"] = {"auto": {}}
+                else:
+                    return None
+
+            return tool_config
+
+        tool_config = _get_tool_config()
+        if tool_config:
+            opts["toolConfig"] = tool_config
+        messages, system_message = to_chat_ctx(chat_ctx, id(self))
+        opts["messages"] = messages
+        if system_message:
+            opts["system"] = [system_message]
+
+        inference_config = {}
+        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
+        if is_given(temperature):
+            inference_config["temperature"] = temperature
+        if is_given(self._opts.top_p):
+            inference_config["topP"] = self._opts.top_p
+
+        opts["inferenceConfig"] = inference_config
+        if is_given(self._opts.additional_request_fields):
+            opts["additionalModelRequestFields"] = self._opts.additional_request_fields
 
         return LLMStream(
             self,
-            model=self._opts.model,
-            aws_access_key_id=self._api_key,
-            aws_secret_access_key=self._api_secret,
-            region_name=self._region,
-            max_output_tokens=self._opts.max_output_tokens,
-            top_p=self._opts.top_p,
-            additional_request_fields=self._opts.additional_request_fields,
             chat_ctx=chat_ctx,
-            fnc_ctx=fnc_ctx,
+            tools=tools,
+            session=self._session,
             conn_options=conn_options,
-            temperature=temperature,
-            tool_choice=tool_choice,
+            extra_kwargs=opts,
         )
 
 
@@ -144,107 +173,39 @@ class LLMStream(llm.LLMStream):
         self,
         llm: LLM,
         *,
-        model: str | TEXT_MODEL,
-        aws_access_key_id: str | None,
-        aws_secret_access_key: str | None,
-        region_name: str,
-        chat_ctx: llm.ChatContext,
+        chat_ctx: ChatContext,
+        session: aioboto3.Session,
         conn_options: APIConnectOptions,
-        fnc_ctx: llm.FunctionContext | None,
-        temperature: float | None,
-        max_output_tokens: int | None,
-        top_p: float | None,
-        tool_choice: Union[ToolChoice, Literal["auto", "required", "none"]],
-        additional_request_fields: dict[str, Any] | None,
+        tools: list[FunctionTool] | None,
+        extra_kwargs: dict[str, Any],
     ) -> None:
-        super().__init__(
-            llm, chat_ctx=chat_ctx, fnc_ctx=fnc_ctx, 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._model = model
+        super().__init__(llm, chat_ctx=chat_ctx, tools=tools, conn_options=conn_options)
         self._llm: LLM = llm
-        self._max_output_tokens = max_output_tokens
-        self._top_p = top_p
-        self._temperature = temperature
-        self._tool_choice = tool_choice
-        self._additional_request_fields = additional_request_fields
-
-    async def _run(self) -> None:
+        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
         self._text: str = ""
-        retryable = True
 
+    async def _run(self) -> None:
+        retryable = True
         try:
-            opts: dict[str, Any] = {}
-            messages, system_instruction = _build_aws_ctx(self._chat_ctx, id(self))
-            messages = _merge_messages(messages)
-
-            def _get_tool_config() -> dict[str, Any] | None:
-                if not (self._fnc_ctx and self._fnc_ctx.ai_functions):
-                    return None
-
-                tools = _build_tools(self._fnc_ctx)
-                config: dict[str, Any] = {"tools": tools}
-
-                if isinstance(self._tool_choice, ToolChoice):
-                    config["toolChoice"] = {"tool": {"name": self._tool_choice.name}}
-                elif self._tool_choice == "required":
-                    config["toolChoice"] = {"any": {}}
-                elif self._tool_choice == "auto":
-                    config["toolChoice"] = {"auto": {}}
-                else:
-                    return None
-
-                return config
-
-            tool_config = _get_tool_config()
-            if tool_config:
-                opts["toolConfig"] = tool_config
-
-            if self._additional_request_fields:
-                opts["additionalModelRequestFields"] = _strip_nones(
-                    self._additional_request_fields
-                )
-            if system_instruction:
-                opts["system"] = [system_instruction]
-
-            inference_config = _strip_nones(
-                {
-                    "maxTokens": self._max_output_tokens,
-                    "temperature": self._temperature,
-                    "topP": self._top_p,
-                }
-            )
-            response = self._client.converse_stream(
-                modelId=self._model,
-                messages=messages,
-                inferenceConfig=inference_config,
-                **_strip_nones(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)
+            async with self._session.client("bedrock-runtime") as client:
+                response = await 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,
+                    )
+
+                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(
@@ -258,93 +219,53 @@ class LLMStream(llm.LLMStream):
             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"]
-        elif "contentBlockStop" in chunk:
-            if self._text:
-                chat_chunk = llm.ChatChunk(
-                    request_id=request_id,
-                    choices=[
-                        llm.Choice(
-                            delta=llm.ChoiceDelta(content=self._text, role="assistant"),
-                            index=chunk["contentBlockStop"]["contentBlockIndex"],
-                        )
-                    ],
+                return llm.ChatChunk(
+                    id=request_id,
+                    delta=llm.ChoiceDelta(content=delta["text"], role="assistant"),
                 )
-                self._text = ""
-                return chat_chunk
-            elif self._tool_call_id:
-                return self._try_build_function(request_id, chunk)
-
-        return None
-
-    def _try_build_function(self, request_id: str, chunk: dict) -> llm.ChatChunk | None:
-        if self._tool_call_id is None:
-            logger.warning("aws bedrock llm: no tool call id in the response")
-            return None
-        if self._fnc_name is None:
-            logger.warning("aws bedrock llm: no function name in the response")
-            return None
-        if self._fnc_raw_arguments is None:
-            logger.warning("aws bedrock llm: no function arguments in the response")
-            return None
-        if self._fnc_ctx is None:
-            logger.warning(
-                "aws bedrock llm: stream tried to run function without function context"
+            else:
+                logger.warning(f"aws bedrock llm: unknown chunk type: {chunk}")
+
+        elif "metadata" in chunk:
+            metadata = chunk["metadata"]
+            return llm.ChatChunk(
+                id=request_id,
+                usage=llm.CompletionUsage(
+                    completion_tokens=metadata["usage"]["outputTokens"],
+                    prompt_tokens=metadata["usage"]["inputTokens"],
+                    total_tokens=metadata["usage"]["totalTokens"],
+                ),
             )
-            return None
-
-        fnc_info = _create_ai_function_info(
-            self._fnc_ctx,
-            self._tool_call_id,
-            self._fnc_name,
-            self._fnc_raw_arguments,
-        )
-
-        self._tool_call_id = self._fnc_name = self._fnc_raw_arguments = None
-        self._function_calls_info.append(fnc_info)
-
-        return llm.ChatChunk(
-            request_id=request_id,
-            choices=[
-                llm.Choice(
+        elif "contentBlockStop" in chunk:
+            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
+                if self._fnc_name is None:
+                    logger.warning("aws bedrock llm: no function name in the response")
+                    return None
+                if self._fnc_raw_arguments is None:
+                    logger.warning("aws bedrock llm: no function arguments in the response")
+                    return None
+                chat_chunk = llm.ChatChunk(
+                    id=request_id,
                     delta=llm.ChoiceDelta(
                         role="assistant",
-                        tool_calls=[fnc_info],
+                        tool_calls=[
+                            FunctionToolCall(
+                                arguments=self._fnc_raw_arguments,
+                                name=self._fnc_name,
+                                call_id=self._tool_call_id,
+                            ),
+                        ],
                     ),
-                    index=chunk["contentBlockStop"]["contentBlockIndex"],
                 )
-            ],
-        )
-
-
-def _merge_messages(
-    messages: list[dict],
-) -> list[dict]:
-    # Anthropic enforces alternating messages
-    combined_messages: list[dict] = []
-    for m in messages:
-        if len(combined_messages) == 0 or m["role"] != combined_messages[-1]["role"]:
-            combined_messages.append(m)
-            continue
-        last_message = combined_messages[-1]
-        if not isinstance(last_message["content"], list) or not isinstance(
-            m["content"], list
-        ):
-            logger.error("message content is not a list")
-            continue
-
-        last_message["content"].extend(m["content"])
-
-    if len(combined_messages) == 0 or combined_messages[0]["role"] != "user":
-        combined_messages.insert(0, {"role": "user", "content": [{"text": "(empty)"}]})
-
-    return combined_messages
-
-
-def _strip_nones(d: dict[str, Any]) -> dict[str, Any]:
-    return {k: v for k, v in d.items() if v is not None}
+                self._tool_call_id = self._fnc_name = self._fnc_raw_arguments = None
+                return chat_chunk
+        return None

livekit/plugins/aws/models.py

@@ -45,4 +45,4 @@ TTS_LANGUAGE = Literal[
     "de-CH",
 ]
 
-TTS_OUTPUT_FORMAT = Literal["pcm", "mp3"]
+TTS_OUTPUT_FORMAT = Literal["mp3"]