diff --git a/astrbot/builtin_stars/astrbot/group_chat_context.py b/astrbot/builtin_stars/astrbot/group_chat_context.py index 6e832d4cb0..f88b2c18f0 100644 --- a/astrbot/builtin_stars/astrbot/group_chat_context.py +++ b/astrbot/builtin_stars/astrbot/group_chat_context.py @@ -47,6 +47,7 @@ def __init__(self, acm: AstrBotConfigManager, context: star.Context) -> None: self._locks: dict[str, asyncio.Lock] = {} self.raw_records: dict[str, deque[str]] = defaultdict(deque) self._record_ids: dict[str, deque[str]] = defaultdict(deque) + self._caption_tasks: set[asyncio.Task[None]] = set() def _get_lock(self, umo: str) -> asyncio.Lock: lock = self._locks.get(umo) @@ -139,26 +140,66 @@ async def remove_session(self, event: AstrMessageEvent) -> int: self._locks.pop(umo, None) return cnt - async def handle_message(self, event: AstrMessageEvent) -> None: + async def handle_message( + self, + event: AstrMessageEvent, + *, + caption_images: bool = True, + ) -> None: + """Record a group message without blocking on image caption requests. + + Args: + event: Incoming group message event. + caption_images: Whether passive group images should be captioned for + later context injection. Messages that already trigger an LLM + reply should set this to ``False`` because the main request + pipeline handles their images separately. + """ if event.get_message_type() != MessageType.GROUP_MESSAGE: return umo = event.unified_msg_origin cfg = self.cfg(event) - final_message = await self._format_message(event, cfg) + final_message, caption_template, pending_images = self._format_message( + event, + cfg, + caption_images=caption_images, + ) + record_id = uuid.uuid4().hex async with self._get_lock(umo): records = self.raw_records[umo] record_ids = self._record_ids[umo] - record_id = uuid.uuid4().hex records.append(final_message) record_ids.append(record_id) _trim_left(records, cfg["group_message_max_cnt"], record_ids) event.set_extra("_group_context_record_id", record_id) event.set_extra("_group_context_raw_idx", len(records) - 1) + if pending_images: + task = asyncio.create_task( + self._fill_image_captions( + umo=umo, + record_id=record_id, + caption_template=caption_template, + pending_images=pending_images, + provider_id=cfg["image_caption_provider_id"], + prompt=cfg["image_caption_prompt"], + ) + ) + self._caption_tasks.add(task) + task.add_done_callback(self._on_caption_task_done) + logger.debug(f"group_chat_context | {umo} | {final_message}") + def _on_caption_task_done(self, task: asyncio.Task[None]) -> None: + """Release a completed caption task and expose unexpected failures.""" + self._caption_tasks.discard(task) + if task.cancelled(): + return + if exc := task.exception(): + logger.error("Group image caption task failed.", exc_info=exc) + async def on_req_llm(self, event: AstrMessageEvent, req: ProviderRequest) -> None: umo = event.unified_msg_origin record_id = event.get_extra("_group_context_record_id", None) @@ -196,29 +237,75 @@ async def on_req_llm(self, event: AstrMessageEvent, req: ProviderRequest) -> Non TextPart(text=_format_group_history_block(records_to_inject)) ) - async def _format_message(self, event: AstrMessageEvent, cfg: dict) -> str: + async def _fill_image_captions( + self, + *, + umo: str, + record_id: str, + caption_template: str, + pending_images: list[tuple[str, str]], + provider_id: str, + prompt: str, + ) -> None: + """Resolve image captions in the background and update one record.""" + results = await asyncio.gather( + *( + self.get_image_caption(image_url, provider_id, prompt) + for _, image_url in pending_images + ), + return_exceptions=True, + ) + + resolved_message = caption_template + for (marker, _), result in zip(pending_images, results, strict=True): + if isinstance(result, BaseException): + logger.error("Failed to get image caption: %s", result) + replacement = " [Image]" + else: + replacement = f" [Image: {result}]" + resolved_message = resolved_message.replace(marker, replacement, 1) + + async with self._get_lock(umo): + record_ids = self._record_ids.get(umo) + records = self.raw_records.get(umo) + if not record_ids or not records or record_id not in record_ids: + return + record_index = record_ids.index(record_id) + records[record_index] = resolved_message + + logger.debug(f"group_chat_context captioned | {umo} | {resolved_message}") + + def _format_message( + self, + event: AstrMessageEvent, + cfg: dict, + *, + caption_images: bool, + ) -> tuple[str, str, list[tuple[str, str]]]: + """Format one record and prepare optional background image captions.""" datetime_str = datetime.datetime.now().strftime("%H:%M:%S") - parts = [f"[{event.message_obj.sender.nickname}/{datetime_str}]: "] + prefix = f"[{event.message_obj.sender.nickname}/{datetime_str}]: " + parts = [prefix] + template_parts = [prefix] + pending_images: list[tuple[str, str]] = [] for comp in event.get_messages(): if isinstance(comp, Plain): - parts.append(f" {comp.text}") + text = f" {comp.text}" + parts.append(text) + template_parts.append(text) elif isinstance(comp, Image): - if cfg["image_caption"]: - try: - url = comp.url if comp.url else comp.file - if not url: - raise Exception("图片 URL 为空") - caption = await self.get_image_caption( - url, - cfg["image_caption_provider_id"], - cfg["image_caption_prompt"], - ) - parts.append(f" [Image: {caption}]") - except Exception as e: - logger.error(f"获取图片描述失败: {e}") + url = comp.url if comp.url else comp.file + should_caption = caption_images and cfg["image_caption"] and bool(url) + parts.append(" [Image]") + if should_caption: + marker = f" __ASTRBOT_IMAGE_CAPTION_{uuid.uuid4().hex}__" + template_parts.append(marker) + pending_images.append((marker, url)) else: - parts.append(" [Image]") + template_parts.append(" [Image]") + if caption_images and cfg["image_caption"] and not url: + logger.error("Failed to get image caption: image URL is empty.") elif isinstance(comp, Json): card_data = comp.data if isinstance(card_data, dict) and isinstance( @@ -249,7 +336,9 @@ async def _format_message(self, event: AstrMessageEvent, cfg: dict) -> str: normalized = " ".join(value.split()) fields.append(f"{label}: {_truncate_reply_text(normalized)}") suffix = f": {'; '.join(fields)}" if fields else "" - parts.append(f" [Shared Card{suffix}]") + text = f" [Shared Card{suffix}]" + parts.append(text) + template_parts.append(text) elif isinstance(comp, At): is_at_self = str(comp.qq) in ( event.get_self_id(), @@ -257,19 +346,25 @@ async def _format_message(self, event: AstrMessageEvent, cfg: dict) -> str: ) if is_at_self: parts.insert(1, "⚠️[DIRECTED AT YOU] ") - parts.append(f" [At: {comp.name}]") + template_parts.insert(1, "⚠️[DIRECTED AT YOU] ") + text = f" [At: {comp.name}]" + parts.append(text) + template_parts.append(text) elif isinstance(comp, Reply): if comp.message_str: - parts.append( - f" [Quote({comp.sender_nickname}: {_truncate_reply_text(comp.message_str)})]" + text = ( + f" [Quote({comp.sender_nickname}: " + f"{_truncate_reply_text(comp.message_str)})]" ) elif comp.chain: chain_desc = _describe_chain(comp.chain) - parts.append(f" [Quote({comp.sender_nickname}: {chain_desc})]") + text = f" [Quote({comp.sender_nickname}: {chain_desc})]" else: - parts.append(" [Quote]") + text = " [Quote]" + parts.append(text) + template_parts.append(text) - return "".join(parts) + return "".join(parts), "".join(template_parts), pending_images _MAX_REPLY_TEXT_LENGTH = 200 diff --git a/astrbot/builtin_stars/astrbot/main.py b/astrbot/builtin_stars/astrbot/main.py index cc2c368a23..86bd9f375d 100644 --- a/astrbot/builtin_stars/astrbot/main.py +++ b/astrbot/builtin_stars/astrbot/main.py @@ -222,7 +222,15 @@ async def on_message(self, event: AstrMessageEvent): # chat context that should be injected into future LLM requests. if not event.get_extra("handlers_parsed_params", {}): try: - await self.group_chat_context.handle_message(event) + # The main LLM pipeline already handles images from messages + # that trigger a reply. Captioning them here would add a + # duplicate foreground model request before the reply starts. + await self.group_chat_context.handle_message( + event, + caption_images=not ( + event.is_at_or_wake_command or need_active + ), + ) except BaseException as e: logger.error(e) diff --git a/tests/unit/test_group_chat_context_wiring.py b/tests/unit/test_group_chat_context_wiring.py index a103d4f913..ef3f6c6020 100644 --- a/tests/unit/test_group_chat_context_wiring.py +++ b/tests/unit/test_group_chat_context_wiring.py @@ -1,10 +1,11 @@ +import asyncio import json from types import SimpleNamespace from unittest.mock import AsyncMock, MagicMock import pytest -from astrbot.api.message_components import Json, Plain +from astrbot.api.message_components import Image, Json, Plain from astrbot.api.provider import LLMResponse from astrbot.builtin_stars.astrbot.group_chat_context import GroupChatContext from astrbot.builtin_stars.astrbot.main import Main @@ -37,6 +38,7 @@ def make_event( umo: str = "aiocqhttp:GroupMessage:user_123_group_456", *, handlers_parsed_params: dict | None = None, + is_at_or_wake_command: bool = False, ): event = MagicMock() event.unified_msg_origin = umo @@ -44,6 +46,7 @@ def make_event( event.message_obj = SimpleNamespace(message=[Plain("hello")]) event.message_str = "hello" event.session_id = "session-1" + event.is_at_or_wake_command = is_at_or_wake_command store, get_extra, set_extra = _make_extras_store() # Simulate WakingCheckStage output: an empty dict means no command matched. @@ -147,7 +150,10 @@ async def test_on_message_does_not_clear_group_context_on_first_enabled_message( pass main.group_chat_context.need_active_reply.assert_awaited_once_with(event) - main.group_chat_context.handle_message.assert_awaited_once_with(event) + main.group_chat_context.handle_message.assert_awaited_once_with( + event, + caption_images=True, + ) main.group_chat_context.remove_session.assert_not_called() @@ -172,7 +178,10 @@ async def test_on_message_records_json_card_and_checks_active_reply(): pass main.group_chat_context.need_active_reply.assert_awaited_once_with(event) - main.group_chat_context.handle_message.assert_awaited_once_with(event) + main.group_chat_context.handle_message.assert_awaited_once_with( + event, + caption_images=True, + ) @pytest.mark.asyncio @@ -279,9 +288,15 @@ async def test_format_message_summarizes_json_card(card_data, expected): event.message_obj = SimpleNamespace(sender=SimpleNamespace(nickname="Alice")) event.get_messages.return_value = [Json(data=card_data)] - formatted = await context._format_message(event, {}) + formatted, template, pending = context._format_message( + event, + {}, + caption_images=True, + ) assert formatted.endswith(expected) + assert template == formatted + assert pending == [] @pytest.mark.asyncio @@ -297,6 +312,160 @@ async def test_format_message_truncates_long_json_card_fields(): ) ] - formatted = await context._format_message(event, {}) + formatted, template, pending = context._format_message( + event, + {}, + caption_images=True, + ) assert f"Description: {'a' * 200}...]" in formatted + assert template == formatted + assert pending == [] + + +@pytest.mark.asyncio +async def test_on_message_skips_duplicate_caption_for_reply_triggering_image(): + main = Main.__new__(Main) + main.context = MagicMock() + main.context.get_config.return_value = { + "provider_ltm_settings": { + "group_icl_enable": True, + "active_reply": {"enable": False}, + }, + } + main.group_chat_context = SimpleNamespace( + need_active_reply=AsyncMock(return_value=False), + handle_message=AsyncMock(), + ) + event = make_event(is_at_or_wake_command=True) + event.message_obj.message = [Image(file="https://example.com/cat.png")] + + async for _ in main.on_message(event): + pass + + main.group_chat_context.handle_message.assert_awaited_once_with( + event, + caption_images=False, + ) + + +@pytest.mark.asyncio +async def test_group_image_caption_runs_after_message_is_recorded(): + context = MagicMock() + group_context = GroupChatContext(MagicMock(), context) + group_context.cfg = MagicMock( + return_value={ + "group_message_max_cnt": 100, + "image_caption": True, + "image_caption_prompt": "describe", + "image_caption_provider_id": "vision-provider", + } + ) + caption_started = asyncio.Event() + release_caption = asyncio.Event() + + async def delayed_caption(*_args, **_kwargs): + caption_started.set() + await release_caption.wait() + return "a cat" + + group_context.get_image_caption = AsyncMock(side_effect=delayed_caption) + event = MagicMock() + event.unified_msg_origin = "aiocqhttp:GroupMessage:user_1_group_1" + event.get_message_type.return_value = MessageType.GROUP_MESSAGE + event.message_obj = SimpleNamespace(sender=SimpleNamespace(nickname="alice")) + event.get_messages.return_value = [Image(file="https://example.com/cat.png")] + + await asyncio.wait_for(group_context.handle_message(event), timeout=1.0) + await asyncio.wait_for(caption_started.wait(), timeout=1.0) + + assert "[Image]" in group_context.raw_records[event.unified_msg_origin][0] + tasks = tuple(group_context._caption_tasks) + assert len(tasks) == 1 + + release_caption.set() + await asyncio.gather(*tasks) + + assert "[Image: a cat]" in group_context.raw_records[event.unified_msg_origin][0] + + +@pytest.mark.asyncio +async def test_reply_triggering_image_does_not_start_group_caption_task(): + context = MagicMock() + group_context = GroupChatContext(MagicMock(), context) + group_context.cfg = MagicMock( + return_value={ + "group_message_max_cnt": 100, + "image_caption": True, + "image_caption_prompt": "describe", + "image_caption_provider_id": "vision-provider", + } + ) + group_context.get_image_caption = AsyncMock(return_value="unused") + event = MagicMock() + event.unified_msg_origin = "aiocqhttp:GroupMessage:user_1_group_1" + event.get_message_type.return_value = MessageType.GROUP_MESSAGE + event.message_obj = SimpleNamespace(sender=SimpleNamespace(nickname="alice")) + event.get_messages.return_value = [Image(file="https://example.com/cat.png")] + + await group_context.handle_message(event, caption_images=False) + + group_context.get_image_caption.assert_not_awaited() + assert not group_context._caption_tasks + assert "[Image]" in group_context.raw_records[event.unified_msg_origin][0] + + +@pytest.mark.asyncio +async def test_group_image_captions_start_concurrently_across_groups(): + context = MagicMock() + group_context = GroupChatContext(MagicMock(), context) + group_context.cfg = MagicMock( + return_value={ + "group_message_max_cnt": 100, + "image_caption": True, + "image_caption_prompt": "describe", + "image_caption_provider_id": "vision-provider", + } + ) + started_urls: list[str] = [] + both_started = asyncio.Event() + release_captions = asyncio.Event() + + async def delayed_caption(image_url, *_args, **_kwargs): + started_urls.append(image_url) + if len(started_urls) == 2: + both_started.set() + await release_captions.wait() + return image_url + + group_context.get_image_caption = AsyncMock(side_effect=delayed_caption) + + def make_image_event(umo: str, image_url: str): + event = MagicMock() + event.unified_msg_origin = umo + event.get_message_type.return_value = MessageType.GROUP_MESSAGE + event.message_obj = SimpleNamespace(sender=SimpleNamespace(nickname="alice")) + event.get_messages.return_value = [Image(file=image_url)] + return event + + event_one = make_image_event("group-1", "https://example.com/one.png") + event_two = make_image_event("group-2", "https://example.com/two.png") + + await asyncio.wait_for( + asyncio.gather( + group_context.handle_message(event_one), + group_context.handle_message(event_two), + ), + timeout=1.0, + ) + await asyncio.wait_for(both_started.wait(), timeout=1.0) + + tasks = tuple(group_context._caption_tasks) + assert len(tasks) == 2 + release_captions.set() + await asyncio.gather(*tasks) + + assert set(started_urls) == { + "https://example.com/one.png", + "https://example.com/two.png", + }