wehub-resource-sync
406 lines
16 KiB
Python
406 lines
16 KiB
Python
# Copyright 2023 LiveKit, Inc.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
from __future__ import annotations
|
|
|
|
import os
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any, Literal, cast
|
|
|
|
import httpx
|
|
|
|
import anthropic
|
|
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.llm.tool_context import Tool
|
|
from livekit.agents.types import (
|
|
DEFAULT_API_CONNECT_OPTIONS,
|
|
NOT_GIVEN,
|
|
APIConnectOptions,
|
|
NotGivenOr,
|
|
)
|
|
from livekit.agents.utils import is_given
|
|
|
|
from .models import ChatModels
|
|
from .utils import CACHE_CONTROL_EPHEMERAL
|
|
|
|
# Claude 4.6+ no longer supports prefilling (trailing assistant messages).
|
|
_NO_PREFILL_PATTERNS = ("claude-sonnet-4-6", "claude-opus-4-6")
|
|
|
|
|
|
def _model_disables_prefill(model: str) -> bool:
|
|
"""Return True if the model does not support assistant message prefilling."""
|
|
return any(model.startswith(p) for p in _NO_PREFILL_PATTERNS)
|
|
|
|
|
|
@dataclass
|
|
class _LLMOptions:
|
|
model: str | ChatModels
|
|
user: NotGivenOr[str]
|
|
temperature: NotGivenOr[float]
|
|
parallel_tool_calls: NotGivenOr[bool]
|
|
tool_choice: NotGivenOr[ToolChoice]
|
|
caching: NotGivenOr[Literal["ephemeral"]]
|
|
top_k: NotGivenOr[int]
|
|
max_tokens: NotGivenOr[int]
|
|
strict_tool_schema: bool
|
|
"""If set to "ephemeral", the system prompt, tools, and chat history will be cached."""
|
|
|
|
|
|
class LLM(llm.LLM):
|
|
def __init__(
|
|
self,
|
|
*,
|
|
model: str | ChatModels = "claude-sonnet-4-6",
|
|
api_key: NotGivenOr[str] = NOT_GIVEN,
|
|
base_url: NotGivenOr[str] = NOT_GIVEN,
|
|
user: NotGivenOr[str] = NOT_GIVEN,
|
|
client: anthropic.AsyncClient | None = None,
|
|
top_k: NotGivenOr[int] = NOT_GIVEN,
|
|
max_tokens: NotGivenOr[int] = NOT_GIVEN,
|
|
temperature: NotGivenOr[float] = NOT_GIVEN,
|
|
parallel_tool_calls: NotGivenOr[bool] = NOT_GIVEN,
|
|
tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
|
|
caching: NotGivenOr[Literal["ephemeral"]] = NOT_GIVEN,
|
|
timeout: NotGivenOr[httpx.Timeout] = NOT_GIVEN,
|
|
_strict_tool_schema: bool = True,
|
|
) -> None:
|
|
"""
|
|
Create a new instance of Anthropic LLM.
|
|
|
|
``api_key`` must be set to your Anthropic API key, either using the argument or by setting
|
|
the ``ANTHROPIC_API_KEY`` environmental variable.
|
|
|
|
model (str | ChatModels): The model to use. Defaults to "claude-sonnet-4-6".
|
|
api_key (str, optional): The Anthropic API key. Defaults to the ANTHROPIC_API_KEY environment variable.
|
|
base_url (str, optional): The base URL for the Anthropic API. Defaults to None.
|
|
user (str, optional): The user for the Anthropic API. Defaults to None.
|
|
client (anthropic.AsyncClient | None): The Anthropic client to use. Defaults to None.
|
|
timeout (httpx.Timeout | None): HTTP timeout configuration for the underlying httpx client.
|
|
Defaults to ``httpx.Timeout(5.0, read=30.0)``, which keeps a tight connect timeout
|
|
while allowing up to 30 s between streamed chunks — long enough for Claude's
|
|
adaptive-thinking phases without masking genuine network stalls.
|
|
Pass a custom ``httpx.Timeout`` to override (e.g. ``httpx.Timeout(5.0, read=60.0)``
|
|
for very large contexts or extended thinking budgets).
|
|
temperature (float, optional): The temperature for the Anthropic API. Defaults to None.
|
|
parallel_tool_calls (bool, optional): Whether to parallelize tool calls. Defaults to None.
|
|
tool_choice (ToolChoice, optional): The tool choice for the Anthropic API. Defaults to "auto".
|
|
caching (Literal["ephemeral"], optional): If set to "ephemeral", caching will be enabled for the system prompt, tools, and chat history.
|
|
""" # noqa: E501
|
|
|
|
super().__init__()
|
|
|
|
self._opts = _LLMOptions(
|
|
model=model,
|
|
user=user,
|
|
temperature=temperature,
|
|
parallel_tool_calls=parallel_tool_calls,
|
|
tool_choice=tool_choice,
|
|
caching=caching,
|
|
top_k=top_k,
|
|
max_tokens=max_tokens,
|
|
strict_tool_schema=_strict_tool_schema,
|
|
)
|
|
anthropic_api_key = api_key if is_given(api_key) else os.environ.get("ANTHROPIC_API_KEY")
|
|
if not anthropic_api_key:
|
|
raise ValueError(
|
|
"Anthropic API key is required, either as argument or set"
|
|
" ANTHROPIC_API_KEY environment variable"
|
|
)
|
|
|
|
self._client = client or anthropic.AsyncClient(
|
|
api_key=anthropic_api_key,
|
|
base_url=base_url if is_given(base_url) else None,
|
|
http_client=httpx.AsyncClient(
|
|
timeout=timeout or httpx.Timeout(5.0, read=30.0),
|
|
follow_redirects=True,
|
|
limits=httpx.Limits(
|
|
max_connections=1000,
|
|
max_keepalive_connections=100,
|
|
keepalive_expiry=120,
|
|
),
|
|
),
|
|
)
|
|
|
|
@property
|
|
def model(self) -> str:
|
|
return self._opts.model
|
|
|
|
@property
|
|
def provider(self) -> str:
|
|
return self._client._base_url.netloc.decode("utf-8")
|
|
|
|
def chat(
|
|
self,
|
|
*,
|
|
chat_ctx: ChatContext,
|
|
tools: list[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.user):
|
|
extra["user"] = self._opts.user
|
|
|
|
if is_given(self._opts.temperature):
|
|
extra["temperature"] = self._opts.temperature
|
|
|
|
if is_given(self._opts.top_k):
|
|
extra["top_k"] = self._opts.top_k
|
|
|
|
extra["max_tokens"] = self._opts.max_tokens if is_given(self._opts.max_tokens) else 1024
|
|
|
|
beta_flag: str | None = None
|
|
if tools:
|
|
from .tools import AnthropicTool
|
|
|
|
tool_ctx = llm.ToolContext(tools)
|
|
tool_schemas = tool_ctx.parse_function_tools(
|
|
"anthropic", strict=self._opts.strict_tool_schema
|
|
)
|
|
|
|
for tool in tool_ctx.provider_tools:
|
|
if isinstance(tool, AnthropicTool):
|
|
tool_schemas.append(tool.to_dict())
|
|
if tool.beta_flag:
|
|
beta_flag = tool.beta_flag
|
|
|
|
extra["tools"] = tool_schemas
|
|
|
|
tool_choice = tool_choice if is_given(tool_choice) else self._opts.tool_choice
|
|
if is_given(tool_choice):
|
|
anthropic_tool_choice: dict[str, Any] | None = {"type": "auto"}
|
|
if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
|
|
anthropic_tool_choice = {
|
|
"type": "tool",
|
|
"name": tool_choice["function"]["name"],
|
|
}
|
|
elif isinstance(tool_choice, str):
|
|
if tool_choice == "required":
|
|
anthropic_tool_choice = {"type": "any"}
|
|
elif tool_choice == "none":
|
|
extra["tools"] = []
|
|
anthropic_tool_choice = None
|
|
if anthropic_tool_choice is not None:
|
|
parallel_tool_calls = (
|
|
parallel_tool_calls
|
|
if is_given(parallel_tool_calls)
|
|
else self._opts.parallel_tool_calls
|
|
)
|
|
if is_given(parallel_tool_calls):
|
|
anthropic_tool_choice["disable_parallel_tool_use"] = not parallel_tool_calls
|
|
extra["tool_choice"] = anthropic_tool_choice
|
|
|
|
# Claude 4.6+ does not support prefilling (trailing assistant messages).
|
|
inject_trailing = _model_disables_prefill(self._opts.model)
|
|
anthropic_ctx, extra_data = chat_ctx.to_provider_format(
|
|
format="anthropic", inject_trailing_user_message=inject_trailing
|
|
)
|
|
messages = cast(list[anthropic.types.MessageParam], anthropic_ctx)
|
|
if extra_data.system_messages:
|
|
extra["system"] = [
|
|
anthropic.types.TextBlockParam(text=content, type="text")
|
|
for content in extra_data.system_messages
|
|
]
|
|
|
|
# add cache control
|
|
if self._opts.caching == "ephemeral":
|
|
if extra.get("system"):
|
|
extra["system"][-1]["cache_control"] = CACHE_CONTROL_EPHEMERAL
|
|
|
|
if extra.get("tools"):
|
|
extra["tools"][-1]["cache_control"] = CACHE_CONTROL_EPHEMERAL
|
|
|
|
seen_assistant = False
|
|
for msg in reversed(messages):
|
|
if (
|
|
msg["role"] == "assistant"
|
|
and (content := msg["content"])
|
|
and not seen_assistant
|
|
):
|
|
content[-1]["cache_control"] = CACHE_CONTROL_EPHEMERAL # type: ignore
|
|
seen_assistant = True
|
|
|
|
elif msg["role"] == "user" and (content := msg["content"]) and seen_assistant:
|
|
content[-1]["cache_control"] = CACHE_CONTROL_EPHEMERAL # type: ignore
|
|
break
|
|
|
|
async def create_anthropic_stream() -> anthropic.AsyncStream[
|
|
anthropic.types.RawMessageStreamEvent
|
|
]:
|
|
if beta_flag:
|
|
stream = await self._client.beta.messages.create(
|
|
betas=[beta_flag],
|
|
messages=messages, # type: ignore[arg-type]
|
|
model=self._opts.model,
|
|
stream=True,
|
|
timeout=conn_options.timeout,
|
|
**extra,
|
|
)
|
|
else:
|
|
stream = await self._client.messages.create(
|
|
messages=messages,
|
|
model=self._opts.model,
|
|
stream=True,
|
|
timeout=conn_options.timeout,
|
|
**extra,
|
|
)
|
|
return cast(anthropic.AsyncStream[anthropic.types.RawMessageStreamEvent], stream)
|
|
|
|
return LLMStream(
|
|
self,
|
|
create_anthropic_stream=create_anthropic_stream,
|
|
chat_ctx=chat_ctx,
|
|
tools=tools or [],
|
|
conn_options=conn_options,
|
|
)
|
|
|
|
|
|
class LLMStream(llm.LLMStream):
|
|
def __init__(
|
|
self,
|
|
llm: LLM,
|
|
*,
|
|
create_anthropic_stream: Callable[
|
|
[], Awaitable[anthropic.AsyncStream[anthropic.types.RawMessageStreamEvent]]
|
|
],
|
|
chat_ctx: llm.ChatContext,
|
|
tools: list[Tool],
|
|
conn_options: APIConnectOptions,
|
|
) -> None:
|
|
super().__init__(llm, chat_ctx=chat_ctx, tools=tools, conn_options=conn_options)
|
|
self._create_anthropic_stream = create_anthropic_stream
|
|
|
|
# current function call that we're waiting for full completion (args are streamed)
|
|
self._tool_call_id: str | None = None
|
|
self._fnc_name: str | None = None
|
|
self._fnc_raw_arguments: str | None = None
|
|
|
|
self._request_id: str = ""
|
|
self._ignoring_cot = False # ignore chain of thought
|
|
self._input_tokens = 0
|
|
self._cache_creation_tokens = 0
|
|
self._cache_read_tokens = 0
|
|
self._output_tokens = 0
|
|
|
|
async def _run(self) -> None:
|
|
retryable = True
|
|
try:
|
|
async with await self._create_anthropic_stream() as stream:
|
|
async for event in stream:
|
|
chat_chunk = self._parse_event(event)
|
|
if chat_chunk is not None:
|
|
self._event_ch.send_nowait(chat_chunk)
|
|
retryable = False
|
|
|
|
# https://docs.anthropic.com/en/docs/build-with-claude/prompt-caching#tracking-cache-performance
|
|
prompt_token = (
|
|
self._input_tokens + self._cache_creation_tokens + self._cache_read_tokens
|
|
)
|
|
self._event_ch.send_nowait(
|
|
llm.ChatChunk(
|
|
id=self._request_id,
|
|
usage=llm.CompletionUsage(
|
|
completion_tokens=self._output_tokens,
|
|
prompt_tokens=prompt_token,
|
|
total_tokens=prompt_token + self._output_tokens,
|
|
prompt_cached_tokens=self._cache_read_tokens,
|
|
cache_creation_tokens=self._cache_creation_tokens,
|
|
cache_read_tokens=self._cache_read_tokens,
|
|
),
|
|
)
|
|
)
|
|
except anthropic.APITimeoutError as e:
|
|
raise APITimeoutError(retryable=retryable) from e
|
|
except anthropic.APIStatusError as e:
|
|
raise APIStatusError(
|
|
e.message,
|
|
status_code=e.status_code,
|
|
request_id=e.request_id,
|
|
body=e.body,
|
|
) from e
|
|
except Exception as e:
|
|
raise APIConnectionError(retryable=retryable) from e
|
|
|
|
def _parse_event(self, event: anthropic.types.RawMessageStreamEvent) -> llm.ChatChunk | None:
|
|
if event.type == "message_start":
|
|
self._request_id = event.message.id
|
|
self._input_tokens = event.message.usage.input_tokens
|
|
self._output_tokens = event.message.usage.output_tokens
|
|
if event.message.usage.cache_creation_input_tokens:
|
|
self._cache_creation_tokens = event.message.usage.cache_creation_input_tokens
|
|
if event.message.usage.cache_read_input_tokens:
|
|
self._cache_read_tokens = event.message.usage.cache_read_input_tokens
|
|
elif event.type == "message_delta":
|
|
self._output_tokens += event.usage.output_tokens
|
|
elif event.type == "content_block_start":
|
|
if event.content_block.type == "tool_use":
|
|
self._tool_call_id = event.content_block.id
|
|
self._fnc_name = event.content_block.name
|
|
self._fnc_raw_arguments = ""
|
|
elif event.type == "content_block_delta":
|
|
delta = event.delta
|
|
if delta.type == "text_delta":
|
|
text = delta.text
|
|
|
|
if self._tools is not None:
|
|
# anthropic may inject COC when using functions
|
|
if text.startswith("<thinking>"):
|
|
self._ignoring_cot = True
|
|
elif self._ignoring_cot and "</thinking>" in text:
|
|
text = text.split("</thinking>")[-1]
|
|
self._ignoring_cot = False
|
|
|
|
if self._ignoring_cot:
|
|
return None
|
|
|
|
return llm.ChatChunk(
|
|
id=self._request_id,
|
|
delta=llm.ChoiceDelta(content=text, role="assistant"),
|
|
)
|
|
elif delta.type == "input_json_delta":
|
|
assert self._fnc_raw_arguments is not None
|
|
self._fnc_raw_arguments += delta.partial_json
|
|
|
|
elif event.type == "content_block_stop":
|
|
if self._tool_call_id is not None:
|
|
assert self._fnc_name is not None
|
|
assert self._fnc_raw_arguments is not None
|
|
|
|
chat_chunk = llm.ChatChunk(
|
|
id=self._request_id,
|
|
delta=llm.ChoiceDelta(
|
|
role="assistant",
|
|
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_raw_arguments = self._fnc_name = None
|
|
return chat_chunk
|
|
|
|
return None
|