diff --git a/Dockerfile b/Dockerfile index 6bab3d5ad71..713d7f8f5ab 100644 --- a/Dockerfile +++ b/Dockerfile @@ -121,8 +121,8 @@ RUN sed -i 's/\r$//' docker/entrypoint.sh && chmod +x docker/entrypoint.sh && \ FROM $LITELLM_BUILD_IMAGE AS liteadmin-builder COPY --from=uvbin /uv /usr/local/bin/uv RUN apk add --no-cache python-3.13 -ADD --checksum=sha256:2f7ae5cdd9d91731c0990e74a58239dc3e3fd2bf28dab23b55eafcdc47aaf87e \ - https://github.com/BerriAI/litellm-admin-agent/archive/ef501e94bc9fbacb9233b922abf71427f030408c.tar.gz /tmp/liteadmin.tar.gz +ADD --checksum=sha256:4fb1847e28e2be30ba76dac94f02c4903291326516a0f4ba772f6ffb8b25d05c \ + https://github.com/BerriAI/litellm-admin-agent/archive/5661c805ce479a0cf11b842ea133cf8d22cd7e34.tar.gz /tmp/liteadmin.tar.gz RUN mkdir /tmp/liteadmin && tar xzf /tmp/liteadmin.tar.gz --strip-components=1 -C /tmp/liteadmin && \ uv venv /opt/liteadmin --python python3.13 && \ uv pip install --python /opt/liteadmin/bin/python --require-hashes -r /tmp/liteadmin/requirements.txt && \ diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a67c1bb7df8..3130d49a5e2 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -972,3 +972,9 @@ def _get_openai_compatible_provider_info( if dynamic_api_key is None and api_key is not None: dynamic_api_key = api_key return model, custom_llm_provider, dynamic_api_key, api_base + + +def get_provider_affinity_header(api_base: str | None) -> str | None: + from litellm.llms.opencode.transformation import OpenCodeProviderAdapter + + return OpenCodeProviderAdapter().get_provider_affinity_header(api_base) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index 02c4cd69159..e022839c5c7 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -44,10 +44,10 @@ def _get_value(value: object, key: str) -> object | None: return getattr(value, key, None) -def _get_provider_affinity_header_name(litellm_params: object | None) -> str | None: +def _get_provider_affinity_header_name(litellm_params: object | None, default_header: str | None = None) -> str | None: header: Final = _get_value(litellm_params, "provider_affinity_header") if litellm_params is not None else None if header is None: - return None + return default_header if not isinstance(header, str): raise TypeError("provider_affinity_header must be a string") return validate_provider_affinity_header_name(header) @@ -84,9 +84,9 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: def add_provider_affinity_header( - headers: Mapping[str, object], litellm_params: object | None + headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers - header_name: Final = _get_provider_affinity_header_name(litellm_params) + header_name: Final = _get_provider_affinity_header_name(litellm_params, default_header) if header_name is None or any(key.lower() == header_name.lower() for key in headers): return dict(headers) diff --git a/litellm/llms/aws_polly/text_to_speech/transformation.py b/litellm/llms/aws_polly/text_to_speech/transformation.py index 01fec02c310..325c9fd364d 100644 --- a/litellm/llms/aws_polly/text_to_speech/transformation.py +++ b/litellm/llms/aws_polly/text_to_speech/transformation.py @@ -178,7 +178,7 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): else: mapped_params["output_format"] = self.DEFAULT_OUTPUT_FORMAT - # Extract engine from model name (e.g., "aws_polly/neural" -> "neural") + # model arrives without the "aws_polly/" prefix, e.g. "standard" engine: Final = self._extract_engine_from_model(model) mapped_params["engine"] = engine @@ -196,18 +196,17 @@ class AWSPollyTextToSpeechConfig(BaseTextToSpeechConfig, BaseAWSLLM): """ Extract engine from model name. + speech() strips the "aws_polly/" prefix before the config is called, + so the bare engine name is what normally arrives here. + Examples: - - aws_polly/neural -> neural - - aws_polly/standard -> standard + - standard -> standard - aws_polly/long-form -> long-form - aws_polly -> neural (default) """ - if "/" in model: - parts: Final = model.split("/") - if len(parts) >= 2: - engine: Final = parts[1].lower() - if engine in self.VALID_ENGINES: - return engine + engine: Final = model.rsplit("/", 1)[-1].lower() + if engine in self.VALID_ENGINES: + return engine return self.DEFAULT_ENGINE def validate_environment( diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a7ca0d5f49e..a692f3f964e 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -52,7 +52,9 @@ from litellm.litellm_core_utils.audio_utils.subtitle_utils import ( from litellm.litellm_core_utils.core_helpers import set_provider_response_headers_in_hidden_params from litellm.litellm_core_utils.dot_notation_indexing import delete_nested_value from litellm.litellm_core_utils.get_litellm_params import AWS_CREDENTIAL_KWARGS_KEYS +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.llm_request_utils import serialize_multipart_form_fields +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.litellm_core_utils.realtime_errors import ( close_after_upstream_handshake_refusal, realtime_error_event, @@ -60,6 +62,7 @@ from litellm.litellm_core_utils.realtime_errors import ( ) from litellm.litellm_core_utils.realtime_streaming import RealTimeStreaming from litellm.litellm_core_utils.url_utils import encode_url_path_segment, join_api_base +from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, ) @@ -2435,13 +2438,13 @@ class BaseLLMHTTPHandler: # Also check for extra_headers in kwargs (from config or direct calls) extra_headers_from_kwargs: Final = kwargs.get("extra_headers", None) # Merge all header sources: forwarded < extra_headers < provider_specific - merged_headers: Final = {} - if forwarded_headers: - merged_headers.update(forwarded_headers) - if extra_headers_from_kwargs: - merged_headers.update(extra_headers_from_kwargs) - if provider_specific_headers: - merged_headers.update(provider_specific_headers) + merged_headers: Final[Mapping[str, object]] = MappingProxyType( + { + **(forwarded_headers or {}), + **(extra_headers_from_kwargs or {}), + **(provider_specific_headers or {}), + } + ) request_params: Final = reduce( delete_nested_value, litellm_params.get("additional_drop_params") or (), @@ -2452,11 +2455,19 @@ class BaseLLMHTTPHandler: and anthropic_messages_provider_config.should_filter_anthropic_beta_headers() ): await aload_beta_headers_config() + affinity_api_base: Final = ( + AnthropicModelInfo.get_api_base(api_base) if custom_llm_provider == "anthropic" else api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) ( headers, api_base, ) = await anthropic_messages_provider_config.avalidate_anthropic_messages_environment( - headers=merged_headers or {}, + headers=( + add_provider_affinity_header(merged_headers, litellm_params, default_affinity_header) + if default_affinity_header is not None + else dict(merged_headers) + ), model=model, messages=messages, optional_params=request_params, diff --git a/litellm/llms/opencode/common_utils.py b/litellm/llms/opencode/common_utils.py new file mode 100644 index 00000000000..65a96411e52 --- /dev/null +++ b/litellm/llms/opencode/common_utils.py @@ -0,0 +1,10 @@ +import re + + +def is_opencode_api_base(api_base: str | None) -> bool: + return bool( + api_base + and re.fullmatch( + r"https://opencode\.ai/(?:zen(?:/go)?|inference/(?:go/)?(?:openai|anthropic))(?:/v1)?/?", api_base + ) + ) diff --git a/litellm/llms/opencode/transformation.py b/litellm/llms/opencode/transformation.py new file mode 100644 index 00000000000..9b62e7aca95 --- /dev/null +++ b/litellm/llms/opencode/transformation.py @@ -0,0 +1,6 @@ +from litellm.llms.opencode.common_utils import is_opencode_api_base + + +class OpenCodeProviderAdapter: + def get_provider_affinity_header(self, api_base: str | None) -> str | None: + return "x-opencode-session" if is_opencode_api_base(api_base) else None diff --git a/litellm/main.py b/litellm/main.py index 32a7c290bb8..327a7a1c4b5 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -89,6 +89,7 @@ from litellm.litellm_core_utils.get_litellm_params import ( with_control_options, ) from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_provider_affinity_header, invalid_provider_routing_hints_message, ) from litellm.litellm_core_utils.get_provider_specific_headers import ( @@ -130,6 +131,7 @@ from litellm.llms.bedrock_mantle.chat.claude_transformation import bedrock_mantl from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler, http2_enabled from litellm.llms.openai.chat.gpt_5_transformation import OpenAIGPT5Config +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.vertex_ai.common_utils import ( VertexAIModelRoute, @@ -5990,19 +5992,28 @@ def completion( }, ) litellm_params: Final = with_control_options(requested_litellm_params, control_options) - if litellm_params.get("provider_affinity_header") is not None: + affinity_api_base: Final = ( + str(client.base_url) + if isinstance(client, openai.OpenAI | openai.AsyncOpenAI) + else get_openai_credentials(api_base=api_base).api_base + if custom_llm_provider == "openai" + else api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) + if litellm_params.get("provider_affinity_header") is not None or default_affinity_header is not None: try: headers = add_provider_affinity_header( headers=headers or litellm.headers or MappingProxyType({}), litellm_params=MappingProxyType( { - "provider_affinity_header": litellm_params["provider_affinity_header"], + "provider_affinity_header": litellm_params.get("provider_affinity_header"), "litellm_session_id": kwargs.get("litellm_session_id"), "session_id": kwargs.get("session_id"), "metadata": metadata, "litellm_metadata": kwargs.get("litellm_metadata"), } ), + default_header=default_affinity_header, ) except ValueError as affinity_error: raise litellm.BadRequestError( diff --git a/litellm/messages/dispatch.py b/litellm/messages/dispatch.py index cdb9a925403..2e9d6a77a62 100644 --- a/litellm/messages/dispatch.py +++ b/litellm/messages/dispatch.py @@ -103,6 +103,7 @@ def _target(request: NativeCall, model: str, provider: str | None) -> Target: _FEATURES: Final[detectors.Detectors] = MappingProxyType( { **detectors.COMMON, + Feature.PROVIDER_AFFINITY: route_detectors.provider_affinity, Feature.SAFEGUARDS_CLASSIFIER: route_detectors.safeguards_classifier, Feature.RESPONSE_CACHE: detectors.response_cache, Feature.MODIFY_PARAMS: detectors.modify_params, diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 09eea7cb323..0b526adf4fa 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -30,6 +30,7 @@ from litellm.integrations.anthropic_cache_control_hook import ( from litellm.litellm_core_utils.asyncify import run_async_function from litellm.litellm_core_utils.completion_timeout import CompletionTimeout, StreamKind from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.get_llm_provider_logic import get_provider_affinity_header from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -39,6 +40,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.llms.openai.common_utils import get_openai_credentials from litellm.llms.openai_like.responses.transformation import OpenAILikeResponsesConfig from litellm.responses.litellm_completion_transformation.compaction_items import ( with_compaction_items_as_summaries, @@ -1466,6 +1468,12 @@ def responses( # get llm provider logic litellm_params: Final = GenericLiteLLMParams(**kwargs) + affinity_api_base: Final = ( + get_openai_credentials(api_base=litellm_params.api_base).api_base + if custom_llm_provider == "openai" + else litellm_params.api_base + ) + default_affinity_header: Final = get_provider_affinity_header(affinity_api_base) try: effective_extra_headers: Final = ( add_provider_affinity_header( @@ -1479,8 +1487,9 @@ def responses( "litellm_metadata": kwargs.get("litellm_metadata"), } ), + default_header=default_affinity_header, ) - if litellm_params.provider_affinity_header is not None + if litellm_params.provider_affinity_header is not None or default_affinity_header is not None else extra_headers ) except ValueError as affinity_error: diff --git a/litellm/rust_bridge/catalog.py b/litellm/rust_bridge/catalog.py index ebb221e98be..74cd28161a4 100644 --- a/litellm/rust_bridge/catalog.py +++ b/litellm/rust_bridge/catalog.py @@ -165,6 +165,7 @@ CHAT_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { } MESSAGES_PYTHON_FEATURES: Final = COMMON_PYTHON_FEATURES | { + Feature.PROVIDER_AFFINITY, Feature.SAFEGUARDS_CLASSIFIER, Feature.PRE_REQUEST_HOOK, Feature.WEBSEARCH_INTERCEPTION, diff --git a/litellm/rust_bridge/messages/detectors.py b/litellm/rust_bridge/messages/detectors.py index 691c04f2daa..b53e80805a5 100644 --- a/litellm/rust_bridge/messages/detectors.py +++ b/litellm/rust_bridge/messages/detectors.py @@ -5,10 +5,15 @@ from __future__ import annotations from typing import Final import litellm -from litellm.rust_bridge.detectors import overridden, present, streaming +from litellm.llms.opencode.common_utils import is_opencode_api_base +from litellm.rust_bridge.detectors import api_base, overridden, present, streaming from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str +def provider_affinity(request: NativeCall) -> bool: + return is_opencode_api_base(api_base(request)) + + def anthropic_workload_identity(request: NativeCall) -> bool: """Python mints an Anthropic federation token when no static credential is configured.""" from litellm.llms.anthropic.common_utils import AnthropicModelInfo diff --git a/litellm/rust_bridge/responses/detectors.py b/litellm/rust_bridge/responses/detectors.py index 392c11f48a8..763b63d6742 100644 --- a/litellm/rust_bridge/responses/detectors.py +++ b/litellm/rust_bridge/responses/detectors.py @@ -9,6 +9,7 @@ from typing import TYPE_CHECKING, Final, cast # noqa: TID251 # tool lists stay from pydantic import TypeAdapter import litellm +from litellm.llms.opencode.common_utils import is_opencode_api_base from litellm.rust_bridge.detectors import api_base, call_callbacks, overridden, present from litellm.rust_bridge.public_call import NativeCall, optional_mapping, optional_sequence, optional_str @@ -25,7 +26,7 @@ def client_headers(request: NativeCall) -> bool: def provider_affinity(request: NativeCall) -> bool: """Python validates and fills the affinity header from the session ID.""" - return present(request, "provider_affinity_header") + return present(request, "provider_affinity_header") or is_opencode_api_base(api_base(request)) def text_format(request: NativeCall) -> bool: diff --git a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py index d7474936425..57e66805678 100644 --- a/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py +++ b/tests/unit/litellm_core_utils/test_get_llm_provider_logic.py @@ -8,6 +8,7 @@ from litellm import CustomLLM from litellm.litellm_core_utils import get_llm_provider_logic from litellm.litellm_core_utils.get_llm_provider_logic import ( get_llm_provider, + get_provider_affinity_header, inferred_provider, is_registered_custom_provider, ) @@ -706,3 +707,11 @@ class TestClaudeModelPatternMatching: set_fallback_generalizations([]) with pytest.raises(litellm.BadRequestError): litellm.get_llm_provider(model="claude-opus-4-9") + + +@pytest.mark.parametrize("prefix", ["openai", "anthropic", "custom_openai", "openai_like"]) +def test_opencode_affinity_resolution_preserves_existing_transport(prefix: str) -> None: + api_base: Final = "https://opencode.ai/zen/go/v1" + model, provider, _, destination = get_llm_provider(model=f"{prefix}/test-model", api_base=api_base) + assert (model, provider, destination) == ("test-model", prefix, api_base) + assert get_provider_affinity_header(destination) == "x-opencode-session" diff --git a/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py b/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py index 3826b22af34..ace82e517ad 100644 --- a/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py +++ b/tests/unit/llms/aws_polly/text_to_speech/test_aws_polly_text_to_speech_transformation.py @@ -1,6 +1,10 @@ -import litellm +import json +from unittest.mock import MagicMock, patch + +import httpx import pytest -from unittest.mock import MagicMock + +import litellm @pytest.mark.asyncio @@ -134,6 +138,34 @@ async def test_aws_polly_tts_with_native_voice(): assert request_body.get("TextType", "text") == "text" +@pytest.mark.asyncio +@pytest.mark.usefixtures("fake_provider_credentials") +@pytest.mark.parametrize("engine", ["standard", "neural", "long-form", "generative"]) +async def test_aws_polly_tts_sends_engine_from_model(engine: str) -> None: + """ + The engine in the model name must reach Polly. speech() strips the + provider prefix before the config reads it, which used to turn every + engine into the neural default. + """ + mock_httpx_response = MagicMock(spec=httpx.Response) + mock_httpx_response.content = b"fake_audio_data" + mock_httpx_response.status_code = 200 + mock_httpx_response.headers = {"content-type": "audio/mpeg"} + + with patch("litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post") as mock_post: + mock_post.return_value = mock_httpx_response + + await litellm.aspeech( + model=f"aws_polly/{engine}", + voice="Joanna", + input="Hello", + aws_region_name="us-east-1", + ) + + request_body = json.loads(mock_post.call_args.kwargs["data"]) + assert request_body["Engine"] == engine + + @pytest.mark.asyncio @pytest.mark.usefixtures("fake_provider_credentials") async def test_aws_polly_tts_with_openai_voice_mapping(): diff --git a/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py b/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py index 4778b6296fc..5f195b7a478 100644 --- a/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py +++ b/tests/unit/llms/aws_polly/text_to_speech/test_transformation.py @@ -1,6 +1,8 @@ import json from typing import Final +import pytest + from litellm.llms.aws_polly.text_to_speech.transformation import AWSPollyTextToSpeechConfig @@ -17,3 +19,17 @@ def test_installed_botocore_signs_the_speech_request() -> None: authorization: Final = headers["Authorization"] assert "/us-west-2/polly/aws4_request" in authorization assert json.loads(body) == {"Text": "ping", "VoiceId": "Joanna"} + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("standard", "standard"), + ("Long-Form", "long-form"), + ("aws_polly/generative", "generative"), + ("aws_polly", "neural"), + ("unknown", "neural"), + ], +) +def test_extract_engine_from_model(model: str, expected: str) -> None: + assert AWSPollyTextToSpeechConfig()._extract_engine_from_model(model) == expected diff --git a/tests/unit/llms/custom_httpx/test_llm_http_handler.py b/tests/unit/llms/custom_httpx/test_llm_http_handler.py index 63cc45b29df..8b32ab2e6b8 100644 --- a/tests/unit/llms/custom_httpx/test_llm_http_handler.py +++ b/tests/unit/llms/custom_httpx/test_llm_http_handler.py @@ -5542,3 +5542,85 @@ async def test_anthropic_messages_dropped_safeguards_send_neither_the_field_nor_ sent: Final = received[0] assert "safeguards" not in json.loads(sent.content) assert "dangerous-tool-use-2026-09-03" not in sent.headers.get("anthropic-beta", ""), dict(sent.headers) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("environment_base", [None, "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"]) +@pytest.mark.parametrize( + ("api_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/inference/go/anthropic", None, "conversation-123"), + ("https://opencode.ai/inference/go/anthropic", "caller-override", "caller-override"), + ("https://other.example", None, None), + ], +) +async def test_messages_opencode_session_header( + stream: bool, + environment_base: str | None, + api_base: str, + explicit_header: str | None, + expected_session: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.delenv("ANTHROPIC_API_BASE", raising=False) + monkeypatch.delenv("ANTHROPIC_BASE_URL", raising=False) + if environment_base: + monkeypatch.setenv(environment_base, "https://opencode.ai/inference/go/anthropic") + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": "OK"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, + } + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"event: message_start\ndata: {json.dumps({'type': 'message_start', 'message': payload})}\n\n" + 'event: message_stop\ndata: {"type":"message_stop"}\n\n', + ) + if stream + else httpx.Response(200, json=payload) + ) + + upstream: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + result: Final = await litellm.anthropic.messages.acreate( + model="anthropic/claude-haiku-5-5", + messages=[{"role": "user", "content": "hello"}], + max_tokens=128000, + api_base=None if environment_base and api_base == "https://opencode.ai/inference/go/anthropic" else api_base, + api_key="test-key", + stream=stream, + litellm_session_id="conversation-123", + client=upstream, + extra_headers={"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + ) + if stream: + chunks: Final = b"".join([chunk async for chunk in result]) + assert b'"text": "OK"' in chunks + assert b"message_stop" in chunks + else: + assert result["content"] == [{"type": "text", "text": "OK"}] + finally: + await upstream.close() + sent: Final = received.get_nowait() + assert str(sent.url) == f"{api_base}/v1/messages" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["x-api-key"] == "test-key" + assert json.loads(sent.content) == { + "model": "claude-haiku-5-5", + "messages": [{"role": "user", "content": "hello"}], + "max_tokens": 128000, + "stream": stream, + } + assert received.empty() diff --git a/tests/unit/llms/opencode/test_transformation.py b/tests/unit/llms/opencode/test_transformation.py new file mode 100644 index 00000000000..96203c23122 --- /dev/null +++ b/tests/unit/llms/opencode/test_transformation.py @@ -0,0 +1,90 @@ +from typing import Final + +import pytest + +from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY +from litellm.litellm_core_utils.provider_affinity import add_provider_affinity_header +from litellm.llms.opencode.transformation import OpenCodeProviderAdapter + + +@pytest.mark.parametrize( + "api_base", + [ + "https://opencode.ai/zen/v1", + "https://opencode.ai/zen/go/v1/", + "https://opencode.ai/inference/openai/v1", + "https://opencode.ai/inference/go/openai/v1", + "https://opencode.ai/inference/go/anthropic/v1", + "https://opencode.ai/inference/anthropic/v1", + ], +) +def test_opencode_defaults_affinity_header_to_the_callers_session(api_base: str) -> None: + headers: Final = {"Authorization": "Bearer test-key"} + params: Final = {"metadata": {"session_id": "conversation-123"}} + assert add_provider_affinity_header( + headers, params, OpenCodeProviderAdapter().get_provider_affinity_header(api_base) + ) == { + **headers, + "x-opencode-session": "conversation-123", + } + assert headers == {"Authorization": "Bearer test-key"} + assert params == {"metadata": {"session_id": "conversation-123"}} + + +@pytest.mark.parametrize( + "api_base", + [ + None, + "https://api.openai.com/v1", + "https://opencode.ai.example/zen/go/v1", + "https://opencode.ai@example.com/zen/go/v1", + "https://example.com/opencode.ai/zen/go/v1", + "https://example.com/v1?upstream=https://opencode.ai/zen/go/v1", + "https://opencode.ai/console/v1", + "http://opencode.ai/zen/go/v1", + ], +) +def test_other_destinations_do_not_receive_an_automatic_session_header(api_base: str | None) -> None: + headers: Final = {"Authorization": "Bearer test-key", "x-existing": "unchanged"} + assert ( + add_provider_affinity_header( + headers, + {"litellm_session_id": "conversation-123"}, + OpenCodeProviderAdapter().get_provider_affinity_header(api_base), + ) + == headers + ) + + +@pytest.mark.parametrize( + "params", + [ + {}, + {"litellm_trace_id": "request-123"}, + {"metadata": {"trace_id": "request-123"}}, + { + "litellm_session_id": "request-123", + "litellm_metadata": {"session_id": "request-123", SESSION_ID_GENERATED_METADATA_KEY: True}, + }, + ], +) +def test_opencode_does_not_invent_a_conversation_id(params: dict[str, object]) -> None: + assert ( + add_provider_affinity_header( + {}, params, OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1") + ) + == {} + ) + + +def test_opencode_preserves_explicit_header_and_affinity_configuration() -> None: + assert add_provider_affinity_header( + {"X-OpenCode-Session": "explicit"}, + {"litellm_session_id": "conversation-123"}, + OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1"), + ) == {"X-OpenCode-Session": "explicit"} + assert add_provider_affinity_header( + {}, + {"litellm_session_id": "conversation-123", "provider_affinity_header": "x-custom-session"}, + OpenCodeProviderAdapter().get_provider_affinity_header("https://opencode.ai/zen/go/v1"), + ) == {"x-custom-session": "conversation-123"} diff --git a/tests/unit/responses/test_responses_api_request_body.py b/tests/unit/responses/test_responses_api_request_body.py index 436367bd82d..36b5d8f135b 100644 --- a/tests/unit/responses/test_responses_api_request_body.py +++ b/tests/unit/responses/test_responses_api_request_body.py @@ -4,6 +4,7 @@ over the wire and surface provider errors correctly. Expected JSON bodies are st in expected_responses_api_request/. """ +import asyncio import copy import json import logging @@ -1231,3 +1232,110 @@ def test_faked_stream_accepts_a_text_timeout(monkeypatch: pytest.MonkeyPatch) -> assert [event.response.output[0].content[0].text for event in response if event.type == "response.completed"] == [ "Done." ] + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("api_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/inference/go/openai/v1", None, "conversation-123"), + ("https://opencode.ai/inference/go/openai/v1", "caller-override", "caller-override"), + ("https://other.example/v1", None, None), + ], +) +@pytest.mark.asyncio +async def test_responses_opencode_session_header( + asynchronous: bool, stream: bool, api_base: str, explicit_header: str | None, expected_session: str | None +) -> None: + received: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + payload: Final = _minimal_responses_api_payload("resp_test", "gpt-6-luna") + + def respond(request: httpx.Request) -> httpx.Response: + received.put_nowait(request) + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text="".join( + f"event: {kind}\ndata: {json.dumps({'type': kind, 'response': payload})}\n\n" + for kind in ("response.created", "response.completed") + ), + ) + if stream + else httpx.Response(200, json=payload) + ) + + args: Final = { + "model": "openai/gpt-6-luna", + "input": "hello", + "api_base": api_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + "stream": stream, + "extra_headers": {"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + } + if asynchronous: + async_client: Final = AsyncHTTPHandler(transport=httpx.MockTransport(respond)) + try: + response: Final = await litellm.aresponses(**args, client=async_client) + if stream: + events: Final = tuple([event.type async for event in response]) + assert "response.completed" in events + else: + assert response.output[0].content[0].text == "Done." + finally: + await async_client.close() + else: + sync_client: Final = HTTPHandler(client=httpx.Client(transport=httpx.MockTransport(respond))) + try: + result: Final = litellm.responses(**args, client=sync_client) + if stream: + assert "response.completed" in tuple(event.type for event in result) + else: + assert result.output[0].content[0].text == "Done." + finally: + sync_client.close() + sent: Final = received.get_nowait() + assert str(sent.url) == f"{api_base}/responses" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "gpt-6-luna", "input": "hello", "stream": stream} + assert received.empty() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source", ["OPENAI_BASE_URL", "OPENAI_API_BASE", "litellm.api_base"]) +@pytest.mark.parametrize("explicit_base", [None, "https://other.example/v1"]) +@pytest.mark.asyncio +async def test_responses_opencode_session_uses_effective_destination( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + asynchronous: bool, + source: str, + explicit_base: str | None, +) -> None: + default_base: Final = "https://opencode.ai/inference/go/openai/v1" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", default_base if source == "litellm.api_base" else None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if source != "litellm.api_base": + monkeypatch.setenv(source, default_base) + endpoint: Final = respx_mock.post(f"{explicit_base or default_base}/responses").respond( + json=_minimal_responses_api_payload("resp_test", "gpt-6-luna") + ) + args: Final = { + "model": "openai/gpt-6-luna", + "input": "hello", + "api_base": explicit_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + } + response: Final = await litellm.aresponses(**args) if asynchronous else litellm.responses(**args) + assert response.output[0].content[0].text == "Done." + assert endpoint.call_count == 1 + sent: Final = endpoint.calls[0].request + assert sent.headers.get("x-opencode-session") == (None if explicit_base else "conversation-123") + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "gpt-6-luna", "input": "hello"} diff --git a/tests/unit/rust_bridge/test_detectors.py b/tests/unit/rust_bridge/test_detectors.py index 726128386d5..78722458cf3 100644 --- a/tests/unit/rust_bridge/test_detectors.py +++ b/tests/unit/rust_bridge/test_detectors.py @@ -116,6 +116,9 @@ def _wrap_nested_level(value: object, depth: int) -> object: (responses_detectors.provider_affinity, {"provider_affinity_header": ""}, True), (responses_detectors.provider_affinity, {"provider_affinity_header": None, "session_id": "session-1"}, False), (responses_detectors.provider_affinity, {}, False), + (responses_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/openai/v1"}, True), + (messages_detectors.provider_affinity, {"api_base": "https://opencode.ai/inference/go/anthropic/v1"}, True), + (messages_detectors.provider_affinity, {"api_base": "https://other.example/v1"}, False), (detectors.streaming, {"stream": True}, True), (detectors.streaming, {"stream": False}, False), (detectors.timeout_variants, {"request_timeout": 5}, True), diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index c3d3575b8b3..70b268bf2a2 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -20,7 +20,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import httpx import pytest import respx -from openai import APITimeoutError, AsyncOpenAI +from openai import APITimeoutError, AsyncOpenAI, OpenAI from openai._legacy_response import LegacyAPIResponse from openai.types.chat.chat_completion import ChatCompletion from pydantic import TypeAdapter @@ -95,6 +95,142 @@ def preserve_litellm_completion_state(monkeypatch: pytest.MonkeyPatch) -> None: WHITE_PNG: Final = (Path(__file__).parents[1] / "white_100x100.png").read_bytes() +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize( + ("client_base", "explicit_header", "expected_session"), + [ + ("https://opencode.ai/zen/go/v1", None, "conversation-123"), + ("https://opencode.ai/zen/go/v1", "caller-override", "caller-override"), + ("https://other.example/v1", None, None), + ], +) +@pytest.mark.asyncio +async def test_completion_sends_session_only_to_the_actual_opencode_destination( + asynchronous: bool, + stream: bool, + client_base: str, + explicit_header: str | None, + expected_session: str | None, + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://opencode.ai/zen/go/v1") + monkeypatch.setenv("OPENAI_API_BASE", "https://opencode.ai/zen/go/v1") + monkeypatch.setattr(litellm, "api_base", "https://opencode.ai/zen/go/v1") + requests: Final[asyncio.Queue[httpx.Request]] = asyncio.Queue() + + def respond(request: httpx.Request) -> httpx.Response: + requests.put_nowait(request) + payload: Final = { + "id": "chatcmpl-test", + "object": "chat.completion.chunk" if stream else "chat.completion", + "created": 1, + "model": "session-test", + "choices": [ + { + "index": 0, + "delta" if stream else "message": {"role": "assistant", "content": "OK"}, + "finish_reason": "stop", + } + ], + } + return ( + httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + text=f"data: {json.dumps(payload)}\n\ndata: [DONE]\n\n", + ) + if stream + else httpx.Response(200, json=payload) + ) + + args: Final = { + "model": "openai/session-test", + "api_base": "https://opencode.ai/zen/go/v1", + "messages": [{"role": "user", "content": "hello"}], + "litellm_session_id": "conversation-123", + "extra_headers": {"X-OpenCode-Session": explicit_header} if explicit_header else {"x-existing": "unchanged"}, + "stream": stream, + } + if asynchronous: + async with AsyncOpenAI( + api_key="test-key", + base_url=client_base, + http_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)), + ) as async_client: + response: Final = await acompletion(**args, client=async_client) + content: Final = ( + "".join([chunk.choices[0].delta.content or "" async for chunk in response]) + if stream + else response.choices[0].message.content + ) + assert content == "OK" + else: + with OpenAI( + api_key="test-key", base_url=client_base, http_client=httpx.Client(transport=httpx.MockTransport(respond)) + ) as sync_client: + result: Final = completion(**args, client=sync_client) + text: Final = ( + "".join(chunk.choices[0].delta.content or "" for chunk in result) + if stream + else result.choices[0].message.content + ) + assert text == "OK" + sent: Final = requests.get_nowait() + assert str(sent.url) == f"{client_base}/chat/completions" + assert sent.headers.get("x-opencode-session") == expected_session + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == { + "model": "session-test", + "messages": args["messages"], + **({"stream": True} if stream else {}), + } + assert requests.empty() + + +@pytest.mark.parametrize("asynchronous", [False, True]) +@pytest.mark.parametrize("source", ["OPENAI_BASE_URL", "OPENAI_API_BASE", "litellm.api_base"]) +@pytest.mark.parametrize("explicit_base", [None, "https://other.example/v1"]) +@pytest.mark.asyncio +async def test_completion_opencode_session_uses_effective_destination( + monkeypatch: pytest.MonkeyPatch, + respx_mock: respx.MockRouter, + asynchronous: bool, + source: str, + explicit_base: str | None, +) -> None: + default_base: Final = "https://opencode.ai/inference/go/openai/v1" + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", default_base if source == "litellm.api_base" else None) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + if source != "litellm.api_base": + monkeypatch.setenv(source, default_base) + endpoint: Final = respx_mock.post(f"{explicit_base or default_base}/chat/completions").respond( + json={ + "id": "chatcmpl-test", + "object": "chat.completion", + "created": 1, + "model": "session-test", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "OK"}, "finish_reason": "stop"}], + } + ) + args: Final = { + "model": "openai/session-test", + "messages": [{"role": "user", "content": "hello"}], + "api_base": explicit_base, + "api_key": "test-key", + "litellm_session_id": "conversation-123", + } + response: Final = await litellm.acompletion(**args) if asynchronous else litellm.completion(**args) + assert response.choices[0].message.content == "OK" + assert endpoint.call_count == 1 + sent: Final = endpoint.calls[0].request + assert sent.headers.get("x-opencode-session") == (None if explicit_base else "conversation-123") + assert sent.headers["authorization"] == "Bearer test-key" + assert json.loads(sent.content) == {"model": "session-test", "messages": args["messages"]} + + @pytest.fixture def openai_api_response(): mock_response_data = { diff --git a/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx index 4f34fe22e58..f7bbf59c0cb 100644 --- a/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.images.integration.test.tsx @@ -56,7 +56,7 @@ function renderChat() { async function uploadImage(user: ReturnType, name = "diagram.png") { const bytes = Uint8Array.from(atob(pixel), (character) => character.charCodeAt(0)); - await user.upload(screen.getByLabelText("Upload images"), new File([bytes], name, { type: "image/png" })); + await user.upload(screen.getByLabelText("Upload files"), new File([bytes], name, { type: "image/png" })); await screen.findByRole("img", { name }); } @@ -117,13 +117,67 @@ describe("chat image workflows", () => { it("should render the permitted models and image controls for a research team member", async () => { renderChat(); expect(await screen.findByRole("button", { name: "research-chat" })).toBeInTheDocument(); - expect(screen.getByRole("button", { name: "Attach images" })).toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Attach files" })).toBeInTheDocument(); await userEvent.setup().click(screen.getByRole("button", { name: "research-chat", exact: true })); expect(await screen.findByRole("button", { name: /^studio-image/ })).toBeInTheDocument(); expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); expect(recordChatQuery).not.toHaveBeenCalled(); }); + it("should upload a document as a file part, preserve it on edit, and keep contents out of metrics", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + await user.upload( + screen.getByLabelText("Upload files"), + new File(["private fixture"], "brief.docx", { + type: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + }), + ); + expect(await screen.findByText("brief.docx")).toBeInTheDocument(); + expect(screen.queryByRole("img", { name: "brief.docx" })).not.toBeInTheDocument(); + await user.click(screen.getByRole("button", { name: "Send" })); + await screen.findByText("I can see the uploaded image."); + const filePart = { + type: "input_file", + filename: "brief.docx", + file_data: + "data:application/vnd.openxmlformats-officedocument.wordprocessingml.document;base64,cHJpdmF0ZSBmaXh0dXJl", + }; + expect(requests[0].body).toMatchObject({ input: [{ role: "user", content: [filePart] }] }); + expect(recordChatQuery).toHaveBeenCalledExactlyOnceWith("media-test-user"); + await user.click(screen.getByRole("button", { name: "Edit message" })); + act(() => + fireEvent.change(screen.getByRole("textbox", { name: "Edit message" }), { + target: { value: "Summarize the document" }, + }), + ); + await user.click(screen.getByRole("button", { name: "Save & Send" })); + await waitFor(() => expect(requests).toHaveLength(2)); + expect(requests[1].body).toMatchObject({ + input: [{ role: "user", content: [{ type: "input_text", text: "Summarize the document" }, filePart] }], + }); + expect(requests[1].body).not.toHaveProperty("previous_response_id"); + }); + + it("should enforce the shared file limit and remove a document before submission", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "research-chat" }); + const files = Array.from( + { length: 5 }, + (_, index) => new File(["text"], `note-${index}.txt`, { type: "text/plain" }), + ); + await user.upload(screen.getByLabelText("Upload files"), files); + expect(await screen.findByRole("alert")).toHaveTextContent("Attach up to 4 files per message."); + expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); + await user.upload(screen.getByLabelText("Upload files"), files[0]); + await user.click(await screen.findByRole("button", { name: "Remove note-0.txt" })); + expect(screen.queryByText("note-0.txt")).not.toBeInTheDocument(); + expect(screen.getByRole("button", { name: "Send" })).toBeDisabled(); + expect(recordChatQuery).not.toHaveBeenCalled(); + }); + it("should restore the selected chat model after switching modes and reopening the page", async () => { const user = userEvent.setup(); const view = renderChat(); diff --git a/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx index ab52bd55c3f..a8521bdaaf7 100644 --- a/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.integration.test.tsx @@ -214,7 +214,7 @@ describe("/ui/chat storage banner", () => { it("keeps the dismiss control amber on hover instead of the ghost variant's foreground", async () => { renderWithProviders(); - const banner = await screen.findByText("Some chat history or images could not be saved or restored."); + const banner = await screen.findByText("Some chat history or attachments could not be saved or restored."); const dismiss = within(banner.parentElement!).getByRole("button"); expect(dismiss).toHaveClass("hover:text-warning/80"); expect(dismiss).not.toHaveClass("hover:text-foreground"); diff --git a/ui/litellm-dashboard/src/app/chat/page.teams.integration.test.tsx b/ui/litellm-dashboard/src/app/chat/page.teams.integration.test.tsx index bc8aed27e36..b43d98ce021 100644 --- a/ui/litellm-dashboard/src/app/chat/page.teams.integration.test.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.teams.integration.test.tsx @@ -4,10 +4,8 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, testQueryClient } from "@/../tests/test-utils"; import { ChatShellProvider } from "@/contexts/ChatShellContext"; import type { ModelGroup } from "@/components/llm_calls/fetch_models"; -import { recordChatQuery } from "@/lib/telemetry/chat"; import ChatConversationPage from "./page"; -vi.mock("@/lib/telemetry/chat", () => ({ recordChatQuery: vi.fn(async () => {}) })); vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn(), replace: vi.fn() }) })); const teams = [ @@ -40,7 +38,6 @@ function enterPrompt(text = "Hello") { beforeEach(() => { localStorage.clear(); testQueryClient.clear(); - vi.mocked(recordChatQuery).mockClear(); requests.length = 0; models = [{ model_group: "shared-chat", mode: "chat", providers: ["openai"], available_teams: teams }]; vi.stubGlobal( @@ -85,7 +82,6 @@ describe("Chat session team selection", () => { enterPrompt(); expect(screen.getByRole("button", { name: "Send", exact: true })).toBeDisabled(); expect(requests).toHaveLength(0); - expect(recordChatQuery).not.toHaveBeenCalled(); }); it.each([ @@ -122,7 +118,6 @@ describe("Chat session team selection", () => { expect(requests[1].body.input).toEqual( expect.arrayContaining([expect.objectContaining({ role: "user", content: "First question" })]), ); - expect(recordChatQuery).toHaveBeenCalledTimes(2); }); it("should send image generation to the chosen team and clear the choice when changing models", async () => { @@ -145,4 +140,38 @@ describe("Chat session team selection", () => { body: { model: "studio-image", prompt: "Draw a circle" }, }); }); + + it("should retain uploaded documents and replay them under the newly selected team", async () => { + const user = userEvent.setup(); + renderChat(); + await screen.findByRole("button", { name: "shared-chat" }); + const file = new File(["team document"], "brief.txt", { type: "text/plain" }); + await user.upload(screen.getByLabelText("Upload files"), file); + await screen.findByRole("button", { name: "Remove brief.txt" }); + enterPrompt("Read the file"); + expect(screen.getByRole("button", { name: "Send", exact: true })).toBeDisabled(); + await user.click(screen.getByRole("combobox", { name: "Team paying for this request" })); + await user.click(await screen.findByRole("option", { name: "Research" })); + await user.click(screen.getByRole("button", { name: "Send", exact: true })); + await screen.findByText("Reply 1"); + await user.click(screen.getByRole("combobox", { name: "Team paying for this request" })); + await user.click(await screen.findByRole("option", { name: "Engineering" })); + enterPrompt("Use that file again"); + await user.click(screen.getByRole("button", { name: "Send", exact: true })); + await screen.findByText("Reply 2"); + + const document = { + type: "input_file", + filename: "brief.txt", + file_data: "data:text/plain;base64,dGVhbSBkb2N1bWVudA==", + }; + const originalMessage = expect.objectContaining({ + role: "user", + content: expect.arrayContaining([document]), + }); + expect(requests.map((request) => request.teamId)).toEqual(["research", "engineering"]); + expect(requests[0].body.input).toEqual(expect.arrayContaining([originalMessage])); + expect(requests[1].body).not.toHaveProperty("previous_response_id"); + expect(requests[1].body.input).toEqual(expect.arrayContaining([originalMessage])); + }); }); diff --git a/ui/litellm-dashboard/src/app/chat/page.tsx b/ui/litellm-dashboard/src/app/chat/page.tsx index 9ead3ba3e9f..0af45cbbcf3 100644 --- a/ui/litellm-dashboard/src/app/chat/page.tsx +++ b/ui/litellm-dashboard/src/app/chat/page.tsx @@ -11,7 +11,7 @@ import { getChatRoutes } from "@/components/chat/ChatShell"; import ChatMessages from "@/components/chat/ChatMessages"; import MCPConnectPicker from "@/components/chat/MCPConnectPicker"; import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; -import type { AssistantMessageUpdate, ChatImageAttachment, ChatMessage } from "@/components/chat/types"; +import type { AssistantMessageUpdate, ChatAttachment, ChatMessage } from "@/components/chat/types"; import ChatModelPicker, { ChatTeamPicker, chatTeamSelection, @@ -22,7 +22,8 @@ import ChatModelPicker, { import MoyaiChatComposer from "@/components/chat/moyai/MoyaiChatComposer"; import MoyaiChatHome from "@/components/chat/moyai/MoyaiChatHome"; import { runChatRequest } from "@/components/chat/runChatRequest"; -import { readChatImage, MAX_CHAT_IMAGES } from "@/components/chat/imageHelpers"; +import { readChatAttachment, MAX_CHAT_ATTACHMENTS, CHAT_ATTACHMENT_ACCEPT } from "@/components/chat/attachmentHelpers"; +import ChatFileChip from "@/components/chat/ChatFileChip"; import { useQueryState } from "nuqs"; import { recordChatQuery } from "@/lib/telemetry/chat"; @@ -30,10 +31,51 @@ const LOCALSTORAGE_MODEL_KEY = "litellm_chat_selected_model"; type ReplayMessage = Pick; -function hasMessageContent(text: string, images: ChatImageAttachment[], mode: ChatMode): boolean { +function hasMessageContent(text: string, images: ChatAttachment[], mode: ChatMode): boolean { return Boolean(text.trim()) || (mode === "chat" && images.length > 0); } +function AttachmentPreview({ + attachments, + disabled, + onRemove, +}: { + attachments: ChatAttachment[]; + disabled: boolean; + onRemove: (id: string) => void; +}) { + if (!attachments.length) return null; + return ( +
+ {attachments.map((attachment) => ( +
+ {attachment.kind === "file" ? ( + + ) : ( + <> + {/* User attachments use local data URLs, not the image optimizer. */} + {/* eslint-disable-next-line @next/next/no-img-element */} + {attachment.name} +

{attachment.name}

+ + )} + +
+ ))} +
+ ); +} + export default function ChatConversationPage() { const [, setConversationIdInUrl] = useQueryState("id", { history: "push" }); const router = useRouter(); @@ -57,11 +99,11 @@ export default function ChatConversationPage() { const [selectedTeamId, setSelectedTeamId] = useState(null); const [models, setModels] = useState([]); const [mode, setMode] = useState("chat"); - const [attachments, setAttachments] = useState([]); + const [attachments, setAttachments] = useState([]); const [attachmentError, setAttachmentError] = useState(null); - const [isReadingImages, setIsReadingImages] = useState(false); - const imageInputRef = useRef(null); - const readImagesRef = useRef(false); + const [isReadingAttachments, setIsReadingAttachments] = useState(false); + const attachmentInputRef = useRef(null); + const readAttachmentsRef = useRef(false); const draftVersionRef = useRef(0); const [isLoadingModels, setIsLoadingModels] = useState(true); @@ -141,7 +183,7 @@ export default function ChatConversationPage() { const selectModel = useCallback( (model: ModelGroup) => { if (modelChatMode(model) === "image" && attachments.length > 0) { - setAttachmentError("Remove the attached images before switching to image generation."); + setAttachmentError("Remove the attached files before switching to image generation."); return; } responsesSessionGenerationRef.current += 1; @@ -159,33 +201,33 @@ export default function ChatConversationPage() { [attachments.length], ); - const addImages = async (files: File[]) => { + const addAttachments = async (files: File[]) => { if (mode !== "chat" || files.length === 0) return; - if (isStreaming || readImagesRef.current) return; - if (attachments.length + files.length > MAX_CHAT_IMAGES) { - setAttachmentError(`Attach up to ${MAX_CHAT_IMAGES} images per message.`); + if (isStreaming || readAttachmentsRef.current) return; + if (attachments.length + files.length > MAX_CHAT_ATTACHMENTS) { + setAttachmentError(`Attach up to ${MAX_CHAT_ATTACHMENTS} files per message.`); return; } const draftVersion = draftVersionRef.current; - readImagesRef.current = true; - setIsReadingImages(true); + readAttachmentsRef.current = true; + setIsReadingAttachments(true); setAttachmentError(null); try { - const images = await Promise.all(files.map(readChatImage)); + const images = await Promise.all(files.map(readChatAttachment)); if (draftVersion === draftVersionRef.current) setAttachments((current) => [...current, ...images]); } catch (error) { if (draftVersion === draftVersionRef.current) - setAttachmentError(error instanceof Error ? error.message : "Could not read these images."); + setAttachmentError(error instanceof Error ? error.message : "Could not read these files."); } finally { - readImagesRef.current = false; - setIsReadingImages(false); + readAttachmentsRef.current = false; + setIsReadingAttachments(false); } }; - const handleImageInput = (event: React.ChangeEvent) => { + const handleAttachmentInput = (event: React.ChangeEvent) => { const files = Array.from(event.target.files ?? []); event.target.value = ""; - void addImages(files); + void addAttachments(files); }; useEffect( @@ -200,7 +242,7 @@ export default function ChatConversationPage() { async ( text: string, historyOverride?: ReplayMessage[], - attachmentOverride?: ChatImageAttachment[], + attachmentOverride?: ChatAttachment[], generationModel?: string, ) => { const trimmed = text.trim(); @@ -208,7 +250,7 @@ export default function ChatConversationPage() { const requestMode = generationModel ? "image" : mode; const model = generationModel ?? selectedModel; if (!hasMessageContent(trimmed, images, requestMode) || !model) return; - if (isStreaming || abortControllerRef.current || readImagesRef.current) return; + if (isStreaming || abortControllerRef.current || readAttachmentsRef.current) return; const teamSelection = chatTeamSelection( models.find((item) => item.model_group === model), selectedTeamId, @@ -318,7 +360,7 @@ export default function ChatConversationPage() { const handleEditAndResend = useCallback( (messageId: string, newContent: string) => { if (!activeConversationId || isStreaming) return; - if (abortControllerRef.current || readImagesRef.current) return; + if (abortControllerRef.current || readAttachmentsRef.current) return; // Compute the truncated history synchronously before the async state update lands, // so handleSend receives the correct pre-edit context rather than the stale closure value. const msgs = activeConversation?.messages ?? []; @@ -411,10 +453,10 @@ export default function ChatConversationPage() { setAttachmentError(null); }; - const isModelSelectionDisabled = isStreaming || isReadingImages; + const isModelSelectionDisabled = isStreaming || isReadingAttachments; const inputBar = (inConversation: boolean) => { - const preparing = isLoadingModels || isReadingImages; + const preparing = isLoadingModels || isReadingAttachments; const model = models.find((item) => item.model_group === selectedModel); const teamSelection = chatTeamSelection(model, selectedTeamId); const canSubmit = hasMessageContent(inputText, attachments, mode) && selectedModel; @@ -427,7 +469,7 @@ export default function ChatConversationPage() { onDrop={(event) => { if (mode === "chat" && event.dataTransfer.files.length) { event.preventDefault(); - void addImages(Array.from(event.dataTransfer.files)); + void addAttachments(Array.from(event.dataTransfer.files)); } }} > @@ -451,31 +493,15 @@ export default function ChatConversationPage() { const images = Array.from(event.clipboardData.files).filter((file) => file.type.startsWith("image/")); if (mode === "chat" && images.length > 0) { event.preventDefault(); - void addImages(images); + void addAttachments(images); } }} attachmentPreview={ - attachments.length > 0 ? ( -
- {attachments.map((image) => ( -
- {image.name} - -

{image.name}

-
- ))} -
- ) : undefined + setAttachments((current) => current.filter((item) => item.id !== id))} + /> } modelPicker={
@@ -498,23 +524,23 @@ export default function ChatConversationPage() { mode === "chat" ? ( <> @@ -553,11 +579,14 @@ export default function ChatConversationPage() { ) } /> - {isReadingImages && ( + {isReadingAttachments && (

- Preparing images… + Preparing files…

)} + {attachments.some((attachment) => attachment.kind === "file") && ( +

File support depends on the selected model and provider.

+ )} {attachmentError && (

{attachmentError} @@ -581,7 +610,7 @@ export default function ChatConversationPage() { <> {storageUnavailable && !storageBannerDismissed && (

- Some chat history or images could not be saved or restored. + Some chat history or attachments could not be saved or restored.
+ ); +} diff --git a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx index 15e6aed14b6..7b22b376884 100644 --- a/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx +++ b/ui/litellm-dashboard/src/components/chat/ChatMessages.tsx @@ -18,7 +18,8 @@ import { useSyntaxTheme } from "@/hooks/useSyntaxTheme"; import ReasoningContent from "@/components/chat_ui/ReasoningContent"; import MCPEventsDisplay from "@/components/chat_ui/MCPEventsDisplay"; import ResponseMetrics from "@/components/chat_ui/ResponseMetrics"; -import { ChatImageAttachment, ChatMessage } from "./types"; +import { ChatAttachment, ChatImageAttachment, ChatMessage } from "./types"; +import ChatFileChip from "./ChatFileChip"; import { isChatImageUrl } from "./imageHelpers"; const REDACTED_KEY_PATTERNS = /token|key|secret|password|auth/i; @@ -127,13 +128,17 @@ function ChatImage({ image, allowImages }: { image: ChatImageAttachment; allowIm ); } -function ChatImages({ images, allowImages }: { images?: ChatImageAttachment[]; allowImages: boolean }) { +function ChatImages({ images, allowImages }: { images?: ChatAttachment[]; allowImages: boolean }) { if (!images?.length) return null; return (
1 ? "grid-cols-2" : "grid-cols-1"}`}> - {images.map((image) => ( - - ))} + {images.map((image) => + image.kind === "file" ? ( + + ) : ( + + ), + )}
); } diff --git a/ui/litellm-dashboard/src/components/chat/attachmentHelpers.test.tsx b/ui/litellm-dashboard/src/components/chat/attachmentHelpers.test.tsx new file mode 100644 index 00000000000..a6ae83864af --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/attachmentHelpers.test.tsx @@ -0,0 +1,72 @@ +import { describe, expect, it } from "vitest"; +import { readChatAttachment, toResponsesMessages } from "./attachmentHelpers"; +import { MAX_CHAT_IMAGE_BYTES } from "./imageHelpers"; +import type { ChatFileAttachment, MCPEvent } from "./types"; + +const file: ChatFileAttachment = { + id: "document-1", + kind: "file", + name: "brief.pdf", + mimeType: "application/pdf", + url: "data:application/pdf;base64,cGRm", +}; + +describe("chat attachments", () => { + it.each([ + ["brief.pdf", "application/pdf"], + ["brief.docx", "application/vnd.openxmlformats-officedocument.wordprocessingml.document"], + ["orders.xlsx", "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet"], + ["notes.txt", "text/plain"], + ["example.py", "text/plain"], + ])("should read %s with its canonical MIME type even when the browser omits it", async (name, mimeType) => { + const attachment = await readChatAttachment(new File(["fixture"], name)); + const expected = { kind: "file", name, mimeType, url: `data:${mimeType};base64,Zml4dHVyZQ==` }; + expect(attachment).toMatchObject(expected); + }); + + it("should reject unsupported, empty, and oversized documents", async () => { + await expect(readChatAttachment(new File(["binary"], "installer.exe"))).rejects.toThrow("Choose an image, PDF"); + await expect(readChatAttachment(new File(["binary"], "file.constructor"))).rejects.toThrow("Choose an image, PDF"); + await expect(readChatAttachment(new File(["legacy"], "legacy.xls"))).rejects.toThrow("Excel (.xlsx)"); + await expect(readChatAttachment(new File([], "empty.pdf"))).rejects.toThrow("empty"); + const large = new File(["pdf"], "large.pdf"); + Object.defineProperty(large, "size", { value: MAX_CHAT_IMAGE_BYTES + 1 }); + await expect(readChatAttachment(large)).rejects.toThrow("10 MB"); + }); + + it("should preserve mixed files, images, and tool-only assistant history for Responses", () => { + const image = { id: "image-1", name: "diagram.png", url: "data:image/png;base64,aW1hZ2U=" }; + const mcpEvents: MCPEvent[] = [ + { + type: "response.output_item.done", + item: { type: "mcp_call", id: "lookup-1", name: "lookup", arguments: "{}", output: "Found it" }, + }, + ]; + expect( + toResponsesMessages([ + { role: "user", content: "Compare", attachments: [image, file] }, + { role: "assistant", content: "I read the file.", mcpEvents }, + { role: "assistant", content: "", mcpEvents }, + ]), + ).toEqual([ + { + role: "user", + content: [ + { type: "input_text", text: "Compare" }, + { type: "input_image", image_url: image.url, detail: "auto" }, + { type: "input_file", filename: file.name, file_data: file.url }, + ], + }, + { role: "assistant", content: "I read the file.", mcpEvents }, + { role: "assistant", content: "", mcpEvents }, + ]); + }); + + it("should fail clearly instead of silently dropping unavailable document contents", () => { + for (const url of ["", "javascript:alert(1)", "https://example.com/private.pdf", "data:text/html;base64,cGRm"]) { + expect(() => toResponsesMessages([{ role: "user", content: "Read", attachments: [{ ...file, url }] }])).toThrow( + "A saved file (brief.pdf) is unavailable", + ); + } + }); +}); diff --git a/ui/litellm-dashboard/src/components/chat/attachmentHelpers.ts b/ui/litellm-dashboard/src/components/chat/attachmentHelpers.ts new file mode 100644 index 00000000000..bd3ff9aaa21 --- /dev/null +++ b/ui/litellm-dashboard/src/components/chat/attachmentHelpers.ts @@ -0,0 +1,112 @@ +import type { MultimodalContent } from "@/components/chat_ui/types"; +import type { ResponsesMessage } from "@/components/llm_calls/responses_api"; +import { v4 as uuid } from "uuid"; +import { CHAT_IMAGE_ACCEPT, MAX_CHAT_IMAGE_BYTES, isChatImageUrl, readChatImage } from "./imageHelpers"; +import type { ChatAttachment, ChatFileAttachment, ChatMessage } from "./types"; + +export const MAX_CHAT_ATTACHMENTS = 4; +const DOCUMENT_TYPES: Record = { + pdf: "application/pdf", + docx: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + xlsx: "application/vnd.openxmlformats-officedocument.spreadsheetml.sheet", + csv: "text/csv", + tsv: "text/tab-separated-values", + txt: "text/plain", + md: "text/markdown", + json: "application/json", + html: "text/html", + xml: "application/xml", +}; +const CODE_EXTENSIONS = [ + "c", + "cpp", + "css", + "go", + "h", + "java", + "js", + "jsx", + "log", + "py", + "rb", + "rs", + "sh", + "sql", + "ts", + "tsx", + "yaml", + "yml", +]; +for (const extension of CODE_EXTENSIONS) DOCUMENT_TYPES[extension] = "text/plain"; + +export const CHAT_ATTACHMENT_ACCEPT = [ + CHAT_IMAGE_ACCEPT, + ...Object.keys(DOCUMENT_TYPES).map((extension) => `.${extension}`), +].join(","); + +export function isChatFileData(file: ChatFileAttachment): boolean { + return Boolean( + Object.values(DOCUMENT_TYPES).includes(file.mimeType) && + file.url.startsWith(`data:${file.mimeType};base64,`) && + /^[A-Za-z0-9+/]+={0,2}$/.test(file.url.split(",")[1] ?? ""), + ); +} + +export async function readChatAttachment(file: File): Promise { + if (file.type.startsWith("image/")) return readChatImage(file); + const extension = file.name.split(".").pop()?.toLowerCase() ?? ""; + const mimeType = Object.hasOwn(DOCUMENT_TYPES, extension) ? DOCUMENT_TYPES[extension] : undefined; + if (!mimeType) throw new Error("Choose an image, PDF, Word (.docx), Excel (.xlsx), text, or code file."); + if (file.size === 0) throw new Error(`${file.name} is empty. Choose another file.`); + if (file.size > MAX_CHAT_IMAGE_BYTES) throw new Error(`${file.name} is too large. Files must be 10 MB or smaller.`); + const url = await new Promise((resolve, reject) => { + const reader = new FileReader(); + reader.onload = () => { + if (typeof reader.result === "string") resolve(reader.result); + else reject(new Error(`Could not read ${file.name}. Choose the file again.`)); + }; + reader.onerror = () => reject(new Error(`Could not read ${file.name}. Choose the file again.`)); + reader.onabort = () => reject(new Error(`Reading ${file.name} was canceled.`)); + reader.readAsDataURL(new Blob([file], { type: mimeType })); + }); + return { id: uuid(), kind: "file", name: file.name, url, mimeType }; +} + +export function toResponsesMessages( + messages: Pick< + ChatMessage, + "role" | "content" | "attachments" | "generatedImages" | "imageGeneration" | "mcpEvents" + >[], +): ResponsesMessage[] { + return messages.flatMap((message): ResponsesMessage[] => { + if (message.role === "tool") return []; + if (message.role === "assistant") { + return message.content.trim() || message.mcpEvents?.length + ? [ + { + role: "assistant", + content: message.content, + ...(message.mcpEvents?.length ? { mcpEvents: message.mcpEvents } : {}), + }, + ] + : []; + } + if (!message.attachments?.length) return [{ role: "user", content: message.content }]; + const content: MultimodalContent[] = []; + if (message.content.trim()) content.push({ type: "input_text", text: message.content }); + for (const attachment of message.attachments) { + if (attachment.kind === "file") { + if (!isChatFileData(attachment)) { + throw new Error(`A saved file (${attachment.name}) is unavailable. Start a new chat and attach it again.`); + } + content.push({ type: "input_file", filename: attachment.name, file_data: attachment.url }); + continue; + } + if (!isChatImageUrl(attachment.url)) { + throw new Error(`A saved image (${attachment.name}) is unavailable. Start a new chat and attach it again.`); + } + content.push({ type: "input_image", image_url: attachment.url, detail: "auto" }); + } + return [{ role: "user", content }]; + }); +} diff --git a/ui/litellm-dashboard/src/components/chat/chatImageStorage.ts b/ui/litellm-dashboard/src/components/chat/chatImageStorage.ts index af8dc40d954..364053dc1ab 100644 --- a/ui/litellm-dashboard/src/components/chat/chatImageStorage.ts +++ b/ui/litellm-dashboard/src/components/chat/chatImageStorage.ts @@ -1,4 +1,4 @@ -import type { ChatImageAttachment, Conversation } from "./types"; +import type { ChatAttachment, Conversation } from "./types"; const DATABASE_NAME = "litellm_chat_media_v1"; const STORE_NAME = "images"; @@ -21,8 +21,8 @@ function openDatabase(): Promise { }; resolve(database); }; - request.onerror = () => reject(request.error ?? new Error("Image storage is unavailable.")); - request.onblocked = () => reject(new Error("Image storage is blocked by another browser tab.")); + request.onerror = () => reject(request.error ?? new Error("Attachment storage is unavailable.")); + request.onblocked = () => reject(new Error("Attachment storage is blocked by another browser tab.")); }).catch((error: unknown) => { databasePromise = undefined; throw error; @@ -31,13 +31,13 @@ function openDatabase(): Promise { return databasePromise; } -export function getConversationImages(conversations: Conversation[]): ChatImageAttachment[] { +export function getConversationImages(conversations: Conversation[]): ChatAttachment[] { return conversations.flatMap((conversation) => conversation.messages.flatMap((message) => [...(message.attachments ?? []), ...(message.generatedImages ?? [])]), ); } -export function persistChatImages(userId: string, images: ChatImageAttachment[]): Promise { +export function persistChatImages(userId: string, images: ChatAttachment[]): Promise { return Promise.all( images .filter((image) => image.url.startsWith("data:")) @@ -52,8 +52,9 @@ export function persistChatImages(userId: string, images: ChatImageAttachment[]) const transaction = database.transaction(STORE_NAME, "readwrite"); transaction.objectStore(STORE_NAME).put({ userId, id: image.id, url: image.url }); transaction.oncomplete = () => resolve(); - transaction.onerror = () => reject(transaction.error ?? new Error("Could not save the image.")); - transaction.onabort = () => reject(transaction.error ?? new Error("Saving the image was canceled.")); + transaction.onerror = () => reject(transaction.error ?? new Error("Could not save the attachment.")); + transaction.onabort = () => + reject(transaction.error ?? new Error("Saving the attachment was canceled.")); }), ) .catch((error: unknown) => { @@ -80,8 +81,8 @@ export async function loadChatImages(userId: string, ids: string[]): Promise resolve(loaded); - transaction.onerror = () => reject(transaction.error ?? new Error("Could not restore saved images.")); - transaction.onabort = () => reject(transaction.error ?? new Error("Restoring saved images was canceled.")); + transaction.onerror = () => reject(transaction.error ?? new Error("Could not restore saved attachments.")); + transaction.onabort = () => reject(transaction.error ?? new Error("Restoring saved attachments was canceled.")); }); } @@ -93,13 +94,13 @@ export async function deleteChatImages(userId: string, ids: string[]): Promise resolve(); - transaction.onerror = () => reject(transaction.error ?? new Error("Could not delete saved images.")); - transaction.onabort = () => reject(transaction.error ?? new Error("Deleting saved images was canceled.")); + transaction.onerror = () => reject(transaction.error ?? new Error("Could not delete saved attachments.")); + transaction.onabort = () => reject(transaction.error ?? new Error("Deleting saved attachments was canceled.")); }); } export function withoutImageData(conversations: Conversation[]): Conversation[] { - const stripData = (images: ChatImageAttachment[] | undefined) => + const stripData = (images: T[] | undefined) => images?.map((image) => (image.url.startsWith("data:") ? { ...image, url: "" } : image)); return conversations.map((conversation) => ({ ...conversation, @@ -112,7 +113,7 @@ export function withoutImageData(conversations: Conversation[]): Conversation[] } export function withRestoredImages(conversations: Conversation[], images: Map): Conversation[] { - const restore = (attachments: ChatImageAttachment[]) => + const restore = (attachments: T[]) => attachments.map((image) => (!image.url && images.has(image.id) ? { ...image, url: images.get(image.id)! } : image)); return conversations.map((conversation) => ({ ...conversation, diff --git a/ui/litellm-dashboard/src/components/chat/design.md b/ui/litellm-dashboard/src/components/chat/design.md index e34fec062c9..d3f5b597b67 100644 --- a/ui/litellm-dashboard/src/components/chat/design.md +++ b/ui/litellm-dashboard/src/components/chat/design.md @@ -202,7 +202,7 @@ If you need "shrink to fit content, cap at N px" rather than "always N px," you ### Input / Textarea -Use `Input` and `Textarea` from `@/components/ui`. Standalone Chat uses `MoyaiChatComposer` for the source editor and toolbar arrangement, preserving submission, cancellation, image attachments and IME handling. Other chat surfaces continue using the shared `ChatComposer` / `InputGroupTextarea`; this visual port is scoped to standalone Chat. +Use `Input` and `Textarea` from `@/components/ui`. Standalone Chat uses `MoyaiChatComposer` for the source editor and toolbar arrangement, preserving submission, cancellation, attachments and IME handling. Uploads share a limit of four files per message and 10 MB per file. Images retain previews; PDF, DOCX, XLSX, text and code files use named file chips and native Responses `input_file` parts. File support depends on the selected model and provider. Attachment bytes stay in the existing user-scoped IndexedDB store; history metadata stays small. Legacy DOC/XLS files are not accepted. Other chat surfaces continue using the shared `ChatComposer` / `InputGroupTextarea`; this visual port is scoped to standalone Chat. ```tsx import { Input } from "@/components/ui/input"; diff --git a/ui/litellm-dashboard/src/components/chat/imageHelpers.test.tsx b/ui/litellm-dashboard/src/components/chat/imageHelpers.test.tsx index 30d4afed04b..5f605a0a850 100644 --- a/ui/litellm-dashboard/src/components/chat/imageHelpers.test.tsx +++ b/ui/litellm-dashboard/src/components/chat/imageHelpers.test.tsx @@ -1,12 +1,7 @@ +import { toResponsesMessages } from "./attachmentHelpers"; import { beforeEach, describe, expect, it, vi } from "vitest"; import OpenAI from "openai"; -import { - generateChatImages, - isChatImageUrl, - MAX_CHAT_IMAGE_BYTES, - readChatImage, - toResponsesMessages, -} from "./imageHelpers"; +import { generateChatImages, isChatImageUrl, MAX_CHAT_IMAGE_BYTES, readChatImage } from "./imageHelpers"; const { generate } = vi.hoisted(() => ({ generate: vi.fn() })); vi.mock("openai", () => ({ diff --git a/ui/litellm-dashboard/src/components/chat/imageHelpers.ts b/ui/litellm-dashboard/src/components/chat/imageHelpers.ts index 9756991451b..c62286a2991 100644 --- a/ui/litellm-dashboard/src/components/chat/imageHelpers.ts +++ b/ui/litellm-dashboard/src/components/chat/imageHelpers.ts @@ -1,10 +1,8 @@ import OpenAI from "openai"; +import { v4 as uuid } from "uuid"; import { getGatewayBaseUrl } from "@/components/networking"; -import type { MultimodalContent } from "@/components/chat_ui/types"; -import type { ResponsesMessage } from "@/components/llm_calls/responses_api"; -import type { ChatImageAttachment, ChatMessage } from "./types"; +import type { ChatImageAttachment } from "./types"; -export const MAX_CHAT_IMAGES = 4; export const MAX_CHAT_IMAGE_BYTES = 10 * 1024 * 1024; export const CHAT_IMAGE_ACCEPT = "image/png,image/jpeg,image/webp,image/gif"; const IMAGE_TYPES = new Set(CHAT_IMAGE_ACCEPT.split(",")); @@ -36,39 +34,7 @@ export async function readChatImage(file: File): Promise { reader.onabort = () => reject(new Error(`Reading ${file.name} was canceled.`)); reader.readAsDataURL(file); }); - return { id: crypto.randomUUID(), name: file.name, url, mimeType: file.type }; -} - -export function toResponsesMessages( - messages: Pick< - ChatMessage, - "role" | "content" | "attachments" | "generatedImages" | "imageGeneration" | "mcpEvents" - >[], -): ResponsesMessage[] { - return messages.flatMap((message): ResponsesMessage[] => { - if (message.role === "tool") return []; - if (message.role === "assistant") { - return message.content.trim() || message.mcpEvents?.length - ? [ - { - role: "assistant", - content: message.content, - ...(message.mcpEvents?.length ? { mcpEvents: message.mcpEvents } : {}), - }, - ] - : []; - } - if (!message.attachments?.length) return [{ role: "user", content: message.content }]; - const content: MultimodalContent[] = []; - if (message.content.trim()) content.push({ type: "input_text", text: message.content }); - for (const attachment of message.attachments) { - if (!isChatImageUrl(attachment.url)) { - throw new Error(`A saved image (${attachment.name}) is unavailable. Start a new chat and attach it again.`); - } - content.push({ type: "input_image", image_url: attachment.url, detail: "auto" }); - } - return [{ role: "user", content }]; - }); + return { id: uuid(), name: file.name, url, mimeType: file.type }; } export async function generateChatImages({ @@ -95,7 +61,7 @@ export async function generateChatImages({ const images = (response.data ?? []).flatMap((image, index): ChatImageAttachment[] => { const url = image.b64_json ? `data:${mimeType};base64,${image.b64_json}` : image.url; if (!url || !isChatImageUrl(url)) return []; - return [{ id: crypto.randomUUID(), name: `Generated image ${index + 1}`, url }]; + return [{ id: uuid(), name: `Generated image ${index + 1}`, url }]; }); if (!images.length) throw new Error("The model did not return an image. Try another prompt or image model."); return images; diff --git a/ui/litellm-dashboard/src/components/chat/runChatRequest.test.ts b/ui/litellm-dashboard/src/components/chat/runChatRequest.test.ts index 48f9c2acb2b..60ab8772112 100644 --- a/ui/litellm-dashboard/src/components/chat/runChatRequest.test.ts +++ b/ui/litellm-dashboard/src/components/chat/runChatRequest.test.ts @@ -1,12 +1,16 @@ +import OpenAI from "openai"; import { beforeEach, describe, expect, it, vi } from "vitest"; import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; -import { generateChatImages, toResponsesMessages } from "./imageHelpers"; +import { generateChatImages } from "./imageHelpers"; +import { toResponsesMessages } from "./attachmentHelpers"; import { runChatRequest } from "./runChatRequest"; import type { MCPEvent } from "./types"; vi.mock("@/components/llm_calls/responses_api", () => ({ makeOpenAIResponsesRequest: vi.fn() })); vi.mock("./imageHelpers", () => ({ generateChatImages: vi.fn(), +})); +vi.mock("./attachmentHelpers", () => ({ toResponsesMessages: vi.fn((messages) => messages), })); @@ -105,6 +109,35 @@ describe("runChatRequest", () => { expect(options.onResponseId).toHaveBeenCalledWith(null); }); + it.each([400, 422, 401, 429, 500])( + "should suggest a supported model only when a file request is rejected with %i", + async (status) => { + const options = requestOptions({ + messages: [ + { + role: "user", + content: "Read the report", + attachments: [ + { + id: "report", + kind: "file", + name: "report.pdf", + mimeType: "application/pdf", + url: "data:application/pdf;base64,JVBERg==", + }, + ], + }, + ], + }); + const error = new OpenAI.APIError(status, { message: "Provider detail" }, "Provider detail", new Headers()); + vi.mocked(makeOpenAIResponsesRequest).mockRejectedValue(error); + await runChatRequest(options); + const hint = + status === 400 || status === 422 ? " Try a model that supports file inputs, or paste the relevant text." : ""; + expect(options.onUpdate).toHaveBeenLastCalledWith({ content: error.message + hint }); + }, + ); + it("should return generated images through the same message update callback", async () => { const options = requestOptions({ mode: "image", model: "image-model" }); const images = [{ id: "image-1", name: "Generated image 1", url: "https://images.example/image.png" }]; @@ -121,6 +154,33 @@ describe("runChatRequest", () => { expect(generateChatImages).toHaveBeenCalledWith(expectedRequest); }); + it("should not suggest file inputs for an image-generation rejection after a file message", async () => { + const options = requestOptions({ + mode: "image", + messages: [ + { + role: "user", + content: "Read the report", + attachments: [ + { + id: "report", + kind: "file", + name: "report.pdf", + mimeType: "application/pdf", + url: "data:application/pdf;base64,JVBERg==", + }, + ], + }, + ], + }); + const error = new OpenAI.APIError(400, { message: "Invalid image size" }, "Invalid image size", new Headers()); + vi.mocked(generateChatImages).mockRejectedValue(error); + + await runChatRequest(options); + + expect(options.onUpdate).toHaveBeenLastCalledWith({ content: error.message }); + }); + it("should discard image results that arrive after the request was stopped", async () => { const controller = new AbortController(); const options = requestOptions({ mode: "image", signal: controller.signal }); diff --git a/ui/litellm-dashboard/src/components/chat/runChatRequest.ts b/ui/litellm-dashboard/src/components/chat/runChatRequest.ts index c8b7fbf0907..0f98e8621ab 100644 --- a/ui/litellm-dashboard/src/components/chat/runChatRequest.ts +++ b/ui/litellm-dashboard/src/components/chat/runChatRequest.ts @@ -1,5 +1,7 @@ +import OpenAI from "openai"; import { makeOpenAIResponsesRequest } from "@/components/llm_calls/responses_api"; -import { generateChatImages, toResponsesMessages } from "./imageHelpers"; +import { generateChatImages } from "./imageHelpers"; +import { toResponsesMessages } from "./attachmentHelpers"; import type { AssistantMessageUpdate, ChatMessage, MCPEvent } from "./types"; interface ChatRequestOptions { @@ -19,6 +21,21 @@ interface ChatRequestOptions { onUpdate: (update: AssistantMessageUpdate) => void; } +function requestErrorMessage( + error: unknown, + messages: ChatRequestOptions["messages"], + mode: ChatRequestOptions["mode"], +): string { + const message = error instanceof Error ? error.message : "Something went wrong. Please try again."; + const hasFile = + mode === "chat" && messages.some((item) => item.attachments?.some((attachment) => attachment.kind === "file")); + const fileRejectionStatus = error instanceof OpenAI.APIError && (error.status === 400 || error.status === 422); + if (hasFile && fileRejectionStatus) { + return message + " Try a model that supports file inputs, or paste the relevant text."; + } + return message; +} + export async function runChatRequest({ model, mode, @@ -85,7 +102,7 @@ export async function runChatRequest({ } catch (error: unknown) { onResponseId(null); const stopped = signal.aborted || (error instanceof Error && error.name === "AbortError"); - const errorMessage = error instanceof Error ? error.message : "Something went wrong. Please try again."; + const errorMessage = requestErrorMessage(error, messages, mode); onUpdate({ content: stopped ? content + " [stopped]" : content || errorMessage }); } finally { if (mcpEvents.length) onUpdate({ mcpEvents }); diff --git a/ui/litellm-dashboard/src/components/chat/types.ts b/ui/litellm-dashboard/src/components/chat/types.ts index 121770ece16..25b232e40c4 100644 --- a/ui/litellm-dashboard/src/components/chat/types.ts +++ b/ui/litellm-dashboard/src/components/chat/types.ts @@ -2,13 +2,24 @@ import type { MCPEvent } from "../mcp_tools/types"; import type { TokenUsage } from "../chat_ui/ResponseMetrics"; export type { MCPEvent }; -export interface ChatImageAttachment { +interface AttachmentData { id: string; name: string; url: string; mimeType?: string; } +export interface ChatImageAttachment extends AttachmentData { + kind?: "image"; +} + +export interface ChatFileAttachment extends AttachmentData { + kind: "file"; + mimeType: string; +} + +export type ChatAttachment = ChatImageAttachment | ChatFileAttachment; + export interface ChatImageGeneration { model: string; prompt: string; @@ -18,7 +29,7 @@ export interface ChatMessage { id: string; role: "user" | "assistant" | "tool"; content: string; - attachments?: ChatImageAttachment[]; + attachments?: ChatAttachment[]; generatedImages?: ChatImageAttachment[]; imageGeneration?: ChatImageGeneration; reasoningContent?: string; diff --git a/ui/litellm-dashboard/src/components/chat/useChatHistory.test.ts b/ui/litellm-dashboard/src/components/chat/useChatHistory.test.ts index a09a3ab056b..e35edc3b4a0 100644 --- a/ui/litellm-dashboard/src/components/chat/useChatHistory.test.ts +++ b/ui/litellm-dashboard/src/components/chat/useChatHistory.test.ts @@ -247,4 +247,33 @@ describe("useChatHistory", () => { expect(second.result.current.conversations).toEqual([]); expect(deleteChatImages).toHaveBeenCalledWith("alice", [image.id]); }); + it("should save, restore, and delete document bytes with the existing user-scoped media store", async () => { + const file = { + id: "file-1", + kind: "file" as const, + name: "brief.docx", + mimeType: "application/vnd.openxmlformats-officedocument.wordprocessingml.document", + url: "data:application/vnd.openxmlformats-officedocument.wordprocessingml.document;base64,ZmlsZQ==", + }; + const first = renderHook(() => useChatHistory(null, "alice")); + let id!: string; + act(() => { + id = first.result.current.createConversation("file-model"); + }); + act(() => { + first.result.current.appendMessage(id, { role: "user", content: "Read this", attachments: [file] }); + }); + await waitFor(() => expect(persistChatImages).toHaveBeenCalledWith("alice", [file])); + await waitFor(() => expect(localStorage.getItem(aliceStorageKey)).toContain("brief.docx")); + expect(localStorage.getItem(aliceStorageKey)).not.toContain(file.url); + first.unmount(); + vi.mocked(loadChatImages).mockResolvedValue(new Map([[file.id, file.url]])); + const restored = renderHook(() => useChatHistory(id, "alice")); + await waitFor(() => expect(restored.result.current.activeConversation?.messages[0].attachments).toEqual([file])); + act(() => { + restored.result.current.deleteConversation(id); + }); + await waitFor(() => expect(deleteChatImages).toHaveBeenCalledWith("alice", [file.id])); + expect(JSON.parse(localStorage.getItem(aliceStorageKey)!)).toEqual([]); + }); }); diff --git a/ui/litellm-dashboard/src/components/chat/useChatHistory.ts b/ui/litellm-dashboard/src/components/chat/useChatHistory.ts index 21c73bdbd98..cc9e9504202 100644 --- a/ui/litellm-dashboard/src/components/chat/useChatHistory.ts +++ b/ui/litellm-dashboard/src/components/chat/useChatHistory.ts @@ -1,3 +1,4 @@ +import { v4 as uuid } from "uuid"; import { useCallback, useEffect, useState } from "react"; import { AssistantMessageUpdate, ChatMessage, Conversation } from "./types"; import { @@ -174,7 +175,7 @@ export function useChatHistory( }, [conversations, userId, historyLoadFailed]); const createConversation = useCallback((model: string): string => { - const id = crypto.randomUUID(); + const id = uuid(); const now = Date.now(); const newConversation: Conversation = { id, @@ -193,7 +194,7 @@ export function useChatHistory( const appendMessage = useCallback((conversationId: string, message: Omit) => { const newMessage: ChatMessage = { ...message, - id: crypto.randomUUID(), + id: uuid(), timestamp: Date.now(), }; diff --git a/ui/litellm-dashboard/src/components/chat_ui/types.ts b/ui/litellm-dashboard/src/components/chat_ui/types.ts index 639a42d6dd6..1cfd8b8f187 100644 --- a/ui/litellm-dashboard/src/components/chat_ui/types.ts +++ b/ui/litellm-dashboard/src/components/chat_ui/types.ts @@ -71,7 +71,9 @@ export interface MessageType { } export interface MultimodalContent { - type: "input_text" | "input_image"; + type: "input_text" | "input_image" | "input_file"; + filename?: string; + file_data?: string; text?: string; image_url?: string; detail?: "low" | "high" | "auto";