77 lines
2.9 KiB
Python
77 lines
2.9 KiB
Python
from collections import Counter
|
|
from io import BytesIO
|
|
|
|
from PIL import Image
|
|
|
|
from app import llm
|
|
|
|
|
|
def test_generated_qa_is_normalized_to_conversational_text():
|
|
content = '{"items":[{"question":" 洗完会不会很干？？ ","answers":["回答：我用着还行。”", "感觉不会，挺清爽的！"]}]}'
|
|
|
|
result = llm._parse_content(content, expected_answers=2)
|
|
|
|
assert result[0]["question"] == "洗完会不会很干"
|
|
assert result[0]["answers"] == ["我用着还行", "感觉不会，挺清爽的"]
|
|
|
|
|
|
def test_generation_prompts_request_varied_buyer_language_without_question_marks():
|
|
qa_prompt = llm._prompt(llm.QAGenerateRequest(product_name="洗面奶"))
|
|
task_prompt = llm._questions_prompt({"product_name": "洗面奶"}, 5)
|
|
task_prompt_terms = (
|
|
"像买家随手输入",
|
|
"末尾不要加问号",
|
|
"不要每句都用",
|
|
)
|
|
|
|
assert all(term in qa_prompt for term in task_prompt_terms)
|
|
assert all(term in task_prompt for term in task_prompt_terms)
|
|
assert "不要写成客服说明" in qa_prompt
|
|
|
|
|
|
def test_multi_image_ocr_retries_per_image_and_deduplicates(monkeypatch):
|
|
attempts = Counter()
|
|
|
|
def recognize(image_data, content_type):
|
|
attempts[image_data] += 1
|
|
if image_data == b"first" and attempts[image_data] == 1:
|
|
raise TimeoutError("temporary timeout")
|
|
if image_data == b"first":
|
|
return [
|
|
{"product_title": "商品A", "question": "共同问题？"},
|
|
{"product_title": "商品A", "question": "第一张的问题？"},
|
|
]
|
|
return [
|
|
{"product_title": "商品A", "question": "共同问题？"},
|
|
{"product_title": "商品A", "question": "第二张的问题？"},
|
|
]
|
|
|
|
monkeypatch.setattr(llm.settings, "llm_api_key", "test-key")
|
|
monkeypatch.setattr(llm, "_compress_ocr_image", lambda data, content_type: (data, content_type))
|
|
monkeypatch.setattr(llm, "_extract_questions_from_image", recognize)
|
|
monkeypatch.setattr(llm.time, "sleep", lambda seconds: None)
|
|
|
|
result = llm.extract_questions_from_images([
|
|
(b"first", "image/png"),
|
|
(b"second", "image/png"),
|
|
])
|
|
|
|
assert attempts == Counter({b"first": 2, b"second": 1})
|
|
assert [item["question"] for item in result] == ["共同问题？", "第一张的问题？", "第二张的问题？"]
|
|
assert [item["source_index"] for item in result] == [0, 0, 1]
|
|
|
|
|
|
def test_ocr_image_compression_limits_dimensions_and_converts_to_jpeg():
|
|
image = Image.new("RGBA", (3200, 2600), (20, 100, 220, 180))
|
|
source = BytesIO()
|
|
image.save(source, format="PNG")
|
|
|
|
compressed, content_type = llm._compress_ocr_image(source.getvalue(), "image/png")
|
|
|
|
with Image.open(BytesIO(compressed)) as result:
|
|
assert max(result.size) <= 2200
|
|
assert result.mode == "RGB"
|
|
assert result.format == "JPEG"
|
|
assert content_type == "image/jpeg"
|
|
assert len(compressed) < len(source.getvalue())
|