import importlib.util import base64 import io import json import sys from pathlib import Path import pytest import torch from PIL import Image def _load_module(): module_name = "easy_media_llm_api_for_tests" sys.modules.pop(module_name, None) module_path = Path(__file__).parents[1] / "utils" / "llm_api.py" spec = importlib.util.spec_from_file_location(module_name, module_path) module = importlib.util.module_from_spec(spec) sys.modules[module_name] = module assert spec.loader is not None spec.loader.exec_module(module) return module class _Response: def __init__(self, payload): self.payload = payload def __enter__(self): return self def __exit__(self, *args): return False def read(self): return json.dumps(self.payload).encode("utf-8") def _png_data_uri(width=256, height=256): buffer = io.BytesIO() Image.new("RGB", (width, height)).save(buffer, format="PNG") return "data:image/png;base64," + base64.b64encode(buffer.getvalue()).decode() def test_minimax_length_to_seconds_rounds_to_nearest_and_clamps(): module = _load_module() assert module.minimax_length_to_seconds(1) == 4 assert module.minimax_length_to_seconds(49) == 4 assert module.minimax_length_to_seconds(50) == 4 assert module.minimax_length_to_seconds(124) == 5 assert module.minimax_length_to_seconds(10_000) == 15 def test_prompt_enhancer_model_options_match_node_contract(): module = _load_module() assert module.PROMPT_ENHANCER_MODELS == [ "h3-context-ir (海螺官方)", "doubao-seed-2-0-pro-260215 (火山引擎)", "glm-5v-turbo (智谱)", "bytedance/doubao-seed-2.0-pro (RunningHub)", "glm-5v-turbo (RunningHub)", "llama.cpp (本地)", ] def test_third_party_video_url_capabilities_match_provider_documentation(): module = _load_module() assert module.prompt_enhancer_supports_video_url(module.VOLCENGINE_MODEL) is True assert module.prompt_enhancer_supports_video_url(module.ZHIPU_MODEL) is True assert module.prompt_enhancer_supports_video_url(module.RUNNINGHUB_DOUBAO_MODEL) assert module.prompt_enhancer_supports_video_url(module.RUNNINGHUB_GLM_MODEL) def test_llamacpp_max_tokens_widget_limits(): module = _load_module() assert module.PROMPT_ENHANCER_MAX_TOKENS[module.LLAMACPP_MODEL] == (512, 768) @pytest.mark.parametrize( ("model_name", "config_key"), [ ("MINIMAX_MODEL", "MINIMAX_API_KEY"), ("VOLCENGINE_MODEL", "VOLCENGINE_API_KEY"), ("ZHIPU_MODEL", "BIGMODEL_API_KEY"), ("RUNNINGHUB_DOUBAO_MODEL", "RUNNINGHUB_API_KEY"), ("RUNNINGHUB_GLM_MODEL", "RUNNINGHUB_API_KEY"), ], ) def test_client_reads_empty_api_key_from_vendor_config( tmp_path, monkeypatch, model_name, config_key, ): module = _load_module() config_path = tmp_path / "config.yaml" config_path.write_text(f"{config_key}: config-secret\n", encoding="utf-8") for environment_name in ( "MINIMAX_API_KEY", "VOLCENGINE_API_KEY", "BIGMODEL_API_KEY", "RUNNINGHUB_API_KEY", "ARK_API_KEY", "ZHIPU_API_KEY", ): monkeypatch.delenv(environment_name, raising=False) client = module.PromptEnhancerClient( getattr(module, model_name), "", config_path=config_path, ) assert client.api_key == "config-secret" def test_explicit_api_key_takes_priority_over_config(tmp_path): module = _load_module() config_path = tmp_path / "config.yaml" config_path.write_text("MINIMAX_API_KEY: config-secret\n", encoding="utf-8") client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "widget-secret", config_path=config_path, ) assert client.api_key == "widget-secret" def test_minimax_upload_configuration_is_fixed(tmp_path): module = _load_module() config_path = tmp_path / "config.yaml" config_path.write_text( "\n".join( [ "MINIMAX_UPLOAD_ENDPOINT: https://upload.example.test/files", "MINIMAX_UPLOAD_PURPOSE: should-not-be-used", ] ), encoding="utf-8", ) client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", config_path=config_path, ) assert client.upload_endpoint == "https://api.minimaxi.com/v1/files/upload" assert client.upload_purpose == "video_generation_input" def test_missing_api_key_points_to_vendor_config_key(tmp_path, monkeypatch): module = _load_module() config_path = tmp_path / "config.yaml" config_path.write_text("VOLCENGINE_API_KEY:\n", encoding="utf-8") for environment_name in ("VOLCENGINE_API_KEY", "ARK_API_KEY"): monkeypatch.delenv(environment_name, raising=False) with pytest.raises( ValueError, match=( "Enter apikey or configure VOLCENGINE_API_KEY " "in config.yaml" ), ): module.PromptEnhancerClient( module.VOLCENGINE_MODEL, "", config_path=config_path, ) def test_image_tensor_data_uris_flattens_batches(): module = _load_module() images = torch.zeros(2, 4, 5, 3) encoded = module.image_tensor_data_uris([images, torch.ones(4, 5, 3)]) assert len(encoded) == 3 assert all(value.startswith("data:image/png;base64,") for value in encoded) def test_image_tensor_data_uris_limits_third_party_images_to_two_megapixels(): module = _load_module() encoded = module.image_tensor_data_uris( torch.zeros(1, 1200, 2000, 3), max_pixels=2_000_000, ) image_bytes = base64.b64decode(encoded[0].split(",", 1)[1]) with Image.open(io.BytesIO(image_bytes)) as image: assert image.width * image.height <= 2_000_000 def test_video_frame_data_uris_samples_at_most_24_frames(): module = _load_module() class Video: def get_components(self): return type("Components", (), {"images": torch.zeros(30, 8, 8, 3)})() encoded = module.video_frame_data_uris([Video()]) assert len(encoded) == 24 assert all(value.startswith("data:image/png;base64,") for value in encoded) def test_zhipu_preserves_public_video_url_without_sampling_frames(): module = _load_module() class Video: def get_stream_source(self): return "https://example.com/reference.mp4" def get_components(self): raise AssertionError("public video URL must not be sampled") assert module.prompt_enhancer_video_inputs(module.ZHIPU_MODEL, [Video()]) == [ "https://example.com/reference.mp4" ] def test_zhipu_samples_local_video_when_only_public_urls_are_supported(): module = _load_module() class Video: def get_stream_source(self): return "/tmp/reference.mp4" def get_components(self): return type("Components", (), {"images": torch.zeros(30, 8, 8, 3)})() encoded = module.prompt_enhancer_video_inputs(module.ZHIPU_MODEL, [Video()]) assert len(encoded) == 24 assert all(value.startswith("data:image/png;base64,") for value in encoded) def test_volcengine_uses_native_data_uri_for_local_video(): module = _load_module() class Video: def get_stream_source(self): return b"video-bytes" def get_container_format(self): return "mp4" def get_components(self): raise AssertionError("native video data URI must not be sampled") encoded = module.prompt_enhancer_video_inputs(module.VOLCENGINE_MODEL, [Video()]) assert len(encoded) == 1 assert encoded[0].startswith("data:video/mp4;base64,") def test_runninghub_uses_native_video_data_uri_with_official_limits(monkeypatch): module = _load_module() calls = [] class Video: def get_stream_source(self): return b"video-bytes" monkeypatch.setattr( module, "video_data_uris", lambda values, **kwargs: calls.append(kwargs) or ["data:video/mp4;base64,AAAA"], ) encoded = module.prompt_enhancer_video_inputs( module.RUNNINGHUB_DOUBAO_MODEL, [Video()], ) assert encoded == ["data:video/mp4;base64,AAAA"] assert calls == [{"max_bytes": 10 * 1024 * 1024, "max_duration": 15}] def test_video_data_uri_serializes_active_trim_instead_of_original_stream(): module = _load_module() class Video: def get_stream_source(self): return b"original-video" def get_active_trim_window(self): return 2.0, 5.0 def save_to(self, output_path): Path(output_path).write_bytes(b"trimmed-video") encoded = module.video_data_uris([Video()]) assert base64.b64decode(encoded[0].split(",", 1)[1]) == b"trimmed-video" def test_video_data_uri_enforces_duration_even_when_file_is_under_size_limit(monkeypatch): module = _load_module() calls = [] class Video: def get_stream_source(self): return b"small-long-video" monkeypatch.setattr(module, "_probe_video_duration", lambda _source, _data: 20.0) monkeypatch.setattr( module, "_limit_video_for_data_uri", lambda source, data, **kwargs: calls.append((source, data, kwargs)) or b"limited", ) encoded = module.video_data_uris( [Video()], max_bytes=10 * 1024 * 1024, max_duration=15, ) assert calls == [( b"small-long-video", b"small-long-video", {"max_bytes": 10 * 1024 * 1024, "max_duration": 15}, )] assert base64.b64decode(encoded[0].split(",", 1)[1]) == b"limited" @pytest.mark.parametrize( ("task_type", "expected_mode", "image_count", "video_count", "audio_count"), [ ("t2v", "t2va", 0, 0, 0), ("i2v", "i2va", 2, 0, 0), ("l2v", "i2va", 1, 0, 0), ("v2v", "r2va", 2, 1, 1), ("r2v", "r2va", 2, 1, 1), ("vi2v", "r2va", 2, 1, 1), ("rv2v", "r2va", 2, 1, 1), ], ) def test_h3_task_type_selects_request_mode_and_media( task_type, expected_mode, image_count, video_count, audio_count, ): module = _load_module() mode, images, videos, audios = module.PromptEnhancerClient._h3_media_for_task_type( task_type, ["image-1", "image-2"], ["video-1"], ["audio-1"], ) assert mode == expected_mode assert len(images) == image_count assert len(videos) == video_count assert len(audios) == audio_count def test_h3_i2va_roles_follow_task_type_and_image_count(): module = _load_module() i2v_content = module.PromptEnhancerClient._h3_content( "Prompt", "i2v", ["first", "last"], [], [] ) l2v_content = module.PromptEnhancerClient._h3_content( "Prompt", "l2v", ["last"], [], [] ) assert [item.get("role") for item in i2v_content[1:]] == [ "first_frame", "last_frame", ] assert l2v_content[1]["role"] == "last_frame" def test_h3_r2va_assigns_roles_by_media_type(): module = _load_module() content = module.PromptEnhancerClient._h3_content( "Prompt", "rv2v", ["image"], ["video"], ["audio"], ) assert [item.get("role") for item in content[1:]] == [ "reference_image", "reference_video", "reference_audio", ] def test_h3_rejects_image_dimensions_before_upload(): module = _load_module() client = module.PromptEnhancerClient(module.MINIMAX_MODEL, "secret") with pytest.raises(ValueError, match="between 256 and 5760"): client._upload_h3_media(_png_data_uri(128, 256), "image", 1, None) def test_minimax_async_request_returns_task_id_and_maps_ratio(): module = _load_module() requests = [] log_calls = [] def opener(request, timeout): requests.append((request, timeout)) if request.full_url.endswith("/v1/files/upload"): return _Response({"file": {"file_id": 101}}) return _Response({"task_id": "task-123"}) client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=opener, ) result = client.enhance( system_prompt="System instructions that H3 must ignore. " * 300, user_prompt="A dancer turns.", task_type="r2v", duration=2, ratio="9:19", seed=7, image_urls=[_png_data_uri()], return_async=True, file_count=2, request_logger=lambda node_name, message=None: log_calls.append( (node_name, message) ), ) assert result.prompt == "" assert result.task_id == "task-123" assert result.file_ids == "101" upload_request, upload_timeout = requests[0] assert upload_timeout == 300.0 assert upload_request.full_url == "https://api.minimaxi.com/v1/files/upload" assert b"video_generation_input" in upload_request.data request, timeout = requests[1] assert timeout == 300.0 assert request.full_url == "https://api.minimaxi.com/v2/h3_context_ir" payload = json.loads(request.data) assert payload["duration"] == 4 assert payload["ratio"] == "9:16" assert payload["content"][0] == { "type": "text", "text": "A dancer turns.", } assert payload["content"][1]["role"] == "reference_image" assert payload["content"][1]["image_url"]["url"] == "mm_file://101" assert request.headers["Authorization"] == "Bearer secret" messages = [message for _node_name, message in log_calls] assert messages[0].startswith("H3 uploading image #1") assert messages[1] == "H3 upload image #1 succeeded: file_id=101" assert "images=1, videos=0, audios=0, files=2" in messages[2] assert messages[3] == "H3 create request succeeded: task_id=task-123" def test_minimax_rejects_only_user_prompt_over_7000_characters(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=lambda request, timeout: requests.append(request), ) with pytest.raises(ValueError, match="cannot exceed 7000 characters"): client.enhance( system_prompt="Ignored system instructions.", user_prompt="x" * 7001, task_type="t2v", duration=5, ratio="16:9", seed=0, ) assert requests == [] def test_minimax_returns_uploaded_file_ids_in_media_order(): module = _load_module() upload_ids = iter([101, 202, 303]) def opener(request, timeout): if request.full_url.endswith("/v1/files/upload"): return _Response({"file": {"file_id": next(upload_ids)}}) return _Response({"task_id": "task-files"}) client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=opener, ) result = client.enhance( system_prompt="", user_prompt="Use all references.", task_type="r2v", duration=5, ratio="adaptive", seed=0, image_urls=[_png_data_uri()], video_urls=["data:video/mp4;base64,AAAA"], audio_urls=["data:audio/wav;base64,AAAA"], return_async=True, ) assert result.file_ids == "101,202,303" def test_minimax_rejects_empty_user_prompt_before_upload_or_create_request(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=lambda request, timeout: requests.append(request), ) with pytest.raises(ValueError, match="user_prompt cannot be empty"): client.enhance( system_prompt="Ignored system instructions.", user_prompt=" ", task_type="r2v", duration=5, ratio="adaptive", seed=0, image_urls=[_png_data_uri()], ) assert requests == [] def test_minimax_sync_request_polls_every_five_seconds(): module = _load_module() payloads = iter( [ {"file": {"file_id": 404}}, {"task_id": "task-456"}, {"task": {"status": "queued"}}, { "task": { "status": "succeeded", "content": {"prompt": "```text\nEnhanced.\n```"}, } }, ] ) sleeps = [] statuses = [] log_messages = [] client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=lambda request, timeout: _Response(next(payloads)), sleeper=sleeps.append, ) result = client.enhance( system_prompt="", user_prompt="A city at night.", task_type="i2v", duration=5, ratio="adaptive", seed=0, image_urls=[_png_data_uri()], poll_interval=5.0, poll_callback=statuses.append, request_logger=lambda _name, message=None: log_messages.append(message), ) assert result.prompt == "Enhanced." assert result.task_id == "task-456" assert result.file_ids == "404" assert sleeps == [5.0, 5.0] assert statuses == ["queued", "succeeded"] assert log_messages[-4:] == [ "H3 create request succeeded: task_id=task-456", "H3 polling result: task_id=task-456, status=queued", "H3 polling result: task_id=task-456, status=succeeded", "H3 task succeeded: task_id=task-456", ] def test_minimax_sync_polling_defaults_to_ten_minutes(): module = _load_module() assert module.PromptEnhancerClient.enhance.__kwdefaults__["poll_timeout"] == 600.0 def test_minimax_sync_polling_stops_at_overall_timeout(): module = _load_module() payloads = iter( [ {"task_id": "task-timeout"}, {"task": {"status": "queued"}}, ] ) now = [0.0] def sleep(seconds): now[0] += seconds client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=lambda request, timeout: _Response(next(payloads)), sleeper=sleep, clock=lambda: now[0], ) with pytest.raises( module.PromptEnhancerApiError, match="timed out after 6s: task-timeout", ): client.enhance( system_prompt="", user_prompt="Prompt", task_type="t2v", duration=5, ratio="16:9", seed=0, poll_interval=5, poll_timeout=6, ) def test_openai_compatible_request_extracts_prompt_without_task_id(): module = _load_module() requests = [] def opener(request, timeout): requests.append(request) return _Response({ "choices": [{"message": {"content": " ```text\nBetter prompt.\n``` "}}] }) client = module.PromptEnhancerClient( module.ZHIPU_MODEL, "secret", opener=opener, ) result = client.enhance( system_prompt="You enhance prompts.", user_prompt="A dog runs.", task_type="i2v", duration=6, ratio="16:9", seed=42, image_urls=["data:image/png;base64,AAAA"], return_async=True, ) assert result.prompt == "Better prompt." assert result.task_id == "" payload = json.loads(requests[0].data) assert payload["seed"] == 42 assert payload["messages"][0] == { "role": "system", "content": "You enhance prompts.", } assert payload["messages"][1]["content"][1]["type"] == "image_url" def test_openai_compatible_request_maps_comfyui_seed_to_int32(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.ZHIPU_MODEL, "secret", opener=lambda request, timeout: ( requests.append(request) or _Response({"choices": [{"message": {"content": "Enhanced"}}]}) ), ) client.enhance( system_prompt="", user_prompt="Prompt", task_type="t2v", duration=5, ratio="16:9", seed=715809353714956, ) payload = json.loads(requests[0].data) assert payload["seed"] == 1514229004 assert 0 <= payload["seed"] <= 0x7FFFFFFF @pytest.mark.parametrize( ("prompt", "expected"), [ ("```text\nA cinematic prompt.\n```", "A cinematic prompt."), (" ```text\r\nA cinematic prompt.\r\n``` ", "A cinematic prompt."), ("```json\n{\"prompt\": \"keep\"}\n```", "```json\n{\"prompt\": \"keep\"}\n```"), ("```text\nMissing closing fence", "```text\nMissing closing fence"), ("Plain prompt", "Plain prompt"), ], ) def test_strip_text_code_fence_only_removes_complete_text_fence(prompt, expected): module = _load_module() assert module.strip_text_code_fence(prompt) == expected def test_minimax_failed_task_is_logged(): module = _load_module() payloads = iter( [ {"task_id": "task-failed"}, {"task": {"status": "failed", "error": {"message": "bad media"}}}, ] ) log_messages = [] client = module.PromptEnhancerClient( module.MINIMAX_MODEL, "secret", opener=lambda request, timeout: _Response(next(payloads)), sleeper=lambda _seconds: None, ) with pytest.raises(module.PromptEnhancerApiError, match="bad media"): client.enhance( system_prompt="", user_prompt="Prompt", task_type="t2v", duration=5, ratio="16:9", seed=0, request_logger=lambda _name, message=None: log_messages.append(message), ) assert log_messages[-2:] == [ "H3 polling result: task_id=task-failed, status=failed", "H3 task failed: task_id=task-failed, error=bad media", ] def test_task_type_controls_openai_compatible_media_parts(): module = _load_module() requests = [] def opener(request, timeout): requests.append(request) return _Response({"choices": [{"message": {"content": "Enhanced"}}]}) client = module.PromptEnhancerClient(module.ZHIPU_MODEL, "secret", opener=opener) client.enhance( system_prompt="", user_prompt="Text only", task_type="t2v", duration=5, ratio="16:9", seed=0, image_urls=["data:image/png;base64,AAAA"], video_urls=["data:video/mp4;base64,AAAA"], audio_urls=["data:audio/wav;base64,AAAA"], ) payload = json.loads(requests[0].data) content = payload["messages"][0]["content"] assert [part["type"] for part in content] == ["text"] assert payload["max_tokens"] == 65536 def test_third_party_max_tokens_is_capped_by_model_limit(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.VOLCENGINE_MODEL, "secret", opener=lambda request, timeout: ( requests.append(request) or _Response({"choices": [{"message": {"content": "Enhanced"}}]}) ), ) client.enhance( system_prompt="", user_prompt="Enhance this", task_type="t2v", duration=5, ratio="adaptive", seed=0, max_tokens=200000, ) assert json.loads(requests[0].data)["max_tokens"] == 131072 def test_native_video_provider_uses_public_video_url_and_omits_audio(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.ZHIPU_MODEL, "secret", opener=lambda request, timeout: ( requests.append(request) or _Response({"choices": [{"message": {"content": "Enhanced"}}]}) ), ) client.enhance( system_prompt="", user_prompt="Edit the video", task_type="v2v", duration=5, ratio="16:9", seed=0, video_urls=["https://example.com/reference.mp4"], audio_urls=["data:audio/wav;base64,AAAA"], ) payload = json.loads(requests[0].data) content = payload["messages"][0]["content"] assert [part["type"] for part in content] == ["text", "video_url"] def test_sampled_video_frames_remain_image_parts_for_native_provider(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.ZHIPU_MODEL, "secret", opener=lambda request, timeout: ( requests.append(request) or _Response({"choices": [{"message": {"content": "Enhanced"}}]}) ), ) client.enhance( system_prompt="", user_prompt="Edit the video", task_type="v2v", duration=5, ratio="16:9", seed=0, video_urls=["data:image/png;base64,AAAA"], ) payload = json.loads(requests[0].data) content = payload["messages"][0]["content"] assert [part["type"] for part in content] == ["text", "image_url"] def test_runninghub_uses_video_url_and_omits_unsupported_seed(): module = _load_module() requests = [] client = module.PromptEnhancerClient( module.RUNNINGHUB_GLM_MODEL, "secret", opener=lambda request, timeout: ( requests.append(request) or _Response({"choices": [{"message": {"content": "Enhanced"}}]}) ), ) client.enhance( system_prompt="", user_prompt="Edit the video", task_type="v2v", duration=5, ratio="adaptive", seed=123, video_urls=["data:video/mp4;base64,AAAA"], ) payload = json.loads(requests[0].data) assert [part["type"] for part in payload["messages"][0]["content"]] == [ "text", "video_url", ] assert "seed" not in payload def test_http_error_is_normalized_with_provider_message(): module = _load_module() def opener(request, timeout): raise module.urllib.error.HTTPError( request.full_url, 429, "Too Many Requests", {}, io.BytesIO(b'{"error":{"message":"rate limited"}}'), ) client = module.PromptEnhancerClient( module.ZHIPU_MODEL, "secret", opener=opener, ) with pytest.raises(module.PromptEnhancerApiError, match="HTTP 429: rate limited"): client.enhance( system_prompt="", user_prompt="Prompt", task_type="t2v", duration=5, ratio="16:9", seed=0, )