diff --git a/src/bub/builtin/hook_impl.py b/src/bub/builtin/hook_impl.py index 336c3d44..ade488f8 100644 --- a/src/bub/builtin/hook_impl.py +++ b/src/bub/builtin/hook_impl.py @@ -2,7 +2,7 @@ from datetime import datetime from difflib import get_close_matches from pathlib import Path -from typing import cast +from typing import Any, cast import typer from loguru import logger @@ -15,7 +15,7 @@ from bub.channels.admission import AdmitDecision, SteeringInbox, TurnSnapshot from bub.channels.base import Channel from bub.channels.contracts import MessageHandler -from bub.channels.message import ChannelMessage, MediaItem +from bub.channels.message import ChannelMessage, MediaItem, audio_format_from_mime_type from bub.envelope import Envelope, content_of, field_of from bub.framework import BubFramework from bub.hooks import hookimpl @@ -63,6 +63,16 @@ DEFAULT_CONTINUE_PROMPT = "Continue the task until all targets are completed." +def _input_audio_part(data_url: str, mime_type: str) -> dict[str, Any] | None: + prefix, separator, data = data_url.partition("base64,") + if not separator or not prefix.startswith("data:audio/") or not data: + return None + return { + "type": "input_audio", + "input_audio": {"data": data, "format": audio_format_from_mime_type(mime_type)}, + } + + class BuiltinImpl: """Default hook implementations for basic runtime operations.""" @@ -198,13 +208,18 @@ async def build_prompt(self, message: ChannelMessage, session_id: str, state: Tu media_parts: list[dict] = [] for item in cast("list[MediaItem]", media): match item.type: - case "image": + case "image" | "video": data_url = await item.get_url() if not data_url: continue - media_parts.append({"type": "image_url", "image_url": {"url": data_url}}) + part_type = f"{item.type}_url" + media_parts.append({"type": part_type, part_type: {"url": data_url}}) + case "audio": + data_url = await item.get_url() + if data_url and (audio_part := _input_audio_part(data_url, item.mime_type)): + media_parts.append(audio_part) case _: - pass # TODO: Not supported for now + pass if media_parts: return [{"type": "text", "text": text}, *media_parts] return text diff --git a/src/bub/builtin/model_runner.py b/src/bub/builtin/model_runner.py index 099f3f68..169a250a 100644 --- a/src/bub/builtin/model_runner.py +++ b/src/bub/builtin/model_runner.py @@ -28,6 +28,7 @@ from bub.builtin.codex_provider import OpenaiCodexProvider, should_use_openai_codex_provider from bub.builtin.settings import AgentSettings, ModelCandidate +from bub.channels.message import audio_mime_type_from_format from bub.errors import BubError, ErrorKind from bub.hooks.interception import ( AgentHooks, @@ -45,6 +46,7 @@ ) TOOL_ARGUMENTS_ADAPTER = TypeAdapter(dict[str, Any]) CompletionResult = ChatCompletion | ParsedChatCompletion[Any] | AsyncIterator[ChatCompletionChunk] +GOOGLE_FILE_CONTENT_PROVIDERS = frozenset({LLMProvider.GEMINI, LLMProvider.VERTEXAI}) def _extra_options(llm: AnyLLM, *, stream: bool) -> dict[str, Any]: @@ -56,6 +58,48 @@ def _extra_options(llm: AnyLLM, *, stream: bool) -> dict[str, Any]: return {} +def _adapt_messages_for_provider(messages: list[dict[str, Any]], provider: LLMProvider) -> list[dict[str, Any]]: + """Translate canonical multimodal blocks when a provider uses a different wire format.""" + if provider not in GOOGLE_FILE_CONTENT_PROVIDERS: + return messages + + adapted_messages: list[dict[str, Any]] = [] + for message in messages: + content = message.get("content") + if not isinstance(content, list): + adapted_messages.append(message) + continue + + adapted_content: list[Any] = [] + changed = False + for part in content: + if not isinstance(part, dict): + adapted_content.append(part) + continue + + if part.get("type") == "video_url": + video_url = part.get("video_url") + url = video_url.get("url") if isinstance(video_url, dict) else None + if isinstance(url, str) and url: + adapted_content.append({"type": "file", "file": {"file_data": url}}) + changed = True + continue + elif part.get("type") == "input_audio": + input_audio = part.get("input_audio") + data = input_audio.get("data") if isinstance(input_audio, dict) else None + audio_format = input_audio.get("format") if isinstance(input_audio, dict) else None + if isinstance(data, str) and data and isinstance(audio_format, str) and audio_format: + mime_type = audio_mime_type_from_format(audio_format) + file_data = f"data:{mime_type};base64,{data}" + adapted_content.append({"type": "file", "file": {"file_data": file_data}}) + changed = True + continue + adapted_content.append(part) + + adapted_messages.append({**message, "content": adapted_content} if changed else message) + return adapted_messages + + class ModelRunner: def __init__(self, settings: AgentSettings, hooks: AgentHooks | None = None) -> None: self.settings = settings @@ -90,12 +134,12 @@ async def completion_response( reasoning_effort: str | None = None, ) -> CompletionResult: tool_payloads = [tool.to_schema() for tool in tools] or None - completion_messages: list[dict[str, Any] | ChatCompletionMessage] = list(messages) clients = list(self.iter_llm_clients(model)) completion_error: Exception | None = None for index, (candidate, llm) in enumerate(clients): try: streaming = llm.SUPPORTS_COMPLETION_STREAMING + completion_messages = _adapt_messages_for_provider(messages, candidate.provider) completion_kwargs = { **self.settings.completion_args, **_extra_options(llm, stream=streaming), diff --git a/src/bub/channels/message.py b/src/bub/channels/message.py index b026fa8c..9027a2a3 100644 --- a/src/bub/channels/message.py +++ b/src/bub/channels/message.py @@ -9,6 +9,32 @@ type MessageKind = Literal["error", "normal", "command"] type MediaType = Literal["image", "audio", "video", "document"] +_AUDIO_FORMAT_TO_MIME_TYPE = { + "aiff": "audio/aiff", + "flac": "audio/flac", + "m4a": "audio/mp4", + "mp3": "audio/mpeg", + "ogg": "audio/ogg", + "wav": "audio/wav", + "webm": "audio/webm", +} +_AUDIO_MIME_TYPE_TO_FORMAT = { + **{mime_type: audio_format for audio_format, mime_type in _AUDIO_FORMAT_TO_MIME_TYPE.items()}, + "audio/x-aiff": "aiff", + "audio/x-flac": "flac", + "audio/x-m4a": "m4a", + "audio/x-wav": "wav", +} + + +def audio_format_from_mime_type(mime_type: str) -> str: + normalized = mime_type.partition(";")[0].strip().lower() + return _AUDIO_MIME_TYPE_TO_FORMAT.get(normalized, normalized.removeprefix("audio/") or "unknown") + + +def audio_mime_type_from_format(audio_format: str) -> str: + return _AUDIO_FORMAT_TO_MIME_TYPE.get(audio_format, f"audio/{audio_format}") + @dataclass class MediaItem: @@ -18,7 +44,7 @@ class MediaItem: mime_type: str filename: str | None = None url: str | None = None - data_fetcher: Callable[[], Awaitable[bytes]] | None = None + data_fetcher: Callable[[], Awaitable[bytes | None]] | None = None async def get_url(self) -> str | None: """Get a URL for the media, fetching data if necessary.""" @@ -26,6 +52,8 @@ async def get_url(self) -> str | None: return self.url if self.data_fetcher is not None: data = await self.data_fetcher() + if data is None: + return None return f"data:{self.mime_type};base64,{base64.b64encode(data).decode('utf-8')}" return None diff --git a/src/bub/channels/telegram.py b/src/bub/channels/telegram.py index 142170a1..d80410e2 100644 --- a/src/bub/channels/telegram.py +++ b/src/bub/channels/telegram.py @@ -378,7 +378,7 @@ async def _parse_audio(self, message: Message) -> tuple[str, dict[str, Any] | No duration = audio.duration or 0 metadata = exclude_none({ "file_id": audio.file_id, - "mime_type": audio.mime_type, + "mime_type": audio.mime_type or "audio/mpeg", "file_size": audio.file_size, "duration": audio.duration, "title": audio.title, @@ -389,12 +389,12 @@ async def _parse_audio(self, message: Message) -> tuple[str, dict[str, Any] | No return f"[Audio: {performer} - {title} ({duration}s)]", metadata return f"[Audio: {title} ({duration}s)]", metadata - async def _download_media(self, file_id: str, file_size: int) -> bytes | None: + async def _download_media(self, file_id: str, file_size: int | None) -> bytes | None: if not file_id: raise ValueError("file_id must not be empty") if self._bot_getter is None: raise RuntimeError("Telegram bot is not configured for media downloads.") - if file_size > 2 * 1024 * 1024: # limit to 2MB + if file_size is not None and file_size > 2 * 1024 * 1024: # limit to 2MB return None bot = self._bot_getter() if bot is None: @@ -441,7 +441,7 @@ async def _parse_video(self, message: Message) -> tuple[str, dict[str, Any] | No "width": video.width, "height": video.height, "duration": video.duration, - "mime_type": video.mime_type, + "mime_type": video.mime_type or "video/mp4", "data_fetcher": lambda: self._download_media(video.file_id, video.file_size), }) return formatted, metadata diff --git a/tests/test_builtin_model_runner.py b/tests/test_builtin_model_runner.py index 374ebdf8..9e849860 100644 --- a/tests/test_builtin_model_runner.py +++ b/tests/test_builtin_model_runner.py @@ -10,12 +10,54 @@ from any_llm.providers.openai.base import BaseOpenAIProvider from any_llm.types.completion import ChatCompletionChunk, ChatCompletionMessageFunctionToolCall, Function -from bub.builtin.model_runner import ModelRunner, tool_invocation_from_native +from bub.builtin.model_runner import ModelRunner, _adapt_messages_for_provider, tool_invocation_from_native from bub.builtin.settings import AgentSettings, ModelCandidate from bub.tape import AsyncTapeStoreAdapter, InMemoryTapeStore, Tape, TapeContext from bub.tools import ToolExecutor +@pytest.mark.parametrize("provider", [LLMProvider.GEMINI, LLMProvider.VERTEXAI]) +def test_adapt_messages_converts_video_url_for_google_providers(provider: LLMProvider) -> None: + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this video"}, + {"type": "video_url", "video_url": {"url": "data:video/mp4;base64,dmlkZW8="}}, + {"type": "input_audio", "input_audio": {"data": "YXVkaW8=", "format": "ogg"}}, + ], + } + ] + + result = _adapt_messages_for_provider(messages, provider) + + assert result == [ + { + "role": "user", + "content": [ + {"type": "text", "text": "describe this video"}, + {"type": "file", "file": {"file_data": "data:video/mp4;base64,dmlkZW8="}}, + {"type": "file", "file": {"file_data": "data:audio/ogg;base64,YXVkaW8="}}, + ], + } + ] + assert messages[0]["content"][1]["type"] == "video_url" + + +def test_adapt_messages_keeps_native_multimodal_blocks_for_openrouter() -> None: + messages = [ + { + "role": "user", + "content": [ + {"type": "video_url", "video_url": {"url": "data:video/mp4;base64,dmlkZW8="}}, + {"type": "input_audio", "input_audio": {"data": "YXVkaW8=", "format": "ogg"}}, + ], + } + ] + + assert _adapt_messages_for_provider(messages, LLMProvider.OPENROUTER) is messages + + @pytest.mark.asyncio async def test_unknown_tool_placeholder_surfaces_error_without_hooks() -> None: tool_call = ChatCompletionMessageFunctionToolCall( diff --git a/tests/test_image_message.py b/tests/test_image_message.py index 132e15b6..e797b541 100644 --- a/tests/test_image_message.py +++ b/tests/test_image_message.py @@ -10,7 +10,7 @@ from bub.builtin.hook_impl import BuiltinImpl from bub.channels.message import ChannelMessage, MediaItem -from bub.channels.telegram import TelegramChannel, _extract_media_items +from bub.channels.telegram import TelegramChannel, TelegramMessageParser, _extract_media_items from bub.framework import BubFramework # --------------------------------------------------------------------------- @@ -30,6 +30,13 @@ async def fetch_bytes() -> bytes: assert item.data_fetcher is fetch_bytes +@pytest.mark.asyncio +async def test_media_item_returns_none_when_fetcher_skips_download() -> None: + item = MediaItem(type="video", mime_type="video/mp4", data_fetcher=_async_return(None)) + + assert await item.get_url() is None + + def test_channel_message_from_batch_merges_media() -> None: m1 = ChannelMessage( session_id="s", @@ -129,6 +136,51 @@ def test_extract_media_items_from_video_metadata() -> None: assert items[0].type == "video" +@pytest.mark.asyncio +async def test_telegram_video_parser_defaults_to_mp4_mime_type() -> None: + parser = TelegramMessageParser() + message = SimpleNamespace( + caption=None, + video=SimpleNamespace( + file_id="vid", + file_size=None, + width=640, + height=480, + duration=3, + mime_type=None, + ), + ) + + content, media = await parser._parse_video(message) # type: ignore[arg-type] + + assert content == "[Video: 3s]" + assert media is not None + assert media["mime_type"] == "video/mp4" + assert callable(media["data_fetcher"]) + + +@pytest.mark.asyncio +async def test_telegram_audio_parser_defaults_to_mpeg_mime_type() -> None: + parser = TelegramMessageParser() + message = SimpleNamespace( + audio=SimpleNamespace( + file_id="aud", + file_size=None, + duration=3, + mime_type=None, + title=None, + performer=None, + ), + ) + + content, media = await parser._parse_audio(message) # type: ignore[arg-type] + + assert content == "[Audio: Unknown (3s)]" + assert media is not None + assert media["mime_type"] == "audio/mpeg" + assert callable(media["data_fetcher"]) + + def test_extract_media_items_from_document_metadata() -> None: metadata = { "type": "document", @@ -301,18 +353,84 @@ async def test_build_prompt_with_multiple_images(tmp_path: Path) -> None: @pytest.mark.asyncio -async def test_build_prompt_with_non_image_media_only_includes_text(tmp_path: Path) -> None: +async def test_build_prompt_returns_video_url_part_with_video_media(tmp_path: Path) -> None: + _, impl = _build_impl(tmp_path) + message = ChannelMessage( + session_id="s", + channel="tg", + content="describe this video", + media=[MediaItem(type="video", mime_type="video/mp4", data_fetcher=_async_return(b"video"))], + ) + + result = await impl.build_prompt(message, session_id="s", state={}) + + assert isinstance(result, list) + assert len(result) == 2 + assert result[0]["type"] == "text" + assert "describe this video" in result[0]["text"] + expected = base64.b64encode(b"video").decode("utf-8") + assert result[1] == { + "type": "video_url", + "video_url": {"url": f"data:video/mp4;base64,{expected}"}, + } + + +@pytest.mark.asyncio +async def test_build_prompt_skips_video_when_download_is_too_large(tmp_path: Path) -> None: + _, impl = _build_impl(tmp_path) + message = ChannelMessage( + session_id="s", + channel="tg", + content="describe this video", + media=[MediaItem(type="video", mime_type="video/mp4", data_fetcher=_async_return(None))], + ) + + result = await impl.build_prompt(message, session_id="s", state={}) + + assert isinstance(result, str) + assert "describe this video" in result + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mime_type", "expected_format"), + [("audio/mpeg", "mp3"), ("audio/ogg", "ogg"), ("audio/x-wav", "wav")], +) +async def test_build_prompt_returns_input_audio_part(tmp_path: Path, mime_type: str, expected_format: str) -> None: + _, impl = _build_impl(tmp_path) + message = ChannelMessage( + session_id="s", + channel="tg", + content="listen to this", + media=[MediaItem(type="audio", mime_type=mime_type, data_fetcher=_async_return(b"audio"))], + ) + + result = await impl.build_prompt(message, session_id="s", state={}) + + assert isinstance(result, list) + assert result[0]["type"] == "text" + assert "listen to this" in result[0]["text"] + assert result[1] == { + "type": "input_audio", + "input_audio": { + "data": base64.b64encode(b"audio").decode("utf-8"), + "format": expected_format, + }, + } + + +@pytest.mark.asyncio +async def test_build_prompt_skips_remote_audio_url(tmp_path: Path) -> None: _, impl = _build_impl(tmp_path) message = ChannelMessage( session_id="s", channel="tg", content="listen to this", - media=[MediaItem(type="audio", mime_type="audio/ogg", data_fetcher=_async_return(b"\xff\xfb"))], + media=[MediaItem(type="audio", mime_type="audio/ogg", url="https://example.com/audio.ogg")], ) result = await impl.build_prompt(message, session_id="s", state={}) - # Non-image media: only returns a text assert isinstance(result, str) assert "listen to this" in result