livekit-plugins-volcenginee 1.3.0__py3-none-any.whl

livekit/plugins/volcengine/__init__.py

@@ -0,0 +1,28 @@
+from .llm import LLM
+from .stt import STT
+from .tts import TTS
+from .realtime import RealtimeModel
+from .version import __version__
+
+__all__ = ["TTS", "LLM", "STT", "RealtimeModel", "__version__"]
+
+from livekit.agents import Plugin
+
+from .log import logger
+
+
+class VolcenginePlugin(Plugin):
+    def __init__(self):
+        super().__init__(__name__, __version__, __package__, logger)
+
+
+Plugin.register_plugin(VolcenginePlugin())
+
+# Cleanup docs of unexported modules
+_module = dir()
+NOT_IN_ALL = [m for m in _module if m not in __all__]
+
+__pdoc__ = {}
+
+for n in NOT_IN_ALL:
+    __pdoc__[n] = False

livekit/plugins/volcengine/llm.py

@@ -0,0 +1,283 @@
+from __future__ import annotations
+
+import os
+from dataclasses import dataclass
+from typing import Any
+
+import httpx
+import openai
+from openai.types.chat import ChatCompletionChunk, ChatCompletionToolChoiceOptionParam
+from openai.types.chat.chat_completion_chunk import Choice
+
+from livekit.agents import APIConnectionError, APIStatusError, APITimeoutError, llm
+from livekit.agents.llm import ToolChoice
+from livekit.agents.llm.chat_context import ChatContext
+from livekit.agents.types import (
+    DEFAULT_API_CONNECT_OPTIONS,
+    NOT_GIVEN,
+    APIConnectOptions,
+    NotGivenOr,
+)
+from livekit.agents.utils import is_given
+
+from .log import logger
+from .utils import to_chat_ctx, to_fnc_ctx
+
+
+@dataclass
+class _LLMOptions:
+    model: str
+    user: NotGivenOr[str]
+    temperature: NotGivenOr[float]
+    parallel_tool_calls: NotGivenOr[bool]
+    tool_choice: NotGivenOr[ToolChoice]
+    store: NotGivenOr[bool]
+    metadata: NotGivenOr[dict[str, str]]
+
+
+class LLM(llm.LLM):
+    def __init__(
+        self,
+        *,
+        model: str = "doubao-1-5-lite-32k-250115",
+        api_key: NotGivenOr[str] = NOT_GIVEN,
+        base_url: NotGivenOr[str] = "https://ark.cn-beijing.volces.com/api/v3/",
+        client: openai.AsyncClient | None = None,
+        user: NotGivenOr[str] = NOT_GIVEN,
+        temperature: NotGivenOr[float] = NOT_GIVEN,
+        parallel_tool_calls: NotGivenOr[bool] = NOT_GIVEN,
+        tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
+        store: NotGivenOr[bool] = NOT_GIVEN,
+        metadata: NotGivenOr[dict[str, str]] = NOT_GIVEN,
+        timeout: httpx.Timeout | None = None,
+    ) -> None:
+        """
+        Create a new instance of Volcengine LLM.
+
+        Args:
+            model: The model to use for the LLM, end
+            api_key: The API key to use for the LLM.
+            base_url: The base URL to use for the LLM.
+        """
+        super().__init__()
+        self._opts = _LLMOptions(
+            model=model,
+            user=user,
+            temperature=temperature,
+            parallel_tool_calls=parallel_tool_calls,
+            tool_choice=tool_choice,
+            store=store,
+            metadata=metadata,
+        )
+        api_key = api_key if is_given(api_key) else os.getenv("VOLCENGINE_LLM_API_KEY")
+        if api_key is None:
+            raise ValueError("api_key is required")
+        self._client = client or openai.AsyncClient(
+            api_key=api_key if is_given(api_key) else None,
+            base_url=base_url if is_given(base_url) else None,
+            max_retries=0,
+            http_client=httpx.AsyncClient(
+                timeout=timeout
+                if timeout
+                else httpx.Timeout(connect=15.0, read=5.0, write=5.0, pool=5.0),
+                follow_redirects=True,
+                limits=httpx.Limits(
+                    max_connections=50,
+                    max_keepalive_connections=50,
+                    keepalive_expiry=120,
+                ),
+            ),
+        )
+
+    def chat(
+        self,
+        *,
+        chat_ctx: ChatContext,
+        tools: list[llm.Tool] | None = None,
+        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
+        parallel_tool_calls: NotGivenOr[bool] = NOT_GIVEN,
+        tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
+        extra_kwargs: NotGivenOr[dict[str, Any]] = NOT_GIVEN,
+    ) -> LLMStream:
+        extra = {}
+
+        if is_given(extra_kwargs):
+            extra.update(extra_kwargs)
+
+        if is_given(self._opts.metadata):
+            extra["metadata"] = self._opts.metadata
+
+        if is_given(self._opts.user):
+            extra["user"] = self._opts.user
+
+        parallel_tool_calls = (
+            parallel_tool_calls
+            if is_given(parallel_tool_calls)
+            else self._opts.parallel_tool_calls
+        )
+        if is_given(parallel_tool_calls):
+            extra["parallel_tool_calls"] = parallel_tool_calls
+
+        tool_choice = tool_choice if is_given(tool_choice) else self._opts.tool_choice  # type: ignore
+        if is_given(tool_choice):
+            oai_tool_choice: ChatCompletionToolChoiceOptionParam
+            if isinstance(tool_choice, dict):
+                oai_tool_choice = {
+                    "type": "function",
+                    "function": {"name": tool_choice["function"]["name"]},
+                }
+                extra["tool_choice"] = oai_tool_choice
+            elif tool_choice in ("auto", "required", "none"):
+                oai_tool_choice = tool_choice
+                extra["tool_choice"] = oai_tool_choice
+        logger.info("llm start", extra={"model": self._opts.model})
+        return LLMStream(
+            self,
+            model=self._opts.model,
+            client=self._client,
+            chat_ctx=chat_ctx,
+            tools=tools or [],
+            conn_options=conn_options,
+            extra_kwargs=extra,
+        )
+
+
+class LLMStream(llm.LLMStream):
+    def __init__(
+        self,
+        llm: LLM,
+        *,
+        model: str,
+        client: openai.AsyncClient,
+        chat_ctx: llm.ChatContext,
+        tools: list[llm.Tool],
+        conn_options: APIConnectOptions,
+        extra_kwargs: dict[str, Any],
+    ) -> None:
+        super().__init__(llm, chat_ctx=chat_ctx, tools=tools, conn_options=conn_options)
+        self._model = model
+        self._client = client
+        self._llm = llm
+        self._extra_kwargs = extra_kwargs
+
+    async def _run(self) -> None:
+        # current function call that we're waiting for full completion (args are streamed)
+        # (defined inside the _run method to make sure the state is reset for each run/attempt)
+        self._oai_stream: openai.AsyncStream[ChatCompletionChunk] | None = None
+        self._tool_call_id: str | None = None
+        self._fnc_name: str | None = None
+        self._fnc_raw_arguments: str | None = None
+        self._tool_index: int | None = None
+        retryable = True
+        first_response = True
+        try:
+            stream: openai.AsyncStream[
+                ChatCompletionChunk
+            ] = await self._client.chat.completions.create(
+                messages=to_chat_ctx(self._chat_ctx, id(self._llm)),
+                tools=to_fnc_ctx(self._tools) if self._tools else openai.NOT_GIVEN,
+                model=self._model,
+                stream_options={"include_usage": True},
+                stream=True,
+                **self._extra_kwargs,
+            )
+
+            async with stream:
+                async for chunk in stream:
+                    for choice in chunk.choices:
+                        chat_chunk = self._parse_choice(chunk.id, choice)
+                        if chat_chunk is not None:
+                            retryable = False
+                            self._event_ch.send_nowait(chat_chunk)
+                        if first_response:
+                            logger.info("llm first response")
+                            first_response = False
+
+                    if chunk.usage is not None:
+                        retryable = False
+                        usage_chunk = llm.ChatChunk(
+                            id=chunk.id,
+                            usage=llm.CompletionUsage(
+                                completion_tokens=chunk.usage.completion_tokens,
+                                prompt_tokens=chunk.usage.prompt_tokens,
+                                total_tokens=chunk.usage.total_tokens,
+                            ),
+                        )
+                        self._event_ch.send_nowait(usage_chunk)
+            logger.info("llm end")
+
+        except openai.APITimeoutError:
+            raise APITimeoutError(retryable=retryable)  # noqa: B904
+        except openai.APIStatusError as e:
+            raise APIStatusError(  # noqa: B904
+                e.message,
+                status_code=e.status_code,
+                request_id=e.request_id,
+                body=e.body,
+                retryable=retryable,
+            )
+        except Exception as e:
+            raise APIConnectionError(retryable=retryable) from e
+
+    def _parse_choice(self, id: str, choice: Choice) -> llm.ChatChunk | None:
+        delta = choice.delta
+
+        if delta is None:
+            return None
+
+        if delta.tool_calls:
+            for tool in delta.tool_calls:
+                if not tool.function:
+                    continue
+
+                call_chunk = None
+                if self._tool_call_id and tool.id and tool.index != self._tool_index:
+                    call_chunk = llm.ChatChunk(
+                        id=id,
+                        delta=llm.ChoiceDelta(
+                            role="assistant",
+                            content=delta.content,
+                            tool_calls=[
+                                llm.FunctionToolCall(
+                                    arguments=self._fnc_raw_arguments or "",
+                                    name=self._fnc_name or "",
+                                    call_id=self._tool_call_id or "",
+                                )
+                            ],
+                        ),
+                    )
+                    self._tool_call_id = self._fnc_name = self._fnc_raw_arguments = None
+
+                if tool.function.name:
+                    self._tool_index = tool.index
+                    self._tool_call_id = tool.id
+                    self._fnc_name = tool.function.name
+                    self._fnc_raw_arguments = tool.function.arguments or ""
+                elif tool.function.arguments:
+                    self._fnc_raw_arguments += tool.function.arguments  # type: ignore
+
+                if call_chunk is not None:
+                    return call_chunk
+
+        if choice.finish_reason in ("tool_calls", "stop") and self._tool_call_id:
+            call_chunk = llm.ChatChunk(
+                id=id,
+                delta=llm.ChoiceDelta(
+                    role="assistant",
+                    content=delta.content,
+                    tool_calls=[
+                        llm.FunctionToolCall(
+                            arguments=self._fnc_raw_arguments or "",
+                            name=self._fnc_name or "",
+                            call_id=self._tool_call_id or "",
+                        )
+                    ],
+                ),
+            )
+            self._tool_call_id = self._fnc_name = self._fnc_raw_arguments = None
+            return call_chunk
+
+        return llm.ChatChunk(
+            id=id,
+            delta=llm.ChoiceDelta(content=delta.content, role="assistant"),
+        )

livekit/plugins/volcengine/log.py

@@ -0,0 +1,3 @@
+from logging import getLogger
+
+logger = getLogger("livekit.plugins.volcengine")