Spaces:
Running
Running
File size: 7,247 Bytes
c47ec10 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | 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()
|