338 lines
11 KiB
Python
338 lines
11 KiB
Python
import base64
|
|
import builtins
|
|
import inspect
|
|
import types
|
|
import unittest
|
|
from unittest import mock
|
|
|
|
from PIL import Image
|
|
|
|
try:
|
|
from . import test_path
|
|
except ImportError:
|
|
import test_path
|
|
from ComfyUI_LLM_API import image_encoding
|
|
from ComfyUI_LLM_API.image_encoding import encode_comfy_image_batch_to_data_urls
|
|
from ComfyUI_LLM_API.llm_request import (
|
|
DEFAULT_MAX_RETRIES,
|
|
DEFAULT_REQUEST_TIMEOUT_SECONDS,
|
|
LLMAPIUserError,
|
|
build_chat_messages,
|
|
create_chat_completion,
|
|
create_openai_client,
|
|
extract_completion_text,
|
|
format_llm_error,
|
|
is_openai_error,
|
|
)
|
|
|
|
|
|
class FakeTensor:
|
|
def __init__(self, array):
|
|
self._array = array
|
|
|
|
def cpu(self):
|
|
return self
|
|
|
|
def numpy(self):
|
|
return self._array
|
|
|
|
|
|
class FakeMessage:
|
|
def __init__(self, content):
|
|
self.content = content
|
|
|
|
|
|
class FakeChoice:
|
|
def __init__(self, content):
|
|
self.message = FakeMessage(content)
|
|
|
|
|
|
class FakeCompletion:
|
|
def __init__(self, content):
|
|
self.choices = [FakeChoice(content)]
|
|
|
|
|
|
class FakeCompletions:
|
|
def __init__(self):
|
|
self.calls = []
|
|
self.fail_first_image_attempt = False
|
|
|
|
def create(self, **kwargs):
|
|
self.calls.append(kwargs)
|
|
if self.fail_first_image_attempt and len(self.calls) == 1:
|
|
raise RuntimeError("failed to process image")
|
|
return FakeCompletion("hello")
|
|
|
|
|
|
class FakeChat:
|
|
def __init__(self):
|
|
self.completions = FakeCompletions()
|
|
|
|
|
|
class FakeClient:
|
|
def __init__(self):
|
|
self.chat = FakeChat()
|
|
|
|
|
|
class LLMRequestTests(unittest.TestCase):
|
|
def test_builds_text_only_chat_messages(self):
|
|
messages = build_chat_messages("system", "prompt")
|
|
|
|
self.assertEqual(
|
|
messages,
|
|
[
|
|
{"role": "system", "content": "system"},
|
|
{"role": "user", "content": "prompt"},
|
|
],
|
|
)
|
|
|
|
def test_builds_chat_messages_with_multiple_images(self):
|
|
image_urls = [
|
|
"data:image/jpeg;base64,aaa",
|
|
"data:image/jpeg;base64,bbb",
|
|
]
|
|
|
|
messages = build_chat_messages("system", "describe", image_urls=image_urls)
|
|
|
|
self.assertEqual(messages[0], {"role": "system", "content": "system"})
|
|
user_content = messages[1]["content"]
|
|
self.assertEqual(user_content[0], {"type": "text", "text": "describe"})
|
|
self.assertEqual(
|
|
user_content[1:],
|
|
[
|
|
{"type": "image_url", "image_url": {"url": image_urls[0]}},
|
|
{"type": "image_url", "image_url": {"url": image_urls[1]}},
|
|
],
|
|
)
|
|
|
|
def test_create_chat_completion_does_not_forward_seed(self):
|
|
client = FakeClient()
|
|
|
|
result = create_chat_completion(
|
|
client=client,
|
|
model="demo-model",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
temperature=0.5,
|
|
has_images=False,
|
|
)
|
|
|
|
self.assertEqual(extract_completion_text(result), "hello")
|
|
self.assertEqual(client.chat.completions.calls[0]["model"], "demo-model")
|
|
self.assertNotIn("seed", client.chat.completions.calls[0])
|
|
self.assertNotIn("execution_seed", client.chat.completions.calls[0])
|
|
|
|
def test_image_completion_retry_rebuilds_messages_with_lower_limits(self):
|
|
client = FakeClient()
|
|
client.chat.completions.fail_first_image_attempt = True
|
|
limits = []
|
|
|
|
def build_messages(limit):
|
|
limits.append(limit)
|
|
return [{"role": "user", "content": f"limit {limit}"}]
|
|
|
|
result = create_chat_completion(
|
|
client=client,
|
|
model="demo-model",
|
|
messages=build_messages(100),
|
|
temperature=0.5,
|
|
has_images=True,
|
|
image_retry_message_factory=build_messages,
|
|
image_retry_limits=(100, 50),
|
|
)
|
|
|
|
self.assertEqual(extract_completion_text(result), "hello")
|
|
self.assertEqual(limits, [100, 100, 50])
|
|
self.assertEqual(len(client.chat.completions.calls), 2)
|
|
self.assertEqual(
|
|
client.chat.completions.calls[1]["messages"],
|
|
[{"role": "user", "content": "limit 50"}],
|
|
)
|
|
|
|
def test_image_completion_requires_explicit_retry_policy(self):
|
|
client = FakeClient()
|
|
|
|
with self.assertRaises(ValueError):
|
|
create_chat_completion(
|
|
client=client,
|
|
model="demo-model",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
temperature=0.5,
|
|
has_images=True,
|
|
image_retry_message_factory=None,
|
|
image_retry_limits=(100,),
|
|
)
|
|
|
|
with self.assertRaises(ValueError):
|
|
create_chat_completion(
|
|
client=client,
|
|
model="demo-model",
|
|
messages=[{"role": "user", "content": "hi"}],
|
|
temperature=0.5,
|
|
has_images=True,
|
|
image_retry_message_factory=lambda _limit: [],
|
|
image_retry_limits=(),
|
|
)
|
|
|
|
def test_extract_completion_text_rejects_missing_or_empty_content(self):
|
|
for content in (None, "", " "):
|
|
with self.subTest(content=content):
|
|
with self.assertRaises(LLMAPIUserError):
|
|
extract_completion_text(FakeCompletion(content))
|
|
|
|
def test_format_llm_error_redacts_api_key(self):
|
|
error = RuntimeError("bad key sk-secret-value failed")
|
|
|
|
formatted = format_llm_error(error, api_key="sk-secret-value")
|
|
|
|
self.assertTrue(formatted.startswith("LLM API Error: "))
|
|
self.assertNotIn("sk-secret-value", formatted)
|
|
self.assertIn("[redacted]", formatted)
|
|
|
|
def test_format_llm_error_scrubs_data_urls_and_truncates_long_messages(self):
|
|
data_url = "data:image/jpeg;base64," + ("a" * 500)
|
|
error = RuntimeError("bad request " + data_url + " " + ("b" * 500))
|
|
|
|
formatted = format_llm_error(error)
|
|
|
|
self.assertIn("[image data redacted]", formatted)
|
|
self.assertNotIn("aaaaa", formatted)
|
|
self.assertLessEqual(len(formatted), 260)
|
|
|
|
def test_create_openai_client_reports_missing_dependency(self):
|
|
real_import = builtins.__import__
|
|
|
|
def import_without_openai(name, *args, **kwargs):
|
|
if name == "openai":
|
|
raise ImportError("No module named openai")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
with mock.patch("builtins.__import__", side_effect=import_without_openai):
|
|
with self.assertRaises(LLMAPIUserError) as context:
|
|
create_openai_client("sk-test", "https://example.test/v1")
|
|
|
|
self.assertIn("requirements.txt", str(context.exception))
|
|
|
|
def test_create_openai_client_uses_lazy_openai_import(self):
|
|
calls = []
|
|
|
|
class FakeOpenAI:
|
|
def __init__(self, **kwargs):
|
|
calls.append(kwargs)
|
|
|
|
fake_openai_module = types.SimpleNamespace(OpenAI=FakeOpenAI)
|
|
real_import = builtins.__import__
|
|
|
|
def import_with_fake_openai(name, *args, **kwargs):
|
|
if name == "openai":
|
|
return fake_openai_module
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
with mock.patch("builtins.__import__", side_effect=import_with_fake_openai):
|
|
client = create_openai_client("sk-test", "https://example.test/v1")
|
|
|
|
self.assertIsInstance(client, FakeOpenAI)
|
|
self.assertEqual(
|
|
calls,
|
|
[
|
|
{
|
|
"api_key": "sk-test",
|
|
"base_url": "https://example.test/v1",
|
|
"timeout": DEFAULT_REQUEST_TIMEOUT_SECONDS,
|
|
"max_retries": DEFAULT_MAX_RETRIES,
|
|
}
|
|
],
|
|
)
|
|
|
|
def test_create_openai_client_accepts_explicit_timeout_policy(self):
|
|
calls = []
|
|
|
|
class FakeOpenAI:
|
|
def __init__(self, **kwargs):
|
|
calls.append(kwargs)
|
|
|
|
fake_openai_module = types.SimpleNamespace(OpenAI=FakeOpenAI)
|
|
real_import = builtins.__import__
|
|
|
|
def import_with_fake_openai(name, *args, **kwargs):
|
|
if name == "openai":
|
|
return fake_openai_module
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
with mock.patch("builtins.__import__", side_effect=import_with_fake_openai):
|
|
create_openai_client(
|
|
"sk-test",
|
|
"https://example.test/v1",
|
|
timeout=30.0,
|
|
max_retries=0,
|
|
)
|
|
|
|
self.assertEqual(calls[0]["timeout"], 30.0)
|
|
self.assertEqual(calls[0]["max_retries"], 0)
|
|
|
|
def test_is_openai_error_detects_openai_errors_when_available(self):
|
|
class FakeOpenAIError(Exception):
|
|
pass
|
|
|
|
fake_openai_module = types.SimpleNamespace(OpenAIError=FakeOpenAIError)
|
|
real_import = builtins.__import__
|
|
|
|
def import_with_fake_openai(name, *args, **kwargs):
|
|
if name == "openai":
|
|
return fake_openai_module
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
with mock.patch("builtins.__import__", side_effect=import_with_fake_openai):
|
|
self.assertTrue(is_openai_error(FakeOpenAIError("boom")))
|
|
self.assertFalse(is_openai_error(RuntimeError("boom")))
|
|
|
|
def test_is_openai_error_returns_false_when_openai_is_missing(self):
|
|
real_import = builtins.__import__
|
|
|
|
def import_without_openai(name, *args, **kwargs):
|
|
if name == "openai":
|
|
raise ImportError("No module named openai")
|
|
return real_import(name, *args, **kwargs)
|
|
|
|
with mock.patch("builtins.__import__", side_effect=import_without_openai):
|
|
self.assertFalse(is_openai_error(RuntimeError("boom")))
|
|
|
|
|
|
class ImageBatchEncodingTests(unittest.TestCase):
|
|
def test_image_encoder_has_no_legacy_size_defaults(self):
|
|
signature = inspect.signature(image_encoding.encode_image_to_safe_jpeg_b64)
|
|
|
|
self.assertEqual(signature.parameters["max_bytes"].default, inspect._empty)
|
|
self.assertEqual(signature.parameters["max_side"].default, inspect._empty)
|
|
self.assertFalse(hasattr(image_encoding, "DEFAULT_MAX_IMAGE_BYTES"))
|
|
self.assertFalse(hasattr(image_encoding, "DEFAULT_MAX_IMAGE_SIDE"))
|
|
|
|
def test_encodes_single_image_to_jpeg_data_url(self):
|
|
image = Image.new("RGB", (8, 8), (20, 40, 60))
|
|
urls = encode_comfy_image_batch_to_data_urls(
|
|
[image],
|
|
max_images=4,
|
|
max_bytes=8_000,
|
|
max_side=64,
|
|
)
|
|
|
|
self.assertEqual(len(urls), 1)
|
|
self.assertTrue(urls[0].startswith("data:image/jpeg;base64,"))
|
|
raw = base64.b64decode(urls[0].split(",", 1)[1])
|
|
self.assertLessEqual(len(raw), 8_000)
|
|
|
|
def test_truncates_image_batch_to_max_images(self):
|
|
images = [Image.new("RGB", (4, 4), (idx, idx, idx)) for idx in range(6)]
|
|
|
|
urls = encode_comfy_image_batch_to_data_urls(
|
|
images,
|
|
max_images=4,
|
|
max_bytes=8_000,
|
|
max_side=64,
|
|
)
|
|
|
|
self.assertEqual(len(urls), 4)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|