Merge pull request #212 from thynomex/codex/fix-issue-208-multi-image

services/openai_backend_api.py

@@ -53,6 +53,10 @@ SEARCH_POLL_INTERVAL_SECS = 3.0
 SEARCH_DONE_STATUS = {"finished_successfully", "finished_partial_completion"}
 SEARCH_CONVERSATION_ID_RE = re.compile(r'"conversation_id"\s*:\s*"([^"]+)"')
 SEARCH_URL_RE = re.compile(r"https?://[^\s\"'<>）)\]}]+")
+FILE_SERVICE_ID_RE = re.compile(r"file-service://([A-Za-z0-9_-]+)")
+FILE_ID_RE = re.compile(r"\b(file[-_](?!service\b)[A-Za-z0-9_-]+)\b")
+SEDIMENT_ID_RE = re.compile(r"sediment://([A-Za-z0-9_-]+)")
+IMAGE_POLL_SETTLE_SECS = 2.0
 CODEX_RESPONSES_INSTRUCTIONS = (
     "Use the image_generation tool to create exactly one image for the user's request. "
     "Return the generated image result."
@@ -1099,39 +1103,78 @@ class OpenAIBackendAPI:
     def _clean_search_url(self, value: Any) -> str:
         return str(value or "").strip().rstrip(".,;，。；")
 
+    @staticmethod
+    def _add_unique(values: list[str], candidates: list[str]) -> None:
+        for candidate in candidates:
+            if candidate and candidate not in values:
+                values.append(candidate)
+
+    @classmethod
+    def _extract_image_reference_ids(cls, payload: Any) -> tuple[list[str], list[str]]:
+        file_ids: list[str] = []
+        sediment_ids: list[str] = []
+
+        def walk(value: Any) -> None:
+            if isinstance(value, str):
+                cls._add_unique(file_ids, FILE_SERVICE_ID_RE.findall(value))
+                cls._add_unique(file_ids, FILE_ID_RE.findall(value))
+                cls._add_unique(sediment_ids, SEDIMENT_ID_RE.findall(value))
+                return
+            if isinstance(value, dict):
+                for item in value.values():
+                    walk(item)
+                return
+            if isinstance(value, list):
+                for item in value:
+                    walk(item)
+
+        walk(payload)
+        return file_ids, sediment_ids
+
+    @classmethod
+    def _has_image_asset_pointer(cls, payload: Any) -> bool:
+        if isinstance(payload, dict):
+            if str(payload.get("content_type") or "") == "image_asset_pointer":
+                return True
+            asset_pointer = str(payload.get("asset_pointer") or "")
+            if asset_pointer.startswith(("file-service://", "sediment://")):
+                return True
+            return any(cls._has_image_asset_pointer(item) for item in payload.values())
+        if isinstance(payload, list):
+            return any(cls._has_image_asset_pointer(item) for item in payload)
+        return False
+
     def _extract_image_tool_records(self, data: Dict[str, Any]) -> list[Dict[str, Any]]:
         """从 conversation 明细里提取图片工具输出记录。"""
         mapping = data.get("mapping") or {}
-        file_pat = re.compile(r"file-service://([A-Za-z0-9_-]+)")
-        sed_pat = re.compile(r"sediment://([A-Za-z0-9_-]+)")
         records = []
         for message_id, node in mapping.items():
             message = (node or {}).get("message") or {}
             author = message.get("author") or {}
             metadata = message.get("metadata") or {}
             content = message.get("content") or {}
-            if author.get("role") != "tool":
+            role = str(author.get("role") or "").strip().lower()
+            if role not in {"tool", "assistant"}:
                 continue
-            if content.get("content_type") != "multimodal_text":
+            is_image_gen = metadata.get("async_task_type") == "image_gen"
+            has_asset_pointer = self._has_image_asset_pointer(content) or self._has_image_asset_pointer(metadata)
+            if role == "assistant" and not (is_image_gen or has_asset_pointer):
                 continue
-            file_ids, sediment_ids = [], []
-            for part in content.get("parts") or []:
-                text = (part.get("asset_pointer") or "") if isinstance(part, dict) else (
-                    part if isinstance(part, str) else "")
-                for hit in file_pat.findall(text):
-                    if hit not in file_ids:
-                        file_ids.append(hit)
-                for hit in sed_pat.findall(text):
-                    if hit not in sediment_ids:
-                        sediment_ids.append(hit)
-            if metadata.get("async_task_type") != "image_gen" and not file_ids and not sediment_ids:
+            file_ids, sediment_ids = self._extract_image_reference_ids({"content": content, "metadata": metadata})
+            if not is_image_gen and not has_asset_pointer and not file_ids and not sediment_ids:
                 continue
             records.append(
                 {"message_id": message_id, "create_time": message.get("create_time") or 0, "file_ids": file_ids,
                  "sediment_ids": sediment_ids})
         return sorted(records, key=lambda item: item["create_time"])
 
-    def _poll_image_results(self, conversation_id: str, timeout_secs: float = 120.0) -> tuple[list[str], list[str]]:
+    def _poll_image_results(
+            self,
+            conversation_id: str,
+            timeout_secs: float = 120.0,
+            initial_file_ids: list[str] | None = None,
+            initial_sediment_ids: list[str] | None = None,
+    ) -> tuple[list[str], list[str]]:
         """Poll the conversation document until image file ids appear or budget runs out.
 
         - Sleeps image_poll_initial_wait_secs first (default 10s, +jitter). ChatGPT
@@ -1147,18 +1190,32 @@ class OpenAIBackendAPI:
         attempt = 0
         interval = float(config.image_poll_interval_secs)
         initial_wait = float(config.image_poll_initial_wait_secs)
+        file_ids: list[str] = []
+        sediment_ids: list[str] = []
+        self._add_unique(file_ids, initial_file_ids or [])
+        self._add_unique(sediment_ids, initial_sediment_ids or [])
+        has_initial_ids = bool(file_ids or sediment_ids)
+        last_hit_key: tuple[tuple[str, ...], tuple[str, ...]] | None = (
+            (tuple(file_ids), tuple(sediment_ids)) if has_initial_ids else None
+        )
         logger.info({
             "event": "image_poll_start",
             "conversation_id": conversation_id,
             "timeout_secs": timeout_secs,
             "initial_wait_secs": initial_wait,
             "interval_secs": interval,
+            "initial_file_ids": file_ids,
+            "initial_sediment_ids": sediment_ids,
         })
 
         def _remaining() -> float:
             return timeout_secs - (time.time() - start)
 
-        if initial_wait > 0:
+        if has_initial_ids:
+            settle_for = min(IMAGE_POLL_SETTLE_SECS, max(0.0, _remaining()))
+            if settle_for > 0:
+                time.sleep(settle_for)
+        elif initial_wait > 0:
             jitter = random.uniform(0, min(2.0, initial_wait * 0.2))
             sleep_for = min(initial_wait + jitter, max(0.0, _remaining()))
             if sleep_for > 0:
@@ -1202,7 +1259,6 @@ class OpenAIBackendAPI:
                     continue
                 break
 
-            file_ids, sediment_ids = [], []
             for record in self._extract_image_tool_records(conversation):
                 for file_id in record["file_ids"]:
                     if file_id not in file_ids:
@@ -1212,14 +1268,20 @@ class OpenAIBackendAPI:
                         sediment_ids.append(sediment_id)
             logger.debug({"event": "image_poll_check", "conversation_id": conversation_id, "attempt": attempt,
                           "file_ids": file_ids, "sediment_ids": sediment_ids})
-            if file_ids:
-                logger.info({"event": "image_poll_hit", "conversation_id": conversation_id, "file_ids": file_ids,
-                             "sediment_ids": sediment_ids})
+            if file_ids or sediment_ids:
+                hit_key = (tuple(file_ids), tuple(sediment_ids))
+                if last_hit_key == hit_key:
+                    logger.info({"event": "image_poll_hit", "conversation_id": conversation_id, "file_ids": file_ids,
+                                 "sediment_ids": sediment_ids})
+                    return file_ids, sediment_ids
+                last_hit_key = hit_key
+                logger.info({"event": "image_poll_hit_pending_settle", "conversation_id": conversation_id,
+                             "file_ids": file_ids, "sediment_ids": sediment_ids})
+                wait = min(IMAGE_POLL_SETTLE_SECS, max(0.0, _remaining()))
+                if wait > 0:
+                    time.sleep(wait)
+                    continue
                 return file_ids, sediment_ids
-            if sediment_ids:
-                logger.info({"event": "image_poll_hit", "conversation_id": conversation_id, "file_ids": [],
-                             "sediment_ids": sediment_ids})
-                return [], sediment_ids
             logger.debug({"event": "image_poll_wait", "conversation_id": conversation_id,
                           "elapsed_secs": round(time.time() - start, 1)})
             wait = min(interval, max(0.0, _remaining()))
@@ -1282,7 +1344,8 @@ class OpenAIBackendAPI:
                 })
                 continue
             if url:
-                urls.append(url)
+                if url not in urls:
+                    urls.append(url)
             else:
                 logger.debug({
                     "event": "image_download_url_empty",
@@ -1290,7 +1353,7 @@ class OpenAIBackendAPI:
                     "conversation_id": conversation_id,
                     "id": file_id,
                 })
-        if urls or not conversation_id:
+        if not conversation_id or not sediment_ids:
             logger.debug({
                 "event": "image_urls_resolved",
                 "conversation_id": conversation_id,
@@ -1312,7 +1375,8 @@ class OpenAIBackendAPI:
                 })
                 continue
             if url:
-                urls.append(url)
+                if url not in urls:
+                    urls.append(url)
             else:
                 logger.debug({
                     "event": "image_download_url_empty",
@@ -1338,12 +1402,42 @@ class OpenAIBackendAPI:
     ) -> list[str]:
         file_ids = [item for item in file_ids if item != "file_upload"]
         sediment_ids = list(sediment_ids)
-        if poll and conversation_id and not file_ids and not sediment_ids:
-            logger.info({"event": "image_resolve_poll_needed", "conversation_id": conversation_id})
-            polled_file_ids, polled_sediment_ids = self._poll_image_results(conversation_id,
-                                                                            config.image_poll_timeout_secs)
-            file_ids.extend(item for item in polled_file_ids if item and item not in file_ids)
-            sediment_ids.extend(item for item in polled_sediment_ids if item and item not in sediment_ids)
+        if poll and conversation_id:
+            logger.info({
+                "event": "image_resolve_poll_needed",
+                "conversation_id": conversation_id,
+                "initial_file_ids": file_ids,
+                "initial_sediment_ids": sediment_ids,
+            })
+            try:
+                polled_file_ids, polled_sediment_ids = self._poll_image_results(
+                    conversation_id,
+                    config.image_poll_timeout_secs,
+                    file_ids,
+                    sediment_ids,
+                )
+            except ImagePollTimeoutError:
+                if not file_ids and not sediment_ids:
+                    raise
+                logger.warning({
+                    "event": "image_resolve_poll_partial_timeout",
+                    "conversation_id": conversation_id,
+                    "file_ids": file_ids,
+                    "sediment_ids": sediment_ids,
+                })
+            except Exception as exc:
+                if not file_ids and not sediment_ids:
+                    raise
+                logger.warning({
+                    "event": "image_resolve_poll_partial_error",
+                    "conversation_id": conversation_id,
+                    "file_ids": file_ids,
+                    "sediment_ids": sediment_ids,
+                    "error": repr(exc),
+                })
+            else:
+                file_ids.extend(item for item in polled_file_ids if item and item not in file_ids)
+                sediment_ids.extend(item for item in polled_sediment_ids if item and item not in sediment_ids)
         return self._resolve_image_urls(conversation_id, file_ids, sediment_ids)
 
     def download_image_bytes(self, urls: list[str]) -> list[bytes]:
@@ -1351,7 +1445,8 @@ class OpenAIBackendAPI:
         for url in urls:
             response = self.session.get(url, timeout=120)
             ensure_ok(response, "image_download")
-            images.append(response.content)
+            if response.content not in images:
+                images.append(response.content)
         return images
 
     def stream_conversation(

services/protocol/conversation.py

@@ -429,12 +429,19 @@ def add_unique(values: list[str], candidates: list[str]) -> None:
             values.append(candidate)
 
 
+FILE_SERVICE_ID_RE = re.compile(r"file-service://([A-Za-z0-9_-]+)")
+FILE_ID_RE = re.compile(r"\b(file[-_](?!service\b)[A-Za-z0-9_-]+)\b")
+SEDIMENT_ID_RE = re.compile(r"sediment://([A-Za-z0-9_-]+)")
+
+
 def extract_conversation_ids(payload: str) -> tuple[str, list[str], list[str]]:
     conversation_match = re.search(r'"conversation_id"\s*:\s*"([^"]+)"', payload)
     conversation_id = conversation_match.group(1) if conversation_match else ""
+    file_ids: list[str] = []
     # Negative lookahead excludes "file-service" (URI prefix, not a real id).
-    file_ids = re.findall(r"(file[-_](?!service\b)[A-Za-z0-9]+)", payload)
-    sediment_ids = re.findall(r"sediment://([A-Za-z0-9_-]+)", payload)
+    add_unique(file_ids, FILE_SERVICE_ID_RE.findall(payload))
+    add_unique(file_ids, FILE_ID_RE.findall(payload))
+    sediment_ids = SEDIMENT_ID_RE.findall(payload)
     return conversation_id, file_ids, sediment_ids
 
 

services/protocol/openai_v1_response.py

@@ -294,14 +294,14 @@ def stream_image_response(
             continue
         items = image_output_items(prompt, output.data)
         if items:
-            item = items[0]
             usage = image_usage(
                 input_text_tokens=count_text_tokens(prompt, model),
                 input_image_tokens=input_image_tokens,
                 output_tokens=count_image_output_items_tokens(output.data, size, quality),
             )
-            yield {"type": "response.output_item.done", "output_index": 0, "item": item}
-            yield response_completed(response_id, model, created, [item], usage)
+            for output_index, item in enumerate(items):
+                yield {"type": "response.output_item.done", "output_index": output_index, "item": item}
+            yield response_completed(response_id, model, created, items, usage)
             return
     raise RuntimeError("image generation failed")
 

test/test_multi_image_results.py

@@ -0,0 +1,180 @@
+from __future__ import annotations
+
+import base64
+import unittest
+from unittest import mock
+
+from services.config import config
+from services.openai_backend_api import OpenAIBackendAPI
+from services.protocol.conversation import ImageOutput, extract_conversation_ids
+from services.protocol.openai_v1_response import stream_image_response
+
+
+def _conversation(file_ids: list[str], sediment_ids: list[str] | None = None) -> dict:
+    parts: list[object] = [
+        {"content_type": "image_asset_pointer", "asset_pointer": f"file-service://{file_id}"}
+        for file_id in file_ids
+    ]
+    parts.extend(f"sediment://{sediment_id}" for sediment_id in (sediment_ids or []))
+    return {
+        "mapping": {
+            "tool": {
+                "message": {
+                    "author": {"role": "tool"},
+                    "create_time": 1,
+                    "metadata": {"async_task_type": "image_gen"},
+                    "content": {"content_type": "multimodal_text", "parts": parts},
+                }
+            }
+        }
+    }
+
+
+class FakeBackend(OpenAIBackendAPI):
+    def __init__(self, conversations: list[dict] | None = None) -> None:
+        self.conversations = conversations or []
+        self.calls = 0
+        self.file_urls: dict[str, str] = {}
+        self.sediment_urls: dict[str, str] = {}
+
+    def _get_conversation(self, conversation_id: str) -> dict:
+        self.calls += 1
+        index = min(self.calls - 1, len(self.conversations) - 1)
+        return self.conversations[index]
+
+    def _get_file_download_url(self, file_id: str) -> str:
+        return self.file_urls.get(file_id, "")
+
+    def _get_attachment_download_url(self, conversation_id: str, attachment_id: str) -> str:
+        return self.sediment_urls.get(attachment_id, "")
+
+
+class MultiImageResultTests(unittest.TestCase):
+    def test_stream_id_extractor_keeps_full_file_ids(self) -> None:
+        payload = (
+            '{"conversation_id":"conv-1"} '
+            'file-service://file-first_123-extra sediment://sed-second_456-extra'
+        )
+
+        conversation_id, file_ids, sediment_ids = extract_conversation_ids(payload)
+
+        self.assertEqual(conversation_id, "conv-1")
+        self.assertEqual(file_ids, ["file-first_123-extra"])
+        self.assertEqual(sediment_ids, ["sed-second_456-extra"])
+
+    def test_conversation_record_extractor_finds_all_generated_assets(self) -> None:
+        backend = FakeBackend()
+        conversation = {
+            "mapping": {
+                "user": {
+                    "message": {
+                        "author": {"role": "user"},
+                        "content": {"parts": ["file-service://file-user-input"]},
+                    }
+                },
+                "tool": {
+                    "message": {
+                        "author": {"role": "tool"},
+                        "create_time": 1,
+                        "metadata": {
+                            "async_task_type": "image_gen",
+                            "nested": {"asset": "file-service://file-second"},
+                        },
+                        "content": {
+                            "content_type": "text",
+                            "parts": [
+                                {"content_type": "image_asset_pointer", "asset_pointer": "file-service://file-first"},
+                                "sediment://sed-first",
+                            ],
+                        },
+                    }
+                },
+                "assistant": {
+                    "message": {
+                        "author": {"role": "assistant"},
+                        "create_time": 2,
+                        "metadata": {},
+                        "content": {
+                            "parts": [
+                                {"content_type": "image_asset_pointer", "asset_pointer": "file-service://file-third"}
+                            ]
+                        },
+                    }
+                },
+            }
+        }
+
+        records = backend._extract_image_tool_records(conversation)
+        file_ids = [file_id for record in records for file_id in record["file_ids"]]
+        sediment_ids = [sediment_id for record in records for sediment_id in record["sediment_ids"]]
+
+        self.assertEqual(file_ids, ["file-first", "file-second", "file-third"])
+        self.assertEqual(sediment_ids, ["sed-first"])
+
+    def test_poll_waits_for_generated_asset_ids_to_settle(self) -> None:
+        backend = FakeBackend([
+            _conversation(["file-one"]),
+            _conversation(["file-one", "file-two"], ["sed-one"]),
+            _conversation(["file-one", "file-two"], ["sed-one"]),
+        ])
+
+        with (
+            mock.patch.dict(config.data, {"image_poll_initial_wait_secs": 0, "image_poll_interval_secs": 0.5}),
+            mock.patch("services.openai_backend_api.time.sleep", lambda _seconds: None),
+        ):
+            file_ids, sediment_ids = backend._poll_image_results("conv-1", timeout_secs=10)
+
+        self.assertEqual(file_ids, ["file-one", "file-two"])
+        self.assertEqual(sediment_ids, ["sed-one"])
+        self.assertEqual(backend.calls, 3)
+
+    def test_resolver_uses_file_and_sediment_urls(self) -> None:
+        backend = FakeBackend()
+        backend.file_urls = {"file-one": "https://files.test/one.png"}
+        backend.sediment_urls = {
+            "sed-one": "https://attachments.test/one.png",
+            "sed-two": "https://attachments.test/two.png",
+        }
+
+        urls = backend._resolve_image_urls("conv-1", ["file-one"], ["sed-one", "sed-two"])
+
+        self.assertEqual(urls, [
+            "https://files.test/one.png",
+            "https://attachments.test/one.png",
+            "https://attachments.test/two.png",
+        ])
+
+    def test_resolver_keeps_stream_ids_when_poll_extension_fails(self) -> None:
+        backend = FakeBackend()
+        backend.file_urls = {"file-one": "https://files.test/one.png"}
+        backend._get_conversation = mock.Mock(side_effect=RuntimeError("poll failed"))
+
+        with mock.patch("services.openai_backend_api.time.sleep", lambda _seconds: None):
+            urls = backend.resolve_conversation_image_urls("conv-1", ["file-one"], [], poll=True)
+
+        self.assertEqual(urls, ["https://files.test/one.png"])
+
+    def test_responses_stream_emits_all_image_output_items(self) -> None:
+        first = base64.b64encode(b"first").decode("ascii")
+        second = base64.b64encode(b"second").decode("ascii")
+        events = list(stream_image_response(
+            [ImageOutput(
+                kind="result",
+                model="gpt-image-2",
+                index=1,
+                total=1,
+                data=[{"b64_json": first}, {"b64_json": second}],
+            )],
+            "draw two options",
+            "gpt-image-2",
+        ))
+
+        done_events = [event for event in events if event.get("type") == "response.output_item.done"]
+        completed = next(event["response"] for event in events if event.get("type") == "response.completed")
+
+        self.assertEqual([event["output_index"] for event in done_events], [0, 1])
+        self.assertEqual([item["result"] for item in completed["output"]], [first, second])
+
+
+if __name__ == "__main__":
+    unittest.main()