feat: add ai image service
This commit is contained in:
@@ -0,0 +1,383 @@
|
||||
import base64
|
||||
import binascii
|
||||
import logging
|
||||
import mimetypes
|
||||
from dataclasses import dataclass, field
|
||||
from pathlib import Path
|
||||
from typing import Any, Iterable, Optional, Tuple
|
||||
from urllib.parse import urljoin, urlparse
|
||||
|
||||
import requests
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
API_AUTO = "auto"
|
||||
API_CHAT = "chat"
|
||||
API_GEMINI = "gemini"
|
||||
API_IMAGES = "images"
|
||||
API_IMAGES_EDITS = "images_edits"
|
||||
|
||||
SUPPORTED_API_TYPES = {
|
||||
API_AUTO,
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES,
|
||||
API_IMAGES_EDITS,
|
||||
}
|
||||
|
||||
_BASE64_KEYS = {"image_base64", "base64", "b64_json", "data"}
|
||||
_IMAGE_EXTENSIONS = {".png", ".jpg", ".jpeg", ".webp", ".gif"}
|
||||
|
||||
|
||||
class AiImageServiceError(RuntimeError):
|
||||
"""Base error for AI image generation service failures."""
|
||||
|
||||
|
||||
class AiConfigError(ValueError):
|
||||
"""Raised when model configuration is missing or invalid."""
|
||||
|
||||
|
||||
@dataclass
|
||||
class AiModelConfig:
|
||||
"""AI image model configuration loaded from ai_models.json."""
|
||||
|
||||
url: str
|
||||
model: str
|
||||
api_key: str
|
||||
api_type: str = API_AUTO
|
||||
timeout_seconds: int = 240
|
||||
connect_timeout_seconds: int = 30
|
||||
extra_body: Any = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data):
|
||||
if not isinstance(data, dict):
|
||||
raise AiConfigError("AI 模型配置必须是对象")
|
||||
return cls(
|
||||
url=str(data.get("url", "") or ""),
|
||||
model=str(data.get("model", "") or ""),
|
||||
api_key=str(data.get("api_key", "") or ""),
|
||||
api_type=str(data.get("api_type", API_AUTO) or API_AUTO),
|
||||
timeout_seconds=_to_int(data.get("timeout_seconds", 240), 240),
|
||||
connect_timeout_seconds=_to_int(data.get("connect_timeout_seconds", 30), 30),
|
||||
extra_body=data.get("extra_body") or {},
|
||||
)
|
||||
|
||||
|
||||
def api_config_errors(config):
|
||||
"""Return human-readable configuration errors without raising."""
|
||||
try:
|
||||
cfg = _coerce_config(config)
|
||||
except AiConfigError as exc:
|
||||
return [str(exc)]
|
||||
errors = []
|
||||
|
||||
if not cfg.url.strip():
|
||||
errors.append("缺少 url")
|
||||
else:
|
||||
scheme = urlparse(cfg.url).scheme.lower()
|
||||
if scheme not in ("http", "https"):
|
||||
errors.append("url 必须是 http(s) 地址")
|
||||
if not cfg.model.strip():
|
||||
errors.append("缺少 model")
|
||||
if not cfg.api_key.strip():
|
||||
errors.append("缺少 api_key")
|
||||
if cfg.api_type not in SUPPORTED_API_TYPES:
|
||||
errors.append("api_type 不支持: {}".format(cfg.api_type))
|
||||
if cfg.timeout_seconds <= 0:
|
||||
errors.append("timeout_seconds 必须大于 0")
|
||||
if cfg.connect_timeout_seconds <= 0:
|
||||
errors.append("connect_timeout_seconds 必须大于 0")
|
||||
if not isinstance(cfg.extra_body, dict):
|
||||
errors.append("extra_body 必须是对象")
|
||||
|
||||
return errors
|
||||
|
||||
|
||||
def validate_api_config(config):
|
||||
errors = api_config_errors(config)
|
||||
if errors:
|
||||
raise AiConfigError("; ".join(errors))
|
||||
|
||||
|
||||
def detect_api_type(url, api_type=API_AUTO):
|
||||
"""Resolve auto/chat/gemini/images/images_edits for a model endpoint."""
|
||||
if api_type and api_type != API_AUTO:
|
||||
if api_type not in SUPPORTED_API_TYPES:
|
||||
raise AiConfigError("api_type 不支持: {}".format(api_type))
|
||||
return api_type
|
||||
|
||||
lower_url = str(url).lower()
|
||||
if "generatecontent" in lower_url or "gemini" in lower_url:
|
||||
return API_GEMINI
|
||||
if "/images/edits" in lower_url:
|
||||
return API_IMAGES_EDITS
|
||||
if "/images" in lower_url:
|
||||
return API_IMAGES
|
||||
return API_CHAT
|
||||
|
||||
|
||||
def normalize_api_url(url, api_type):
|
||||
"""Return an endpoint URL for the selected API type.
|
||||
|
||||
Full endpoint URLs are kept. Base URLs are completed using OpenAI-compatible
|
||||
defaults or Gemini generateContent where appropriate.
|
||||
"""
|
||||
raw = str(url).strip()
|
||||
if not raw:
|
||||
return raw
|
||||
|
||||
endpoint = {
|
||||
API_CHAT: "chat/completions",
|
||||
API_IMAGES: "images/generations",
|
||||
API_IMAGES_EDITS: "images/edits",
|
||||
}.get(api_type)
|
||||
|
||||
lower_path = urlparse(raw).path.lower().rstrip("/")
|
||||
if api_type == API_GEMINI:
|
||||
if "generatecontent" in lower_path:
|
||||
return raw
|
||||
return _join_url(raw, "v1beta/models/{model}:generateContent")
|
||||
|
||||
if endpoint is None:
|
||||
return raw
|
||||
if lower_path.endswith("/" + endpoint):
|
||||
return raw
|
||||
if lower_path.endswith("/v1"):
|
||||
return _join_url(raw, endpoint)
|
||||
return _join_url(raw, "v1/" + endpoint)
|
||||
|
||||
|
||||
def image_to_data_url(image_path):
|
||||
"""Read a local image file and return data:<mime>;base64,..."""
|
||||
path = Path(image_path)
|
||||
mime_type = mimetypes.guess_type(str(path))[0] or "application/octet-stream"
|
||||
with open(str(path), "rb") as f:
|
||||
encoded = base64.b64encode(f.read()).decode("ascii")
|
||||
return "data:{};base64,{}".format(mime_type, encoded)
|
||||
|
||||
|
||||
def decode_image_data_url(data_url):
|
||||
marker = ";base64,"
|
||||
if marker not in data_url:
|
||||
raise AiImageServiceError("Unsupported data URL image format")
|
||||
return base64.b64decode(data_url.split(marker, 1)[1])
|
||||
|
||||
|
||||
def build_payload(config, prompt, image_data_url, resolution="1K", aspect_ratio="1:1"):
|
||||
"""Build JSON body for chat/gemini/images request types."""
|
||||
cfg = _coerce_config(config)
|
||||
api_type = detect_api_type(cfg.url, cfg.api_type)
|
||||
|
||||
if api_type == API_CHAT:
|
||||
payload = {
|
||||
"model": cfg.model,
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": prompt},
|
||||
{"type": "image_url", "image_url": {"url": image_data_url}},
|
||||
],
|
||||
}
|
||||
],
|
||||
"stream": False,
|
||||
}
|
||||
elif api_type == API_GEMINI:
|
||||
mime_type, data = _split_data_url(image_data_url)
|
||||
payload = {
|
||||
"contents": [
|
||||
{
|
||||
"parts": [
|
||||
{"text": prompt},
|
||||
{"inlineData": {"mimeType": mime_type, "data": data}},
|
||||
]
|
||||
}
|
||||
],
|
||||
"generationConfig": {"responseModalities": ["TEXT", "IMAGE"]},
|
||||
}
|
||||
elif api_type == API_IMAGES:
|
||||
payload = {
|
||||
"model": cfg.model,
|
||||
"prompt": prompt,
|
||||
"image_urls": [image_data_url],
|
||||
"aspect_ratio": aspect_ratio,
|
||||
"resolution": resolution,
|
||||
"n": 1,
|
||||
}
|
||||
else:
|
||||
raise AiConfigError("build_payload 不支持 {}".format(api_type))
|
||||
|
||||
payload.update(cfg.extra_body)
|
||||
return payload
|
||||
|
||||
|
||||
def build_multipart_fields(config, prompt, image_path, resolution="1K"):
|
||||
"""Build data/files for images_edits requests.
|
||||
|
||||
Caller owns and must close the returned file object.
|
||||
"""
|
||||
cfg = _coerce_config(config)
|
||||
handle = open(str(image_path), "rb")
|
||||
mime_type = mimetypes.guess_type(str(image_path))[0] or "application/octet-stream"
|
||||
data = {
|
||||
"model": cfg.model,
|
||||
"prompt": prompt,
|
||||
"n": "1",
|
||||
"size": _resolution_to_size(resolution),
|
||||
}
|
||||
data.update(cfg.extra_body)
|
||||
files = {"image": (Path(image_path).name, handle, mime_type)}
|
||||
return data, files
|
||||
|
||||
|
||||
def extract_image_from_response(data, session=None, timeout=60):
|
||||
"""Return image bytes found in an AI JSON response, or None.
|
||||
|
||||
Searches data-url/base64 fields first, then image URLs.
|
||||
"""
|
||||
for key, value in _walk_json_items(data):
|
||||
if not isinstance(value, str):
|
||||
continue
|
||||
text = value.strip()
|
||||
if text.startswith("data:image/"):
|
||||
return decode_image_data_url(text)
|
||||
if key and key.lower() in _BASE64_KEYS and _looks_like_base64(text):
|
||||
try:
|
||||
return base64.b64decode(text)
|
||||
except (TypeError, ValueError, binascii.Error):
|
||||
pass
|
||||
|
||||
for _key, value in _walk_json_items(data):
|
||||
if isinstance(value, str) and _is_image_url(value):
|
||||
return download_image_url(value, session=session, timeout=timeout)
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def download_image_url(url, session=None, timeout=60):
|
||||
client = session or requests.Session()
|
||||
if hasattr(client, "trust_env"):
|
||||
client.trust_env = False
|
||||
response = client.get(url, timeout=timeout)
|
||||
response.raise_for_status()
|
||||
return response.content
|
||||
|
||||
|
||||
class ImageApiClient:
|
||||
"""Small HTTP client for AI image APIs.
|
||||
|
||||
GUI code should not use requests directly; it should call this service from
|
||||
worker threads and send progress back through Qt signals.
|
||||
"""
|
||||
|
||||
def __init__(self, config, session=None):
|
||||
self.config = _coerce_config(config)
|
||||
self.session = session or requests.Session()
|
||||
if hasattr(self.session, "trust_env"):
|
||||
self.session.trust_env = False
|
||||
|
||||
def generate(self, prompt, image_path, resolution="1K", aspect_ratio="1:1"):
|
||||
validate_api_config(self.config)
|
||||
api_type = detect_api_type(self.config.url, self.config.api_type)
|
||||
url = normalize_api_url(self.config.url, api_type)
|
||||
if api_type == API_GEMINI:
|
||||
url = url.replace("{model}", self.config.model)
|
||||
|
||||
headers = {"Authorization": "Bearer {}".format(self.config.api_key)}
|
||||
timeout = (self.config.connect_timeout_seconds, self.config.timeout_seconds)
|
||||
|
||||
if api_type == API_IMAGES_EDITS:
|
||||
data, files = build_multipart_fields(self.config, prompt, image_path, resolution)
|
||||
try:
|
||||
response = self.session.post(
|
||||
url,
|
||||
headers=headers,
|
||||
data=data,
|
||||
files=files,
|
||||
timeout=timeout,
|
||||
)
|
||||
finally:
|
||||
files["image"][1].close()
|
||||
else:
|
||||
data_url = image_to_data_url(image_path)
|
||||
payload = build_payload(self.config, prompt, data_url, resolution, aspect_ratio)
|
||||
headers["Content-Type"] = "application/json"
|
||||
response = self.session.post(url, headers=headers, json=payload, timeout=timeout)
|
||||
|
||||
response.raise_for_status()
|
||||
payload = response.json()
|
||||
image_bytes = extract_image_from_response(
|
||||
payload,
|
||||
session=self.session,
|
||||
timeout=self.config.timeout_seconds,
|
||||
)
|
||||
if not image_bytes:
|
||||
raise AiImageServiceError("AI 响应中未找到图片")
|
||||
return image_bytes
|
||||
|
||||
|
||||
def _coerce_config(config):
|
||||
if isinstance(config, AiModelConfig):
|
||||
return config
|
||||
return AiModelConfig.from_dict(config)
|
||||
|
||||
|
||||
def _to_int(value, default):
|
||||
try:
|
||||
return int(value or default)
|
||||
except (TypeError, ValueError):
|
||||
return 0
|
||||
|
||||
|
||||
def _join_url(base_url, suffix):
|
||||
base = str(base_url).rstrip("/") + "/"
|
||||
return urljoin(base, suffix)
|
||||
|
||||
|
||||
def _split_data_url(data_url):
|
||||
prefix, encoded = data_url.split(",", 1)
|
||||
mime_type = prefix[len("data:"):].split(";", 1)[0]
|
||||
return mime_type, encoded
|
||||
|
||||
|
||||
def _resolution_to_size(resolution):
|
||||
mapping = {
|
||||
"512": "512x512",
|
||||
"512px": "512x512",
|
||||
"1K": "1024x1024",
|
||||
"2K": "2048x2048",
|
||||
"4K": "4096x4096",
|
||||
}
|
||||
return mapping.get(str(resolution), str(resolution))
|
||||
|
||||
|
||||
def _walk_json_items(value, key=None):
|
||||
# type: (Any, Optional[str]) -> Iterable[Tuple[Optional[str], Any]]
|
||||
yield key, value
|
||||
if isinstance(value, dict):
|
||||
for child_key, child_value in value.items():
|
||||
for item in _walk_json_items(child_value, str(child_key)):
|
||||
yield item
|
||||
elif isinstance(value, list):
|
||||
for child_value in value:
|
||||
for item in _walk_json_items(child_value, key):
|
||||
yield item
|
||||
|
||||
|
||||
def _looks_like_base64(text):
|
||||
if len(text) < 8:
|
||||
return False
|
||||
try:
|
||||
base64.b64decode(text, validate=True)
|
||||
return True
|
||||
except (TypeError, ValueError, binascii.Error):
|
||||
return False
|
||||
|
||||
|
||||
def _is_image_url(text):
|
||||
parsed = urlparse(text.strip())
|
||||
if parsed.scheme.lower() not in ("http", "https"):
|
||||
return False
|
||||
suffix = Path(parsed.path).suffix.lower()
|
||||
return suffix in _IMAGE_EXTENSIONS
|
||||
@@ -1048,7 +1048,7 @@
|
||||
- [x] 前置依赖：`requirements.txt` 增加 Python 3.7 兼容的 `requests` / `urllib3` / `openpyxl` 锁定版本，并同步 `docs/03-technical-stack.md`
|
||||
- [x] `core/models.py` 新增 `OutfitTask` / `OutfitResult`（纯 dataclass，Python 3.7 兼容，不依赖 PySide6）
|
||||
- [x] `services/excel_service.py`：读行 → `List[OutfitTask]`、写回 D/E/F、占用检测、跳过「完成」/按设置重试「失败」/空字段安全跳过、每行即存
|
||||
- [ ] `services/ai_image_service.py`：移植旧项目 `ImageApiClient`（多模型、多请求格式 `chat/gemini/images/images_edits`、传图 data-url、递归取图、URL 归一化、字段校验）；注意 PEP585 类型注解改 Python 3.7 写法
|
||||
- [x] `services/ai_image_service.py`：移植旧项目 `ImageApiClient`（多模型、多请求格式 `chat/gemini/images/images_edits`、传图 data-url、递归取图、URL 归一化、字段校验）；注意 PEP585 类型注解改 Python 3.7 写法
|
||||
- [ ] `core/ai_outfit.py`：单行生成纯逻辑编排（提示词渲染 + 调用 + 保存 JPG + 产出 `OutfitResult`）
|
||||
- [ ] 单测：Excel 读写、提示词渲染、取图、命名去重（API 用 mock）；可在 Python 3.7 运行、不依赖 GUI
|
||||
|
||||
|
||||
@@ -0,0 +1,255 @@
|
||||
"""Tests for AI image service helpers and client request construction."""
|
||||
import base64
|
||||
import shutil
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
sys.path.insert(0, str(Path(__file__).parent.parent / "src"))
|
||||
|
||||
|
||||
class TestAiImageServiceHelpers(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.image_path = self.tmp / "sample.png"
|
||||
self.image_bytes = b"\x89PNG\r\n\x1a\nfake-png"
|
||||
self.image_path.write_bytes(self.image_bytes)
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(str(self.tmp), ignore_errors=True)
|
||||
|
||||
def test_api_config_errors_require_core_fields(self):
|
||||
from services.ai_image_service import api_config_errors
|
||||
|
||||
errors = api_config_errors({"url": "", "model": "", "api_key": ""})
|
||||
|
||||
self.assertIn("缺少 url", errors)
|
||||
self.assertIn("缺少 model", errors)
|
||||
self.assertIn("缺少 api_key", errors)
|
||||
|
||||
def test_api_config_errors_reject_non_object_config(self):
|
||||
from services.ai_image_service import api_config_errors
|
||||
|
||||
self.assertEqual(api_config_errors(["bad"]), ["AI 模型配置必须是对象"])
|
||||
|
||||
def test_api_config_errors_reject_invalid_extra_body_and_timeout(self):
|
||||
from services.ai_image_service import api_config_errors
|
||||
|
||||
errors = api_config_errors({
|
||||
"url": "https://api.example.test",
|
||||
"model": "m",
|
||||
"api_key": "k",
|
||||
"timeout_seconds": "bad",
|
||||
"extra_body": ["bad"],
|
||||
})
|
||||
|
||||
self.assertIn("timeout_seconds 必须大于 0", errors)
|
||||
self.assertIn("extra_body 必须是对象", errors)
|
||||
|
||||
def test_detect_api_type_from_url(self):
|
||||
from services.ai_image_service import (
|
||||
API_CHAT,
|
||||
API_GEMINI,
|
||||
API_IMAGES_EDITS,
|
||||
detect_api_type,
|
||||
)
|
||||
|
||||
self.assertEqual(detect_api_type("https://x/v1/chat/completions"), API_CHAT)
|
||||
self.assertEqual(detect_api_type("https://x/v1/images/edits"), API_IMAGES_EDITS)
|
||||
self.assertEqual(detect_api_type("https://x/v1beta/models/m:generateContent"), API_GEMINI)
|
||||
|
||||
def test_normalize_openai_urls(self):
|
||||
from services.ai_image_service import API_CHAT, API_IMAGES_EDITS, normalize_api_url
|
||||
|
||||
self.assertEqual(
|
||||
normalize_api_url("https://api.example.test", API_CHAT),
|
||||
"https://api.example.test/v1/chat/completions",
|
||||
)
|
||||
self.assertEqual(
|
||||
normalize_api_url("https://api.example.test/v1", API_IMAGES_EDITS),
|
||||
"https://api.example.test/v1/images/edits",
|
||||
)
|
||||
|
||||
def test_image_to_data_url(self):
|
||||
from services.ai_image_service import image_to_data_url
|
||||
|
||||
data_url = image_to_data_url(self.image_path)
|
||||
|
||||
self.assertTrue(data_url.startswith("data:image/png;base64,"))
|
||||
encoded = data_url.split(",", 1)[1]
|
||||
self.assertEqual(base64.b64decode(encoded), self.image_bytes)
|
||||
|
||||
def test_build_chat_payload(self):
|
||||
from services.ai_image_service import build_payload
|
||||
|
||||
payload = build_payload(
|
||||
{
|
||||
"url": "https://api.example.test/v1/chat/completions",
|
||||
"model": "m1",
|
||||
"api_key": "k",
|
||||
"api_type": "chat",
|
||||
"extra_body": {"temperature": 0},
|
||||
},
|
||||
"生成穿搭",
|
||||
"data:image/png;base64,abc",
|
||||
)
|
||||
|
||||
self.assertEqual(payload["model"], "m1")
|
||||
self.assertEqual(payload["messages"][0]["content"][0]["text"], "生成穿搭")
|
||||
self.assertEqual(
|
||||
payload["messages"][0]["content"][1]["image_url"]["url"],
|
||||
"data:image/png;base64,abc",
|
||||
)
|
||||
self.assertEqual(payload["temperature"], 0)
|
||||
|
||||
def test_build_gemini_payload(self):
|
||||
from services.ai_image_service import build_payload
|
||||
|
||||
payload = build_payload(
|
||||
{
|
||||
"url": "https://generativelanguage.googleapis.com/v1beta/models/m:generateContent",
|
||||
"model": "gemini",
|
||||
"api_key": "k",
|
||||
"api_type": "gemini",
|
||||
},
|
||||
"prompt",
|
||||
"data:image/jpeg;base64,aW1n",
|
||||
)
|
||||
|
||||
inline = payload["contents"][0]["parts"][1]["inlineData"]
|
||||
self.assertEqual(inline["mimeType"], "image/jpeg")
|
||||
self.assertEqual(inline["data"], "aW1n")
|
||||
|
||||
def test_build_images_payload(self):
|
||||
from services.ai_image_service import build_payload
|
||||
|
||||
payload = build_payload(
|
||||
{
|
||||
"url": "https://api.example.test/v1/images/generations",
|
||||
"model": "img",
|
||||
"api_key": "k",
|
||||
"api_type": "images",
|
||||
},
|
||||
"prompt",
|
||||
"data:image/png;base64,abc",
|
||||
resolution="1K",
|
||||
)
|
||||
|
||||
self.assertEqual(payload["image_urls"], ["data:image/png;base64,abc"])
|
||||
self.assertEqual(payload["aspect_ratio"], "1:1")
|
||||
self.assertEqual(payload["resolution"], "1K")
|
||||
|
||||
def test_build_multipart_fields(self):
|
||||
from services.ai_image_service import build_multipart_fields
|
||||
|
||||
data, files = build_multipart_fields(
|
||||
{
|
||||
"url": "https://api.example.test/v1/images/edits",
|
||||
"model": "edit",
|
||||
"api_key": "k",
|
||||
"api_type": "images_edits",
|
||||
},
|
||||
"prompt",
|
||||
self.image_path,
|
||||
resolution="1K",
|
||||
)
|
||||
try:
|
||||
self.assertEqual(data["model"], "edit")
|
||||
self.assertEqual(data["size"], "1024x1024")
|
||||
self.assertEqual(files["image"][0], "sample.png")
|
||||
self.assertEqual(files["image"][2], "image/png")
|
||||
finally:
|
||||
files["image"][1].close()
|
||||
|
||||
def test_extract_data_url_from_response(self):
|
||||
from services.ai_image_service import extract_image_from_response
|
||||
|
||||
encoded = base64.b64encode(b"image-bytes").decode("ascii")
|
||||
payload = {"choices": [{"message": {"content": "x", "image": "data:image/png;base64," + encoded}}]}
|
||||
|
||||
self.assertEqual(extract_image_from_response(payload), b"image-bytes")
|
||||
|
||||
def test_extract_base64_from_response(self):
|
||||
from services.ai_image_service import extract_image_from_response
|
||||
|
||||
encoded = base64.b64encode(b"image-bytes").decode("ascii")
|
||||
payload = {"data": [{"b64_json": encoded}]}
|
||||
|
||||
self.assertEqual(extract_image_from_response(payload), b"image-bytes")
|
||||
|
||||
def test_extract_image_url_downloads(self):
|
||||
from services.ai_image_service import extract_image_from_response
|
||||
|
||||
session = _FakeSession()
|
||||
payload = {"data": [{"url": "https://cdn.example.test/out.jpg"}]}
|
||||
|
||||
self.assertEqual(extract_image_from_response(payload, session=session), b"downloaded")
|
||||
self.assertEqual(session.last_url, "https://cdn.example.test/out.jpg")
|
||||
|
||||
|
||||
class TestImageApiClient(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = Path(tempfile.mkdtemp())
|
||||
self.image_path = self.tmp / "sample.png"
|
||||
self.image_path.write_bytes(b"\x89PNG\r\n\x1a\nfake-png")
|
||||
|
||||
def tearDown(self):
|
||||
shutil.rmtree(str(self.tmp), ignore_errors=True)
|
||||
|
||||
def test_generate_posts_chat_payload_and_extracts_image(self):
|
||||
from services.ai_image_service import ImageApiClient
|
||||
|
||||
session = _FakeSession()
|
||||
config = {
|
||||
"url": "https://api.example.test",
|
||||
"model": "model-x",
|
||||
"api_key": "secret",
|
||||
"api_type": "chat",
|
||||
}
|
||||
|
||||
result = ImageApiClient(config, session=session).generate("prompt", self.image_path)
|
||||
|
||||
self.assertEqual(result, b"image-bytes")
|
||||
self.assertEqual(session.last_post_url, "https://api.example.test/v1/chat/completions")
|
||||
self.assertEqual(session.last_headers["Authorization"], "Bearer secret")
|
||||
self.assertEqual(session.last_json["model"], "model-x")
|
||||
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, payload=None, content=b"downloaded"):
|
||||
self._payload = payload or {
|
||||
"data": [
|
||||
{"b64_json": base64.b64encode(b"image-bytes").decode("ascii")}
|
||||
]
|
||||
}
|
||||
self.content = content
|
||||
|
||||
def raise_for_status(self):
|
||||
return None
|
||||
|
||||
def json(self):
|
||||
return self._payload
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self):
|
||||
self.trust_env = True
|
||||
self.last_url = None
|
||||
self.last_post_url = None
|
||||
self.last_headers = None
|
||||
self.last_json = None
|
||||
|
||||
def get(self, url, timeout=60):
|
||||
self.last_url = url
|
||||
return _FakeResponse(content=b"downloaded")
|
||||
|
||||
def post(self, url, headers=None, json=None, data=None, files=None, timeout=None):
|
||||
self.last_post_url = url
|
||||
self.last_headers = headers
|
||||
self.last_json = json
|
||||
return _FakeResponse()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user