feat: 完成T-301 AI生成接口
新增app/ai.py,按config/ai_models.json读取默认文本和图片模型,提供gen_title与gen_cover通用HTTP接口。 支持chat JSON与images_edits multipart、失败重试、错误脱敏、图片URL/base64解析、resolution resize和jpg_quality保存。 新增tests/test_ai.py覆盖文本重试、缺字段错误不泄露Key、封面JPEG保存;同步任务看板、API、技术栈、架构、current-state和progress,并记录Tab①真机冒烟结果。
This commit is contained in:
@@ -0,0 +1,164 @@
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
from unittest import mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(__file__))
|
||||
|
||||
from _helpers import TempDirMixin
|
||||
|
||||
from app import ai, appconfig
|
||||
|
||||
|
||||
class _Response:
|
||||
def __init__(self, payload):
|
||||
self.payload = payload
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
def read(self, size=-1):
|
||||
return json.dumps(self.payload).encode("utf-8")
|
||||
|
||||
|
||||
class AITests(TempDirMixin, unittest.TestCase):
|
||||
def _write_models(self, path, text=None, image=None):
|
||||
text = text or {
|
||||
"name": "Text",
|
||||
"category": "text",
|
||||
"enabled": True,
|
||||
"url": "https://example.invalid/v1/chat/completions",
|
||||
"model": "text-model",
|
||||
"api_key": "sk-text-secret",
|
||||
"api_type": "chat",
|
||||
"connect_timeout_seconds": 1,
|
||||
"timeout_seconds": 1,
|
||||
"extra_body": {"temperature": 0},
|
||||
}
|
||||
image = image or {
|
||||
"name": "Image",
|
||||
"category": "image",
|
||||
"enabled": True,
|
||||
"url": "https://example.invalid/v1/chat/completions",
|
||||
"model": "image-model",
|
||||
"api_key": "sk-image-secret",
|
||||
"api_type": "auto",
|
||||
"connect_timeout_seconds": 1,
|
||||
"timeout_seconds": 1,
|
||||
"extra_body": {},
|
||||
}
|
||||
appconfig.save_ai_models_config({"models": [text, image]}, path=path)
|
||||
|
||||
def _config(self):
|
||||
cfg = appconfig.default_config()
|
||||
cfg["ai"]["default_text_model"] = "Text"
|
||||
cfg["ai"]["default_image_model"] = "Image"
|
||||
cfg["ai"]["retry"] = 1
|
||||
cfg["ai"]["resolution"] = "512"
|
||||
cfg["ai"]["jpg_quality"] = 80
|
||||
return cfg
|
||||
|
||||
def test_gen_title_uses_configured_model_and_retries(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
models_path = os.path.join(temp_dir, "ai_models.json")
|
||||
self._write_models(models_path)
|
||||
calls = []
|
||||
|
||||
def fake_urlopen(request, timeout=None):
|
||||
calls.append((request, timeout))
|
||||
if len(calls) == 1:
|
||||
raise ai.urllib.error.URLError("temporary")
|
||||
return _Response({"choices": [{"message": {"content": " 新标题 "}}]})
|
||||
|
||||
with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
title = ai.gen_title(
|
||||
"优化标题",
|
||||
"旧标题",
|
||||
config=self._config(),
|
||||
models_path=models_path,
|
||||
)
|
||||
|
||||
self.assertEqual("新标题", title)
|
||||
self.assertEqual(2, len(calls))
|
||||
body = json.loads(calls[-1][0].data.decode("utf-8"))
|
||||
self.assertEqual("text-model", body["model"])
|
||||
self.assertEqual(0, body["temperature"])
|
||||
self.assertNotIn("sk-text-secret", body["messages"][1]["content"])
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_missing_model_fields_raise_clear_error_without_secret(self):
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
models_path = os.path.join(temp_dir, "ai_models.json")
|
||||
self._write_models(
|
||||
models_path,
|
||||
text={
|
||||
"name": "Text",
|
||||
"category": "text",
|
||||
"enabled": True,
|
||||
"url": "",
|
||||
"model": "text-model",
|
||||
"api_key": "sk-text-secret",
|
||||
"api_type": "chat",
|
||||
"connect_timeout_seconds": 1,
|
||||
"timeout_seconds": 1,
|
||||
"extra_body": {},
|
||||
},
|
||||
)
|
||||
|
||||
with self.assertRaises(ai.AIError) as raised:
|
||||
ai.gen_title("prompt", "old", config=self._config(), models_path=models_path)
|
||||
|
||||
message = str(raised.exception)
|
||||
self.assertIn("缺少字段: url", message)
|
||||
self.assertNotIn("sk-text-secret", message)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
def test_gen_cover_saves_jpeg_with_resolution_and_quality(self):
|
||||
try:
|
||||
from PIL import Image
|
||||
except ImportError:
|
||||
self.skipTest("Pillow not installed")
|
||||
|
||||
with self.make_temp_dir() as temp_dir:
|
||||
models_path = os.path.join(temp_dir, "ai_models.json")
|
||||
self._write_models(models_path)
|
||||
old_cover = os.path.join(temp_dir, "old.jpg")
|
||||
output = os.path.join(temp_dir, "new.jpg")
|
||||
Image.new("RGB", (16, 16), (20, 30, 40)).save(old_cover, "JPEG")
|
||||
generated = io.BytesIO()
|
||||
Image.new("RGB", (8, 8), (200, 120, 80)).save(generated, "PNG")
|
||||
b64_image = base64.b64encode(generated.getvalue()).decode("ascii")
|
||||
|
||||
def fake_urlopen(request, timeout=None):
|
||||
body = json.loads(request.data.decode("utf-8"))
|
||||
self.assertEqual("image-model", body["model"])
|
||||
self.assertIn("目标分辨率:512", body["messages"][0]["content"][0]["text"])
|
||||
return _Response({"data": [{"b64_json": b64_image}]})
|
||||
|
||||
with mock.patch("app.ai.urllib.request.urlopen", side_effect=fake_urlopen):
|
||||
result = ai.gen_cover(
|
||||
"生成封面",
|
||||
old_cover,
|
||||
output,
|
||||
config=self._config(),
|
||||
models_path=models_path,
|
||||
)
|
||||
|
||||
self.assertEqual(os.path.abspath(output), result)
|
||||
with Image.open(output) as saved:
|
||||
self.assertEqual((512, 512), saved.size)
|
||||
self.assertEqual("JPEG", saved.format)
|
||||
|
||||
self.assert_removed(temp_dir)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user