from __future__ import annotations import unittest from unittest import mock import json import base64 from services.config import config from services.protocol import openai_v1_chat_complete, openai_v1_response from services.protocol.chat_completion_cache import chat_completion_cache from services.protocol.conversation import iter_conversation_payloads, sanitize_output_text from utils.helper import extract_image_from_message_content PNG_1X1 = base64.b64decode( "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAIAAACQd1PeAAAADUlEQVR4nGP8z8BQDwAFgwJ/luzl4wAAAABJRU5ErkJggg==" ) PNG_1X1_DATA_URL = "data:image/png;base64," + base64.b64encode(PNG_1X1).decode("ascii") class ChatCompletionCacheTests(unittest.TestCase): def setUp(self) -> None: self.old_cache_settings = config.data.get("chat_completion_cache") config.data["chat_completion_cache"] = { "enabled": True, "ttl_seconds": 60, "max_entries": 32, "dedupe_inflight": True, "stream_cache": True, "normalize_messages": True, "drop_adjacent_duplicates": True, "drop_assistant_history": False, } chat_completion_cache.clear() def tearDown(self) -> None: if self.old_cache_settings is None: config.data.pop("chat_completion_cache", None) else: config.data["chat_completion_cache"] = self.old_cache_settings chat_completion_cache.clear() def test_repeated_non_stream_text_completion_uses_cache(self) -> None: calls = 0 def fake_collect_text(_backend, _request): nonlocal calls calls += 1 return f"cached answer {calls}" body = { "model": "auto", "messages": [{"role": "user", "content": "cache this exact prompt"}], } with ( mock.patch("services.protocol.openai_v1_chat_complete.text_backend", return_value=object()), mock.patch("services.protocol.openai_v1_chat_complete.collect_text", side_effect=fake_collect_text), ): first = openai_v1_chat_complete.handle(body) second = openai_v1_chat_complete.handle(body) self.assertEqual(calls, 1) self.assertEqual( first["choices"][0]["message"]["content"], second["choices"][0]["message"]["content"], ) def test_repeated_stream_text_completion_replays_cached_chunks(self) -> None: calls = 0 def fake_stream_text_deltas(_backend, _request): nonlocal calls calls += 1 yield "streamed" yield " answer" body = { "model": "auto", "stream": True, "messages": [{"role": "user", "content": "stream cache this exact prompt"}], } with ( mock.patch("services.protocol.openai_v1_chat_complete.text_backend", return_value=object()), mock.patch( "services.protocol.openai_v1_chat_complete.stream_text_deltas", side_effect=fake_stream_text_deltas, ), ): first = list(openai_v1_chat_complete.handle(body)) second = list(openai_v1_chat_complete.handle(body)) self.assertEqual(calls, 1) self.assertEqual(first, second) content = "".join(str(chunk["choices"][0]["delta"].get("content") or "") for chunk in second) self.assertEqual(content, "streamed answer") def test_adjacent_duplicate_messages_are_removed_before_upstream_call(self) -> None: captured_messages = [] def fake_collect_text(_backend, request): captured_messages.extend(request.messages or []) return "ok" body = { "model": "auto", "messages": [ {"role": "user", "content": "repeat me"}, {"role": "user", "content": "repeat me"}, {"role": "assistant", "content": "old answer"}, {"role": "user", "content": "next prompt"}, ], } with ( mock.patch("services.protocol.openai_v1_chat_complete.text_backend", return_value=object()), mock.patch("services.protocol.openai_v1_chat_complete.collect_text", side_effect=fake_collect_text), ): openai_v1_chat_complete.handle(body) self.assertEqual( captured_messages, [ {"role": "user", "content": "repeat me"}, {"role": "assistant", "content": "old answer"}, {"role": "user", "content": "next prompt"}, ], ) def test_chat_completion_usage_includes_cached_tokens(self) -> None: with ( mock.patch("services.protocol.openai_v1_chat_complete.text_backend", return_value=object()), mock.patch("services.protocol.openai_v1_chat_complete.collect_text", return_value="ok"), ): response = openai_v1_chat_complete.handle({ "model": "auto", "messages": [{"role": "user", "content": "usage shape"}], }) details = response["usage"]["prompt_tokens_details"] self.assertEqual(details["cached_tokens"], 0) output_details = response["usage"]["completion_tokens_details"] self.assertEqual(output_details["reasoning_tokens"], 0) def test_responses_completed_usage_includes_cached_tokens(self) -> None: with ( mock.patch("services.protocol.openai_v1_response.text_backend", return_value=object()), mock.patch("services.protocol.openai_v1_response.stream_text_deltas", return_value=iter(["ok"])), ): response = openai_v1_response.handle({ "model": "auto", "input": "usage shape", }) details = response["usage"]["input_tokens_details"] self.assertEqual(details["cached_tokens"], 0) output_details = response["usage"]["output_tokens_details"] self.assertEqual(output_details["reasoning_tokens"], 0) def test_repeated_responses_text_request_uses_cache(self) -> None: calls = 0 def fake_stream_text_deltas(_backend, _request): nonlocal calls calls += 1 yield f"response cache {calls}" body = { "model": "auto", "input": "cache this responses prompt", "stream": True, } with ( mock.patch("services.protocol.openai_v1_response.text_backend", return_value=object()), mock.patch("services.protocol.openai_v1_response.stream_text_deltas", side_effect=fake_stream_text_deltas), ): first = list(openai_v1_response.handle(body)) second = list(openai_v1_response.handle(body)) self.assertEqual(calls, 1) self.assertEqual(first, second) def test_output_sanitizer_removes_chatgpt_annotation_markup(self) -> None: text = ( "Repo: \ue200url\ue202basketikun/chatgpt2api" "\ue202https://github.com/basketikun/chatgpt2api\ue201 " "details \ue200cite\ue202turn0search0\ue201." ) self.assertEqual( sanitize_output_text(text), "Repo: basketikun/chatgpt2api (https://github.com/basketikun/chatgpt2api) details.", ) def test_output_sanitizer_preserves_annotated_entity_text(self) -> None: text = ( "The character is from \ue200entity\ue202Invincible\ue201, " "which is based on the comic series \ue200entity\ue202Invincible\ue201." ) self.assertEqual( sanitize_output_text(text), "The character is from Invincible, which is based on the comic series Invincible.", ) def test_output_sanitizer_preserves_readable_cite_label(self) -> None: text = "The character is \ue200cite\ue202Invincible\ue202turn0search0\ue201." self.assertEqual(sanitize_output_text(text), "The character is Invincible.") def test_stream_sanitizer_does_not_emit_partial_annotation_or_repeat_prefix(self) -> None: events = [ {"p": "/message/content/parts/0", "o": "append", "v": "Repo: \ue200url\ue202chat"}, {"p": "/message/content/parts/0", "o": "append", "v": "gpt2api\ue202turn0search0\ue201 done \ue200cite\ue202turn0\ue201."}, "[DONE]", ] payloads = [json.dumps(event, ensure_ascii=False) if isinstance(event, dict) else event for event in events] deltas = [ str(event.get("delta") or "") for event in iter_conversation_payloads(iter(payloads)) if event.get("type") == "conversation.delta" ] self.assertEqual("".join(deltas), "Repo: chatgpt2api done.") self.assertFalse(any("\ue200" in delta or "\ue202" in delta or "\ue201" in delta for delta in deltas)) def test_responses_tools_add_honest_no_tool_guard(self) -> None: model, messages = openai_v1_response.text_response_parts({ "model": "auto", "input": "run echo hi", "tools": [{"type": "function", "name": "shell"}], }) self.assertEqual(model, "auto") self.assertEqual(messages[0]["role"], "system") self.assertIn("cannot execute local tools", str(messages[0]["content"])) def test_chat_completions_accepts_remote_image_url(self) -> None: class FakeImageResponse: status_code = 200 headers = {"content-type": "image/png", "content-length": str(len(PNG_1X1))} content = PNG_1X1 with mock.patch("utils.helper.requests.get", return_value=FakeImageResponse()) as request_get: model, messages = openai_v1_chat_complete.text_chat_parts({ "model": "auto", "messages": [{ "role": "user", "content": [ {"type": "text", "text": "Describe this"}, {"type": "image_url", "image_url": {"url": "https://example.test/image.png"}}, ], }], }) request_get.assert_called_once() self.assertEqual(model, "auto") content = messages[0]["content"] self.assertEqual(content[0], {"type": "text", "text": "Describe this"}) self.assertEqual(content[1]["type"], "image") self.assertEqual(content[1]["data"], PNG_1X1) self.assertEqual(content[1]["mime"], "image/png") def test_responses_text_request_preserves_input_image(self) -> None: captured = {} def fake_stream_text_deltas(_backend, request): captured["messages"] = request.messages yield "red" body = { "model": "auto", "input": [ {"type": "input_text", "text": "What color is this image?"}, {"type": "input_image", "image_url": PNG_1X1_DATA_URL}, ], } with ( mock.patch("services.protocol.openai_v1_response.text_backend", return_value=object()), mock.patch("services.protocol.openai_v1_response.stream_text_deltas", side_effect=fake_stream_text_deltas), ): response = openai_v1_response.handle(body) self.assertEqual(response["output"][0]["content"][0]["text"], "red") content = captured["messages"][0]["content"] self.assertEqual(content[0], {"type": "text", "text": "What color is this image?"}) self.assertEqual(content[1]["type"], "image") self.assertEqual(content[1]["mime"], "image/png") self.assertEqual(content[1]["data"], PNG_1X1) self.assertGreater(response["usage"]["input_tokens_details"]["image_tokens"], 0) def test_responses_text_request_accepts_remote_input_image_url(self) -> None: class FakeImageResponse: status_code = 200 headers = {"content-type": "image/png", "content-length": str(len(PNG_1X1))} content = PNG_1X1 with mock.patch("utils.helper.requests.get", return_value=FakeImageResponse()) as request_get: _model, messages = openai_v1_response.text_response_parts({ "model": "auto", "input": [{ "type": "message", "role": "user", "content": [ {"type": "input_text", "text": "Describe this"}, {"type": "input_image", "image_url": {"url": "https://example.test/image.png"}}, ], }], }) request_get.assert_called_once() content = messages[0]["content"] self.assertEqual(content[0], {"type": "text", "text": "Describe this"}) self.assertEqual(content[1]["type"], "image") self.assertEqual(content[1]["data"], PNG_1X1) self.assertEqual(content[1]["mime"], "image/png") def test_image_extractor_supports_extra_image_object_shapes(self) -> None: encoded = base64.b64encode(PNG_1X1).decode("ascii") images = extract_image_from_message_content([ {"type": "image", "data": PNG_1X1, "mime": "image/png"}, {"type": "input_image", "base64": encoded, "mime_type": "image/png"}, {"type": "input_image", "source": {"type": "base64", "data": encoded, "media_type": "image/png"}}, ]) self.assertEqual(len(images), 3) self.assertEqual([image[1] for image in images], ["image/png", "image/png", "image/png"]) self.assertTrue(all(image[0] == PNG_1X1 for image in images)) if __name__ == "__main__": unittest.main()