mediaplayer/sentence_api/tests/test_ai_tutor.py

import pytest

import json
import re
import threading

from sentence_api.ai_tutor import AITutor, AITutorError, _apply_translation_terms, _parse_translation_payload


class FakeResponse:
    def __init__(self, status_code=200, body=None):
        self.status_code = status_code
        self.text = str(body)
        self._body = body or {}

    def json(self):
        if self._body is None:
            raise ValueError("invalid JSON")
        return self._body


def test_ai_tutor_parses_chat_completion(monkeypatch):
    requests_payloads = []

    def fake_post(endpoint, headers, json, timeout):
        requests_payloads.append((endpoint, headers, json, timeout))
        return FakeResponse(
            body={
                "choices": [
                    {"message": {"content": " 这句话的意思是…… "}}
                ]
            }
        )

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test/chat", "secret", "test-model", 1.5)

    explanation = tutor.explain(
        "The early bird catches the worm.",
        course_title="English idioms",
        sentence_index=1,
        language="en",
    )

    endpoint, headers, payload, timeout = requests_payloads[0]
    assert explanation == "这句话的意思是……"
    assert endpoint == "https://ai.test/chat"
    assert headers["Authorization"] == "Bearer secret"
    assert payload["model"] == "test-model"
    assert "The early bird catches the worm." in payload["messages"][1]["content"]
    assert timeout == 1.5


def test_ai_tutor_resolves_openai_compatible_base_url(monkeypatch):
    endpoints = []

    def fake_post(endpoint, **kwargs):
        endpoints.append(endpoint)
        return FakeResponse(body={"choices": [{"message": {"content": "ok"}}]})

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test:8080/v1", "secret", "test-model", 1)

    assert tutor.explain("Hello.") == "ok"
    assert endpoints == ["https://ai.test:8080/v1/chat/completions"]
    assert tutor.request_url == "https://ai.test:8080/v1/chat/completions"


def test_ai_tutor_rejects_unexpected_response(monkeypatch):
    monkeypatch.setattr(
        "sentence_api.ai_tutor.requests.post",
        lambda *args, **kwargs: FakeResponse(body={"choices": []}),
    )
    tutor = AITutor("https://ai.test/chat", "secret", "test-model", 1)

    with pytest.raises(AITutorError):
        tutor.explain("Hello.")



def test_ai_tutor_translates_full_document(monkeypatch):
    requests_payloads = []

    def fake_post(endpoint, headers, json, timeout):
        requests_payloads.append((endpoint, headers, json, timeout))
        return FakeResponse(body={"choices": [{"message": {"content": '{"0":"你好。","1":"见到你很高兴。"}'}}]})

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test/v1", "secret", "test-model", 1.5)

    translated = tutor.translate_sentences(
        [(0, "Hello."), (1, "Nice to meet you.")],
        course_title="Greetings",
    )

    assert translated == {0: "你好。", 1: "见到你很高兴。"}
    _endpoint, _headers, payload, timeout = requests_payloads[0]
    user_prompt = payload["messages"][1]["content"]
    assert "[0] Hello." in user_prompt
    assert "[1] Nice to meet you." in user_prompt
    assert timeout == 45.0


def _translation_response_from_prompt(payload):
    user_prompt = payload["messages"][1]["content"]
    section = user_prompt.split("待翻译句子：", 1)[1]
    indexes = [int(match) for match in re.findall(r"^\[(\d+)\]", section, re.M)]
    body = {str(index): f"中文{index}" for index in indexes}
    return FakeResponse(
        body={"choices": [{"message": {"content": json.dumps(body, ensure_ascii=False)}}]}
    )


def test_ai_tutor_translates_long_input_in_batches(monkeypatch):
    requests_payloads = []
    lock = threading.Lock()

    def fake_post(endpoint, headers, json, timeout):
        with lock:
            requests_payloads.append(json)
        return _translation_response_from_prompt(json)

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test/v1", "secret", "test-model", 1)
    sentences = [(index, f"This is sentence {index} with enough text.") for index in range(30)]

    translated = tutor.translate_sentences(sentences, course_title="Long talk")

    assert translated == {index: f"中文{index}" for index in range(30)}
    assert len(requests_payloads) >= 3
    first_prompt = requests_payloads[0]["messages"][1]["content"]
    assert "前文" in first_prompt or "后文" in first_prompt or "待翻译句子：" in first_prompt


def test_ai_tutor_skips_cached_sentences_and_persists_windows(monkeypatch):
    persisted = []

    def fake_post(endpoint, headers, json, timeout):
        return _translation_response_from_prompt(json)

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test/v1", "secret", "test-model", 1)
    translated = tutor.translate_sentences(
        [(0, "Hello."), (1, "Thanks."), (2, "Bye.")],
        existing={0: "你好。"},
        persist=persisted.append,
        course_title="Greetings",
    )

    assert translated[0] == "你好。"
    assert translated[1] == "中文1"
    assert translated[2] == "中文2"
    assert persisted
    assert 0 not in persisted[0]


def test_ai_tutor_splits_invalid_windows_and_retries(monkeypatch):
    calls = {"count": 0}

    def fake_post(endpoint, headers, json, timeout):
        calls["count"] += 1
        if calls["count"] == 1:
            return FakeResponse(body={"choices": [{"message": {"content": "not-json"}}]})
        return _translation_response_from_prompt(json)

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test/v1", "secret", "test-model", 1)
    sentences = [(index, f"This is sentence {index} with enough text.") for index in range(12)]

    translated = tutor.translate_sentences(sentences)

    assert translated == {index: f"中文{index}" for index in range(12)}
    assert calls["count"] >= 3


def test_ai_tutor_includes_name_terms_and_rewrites_output(monkeypatch):
    requests_payloads = []

    def fake_post(endpoint, headers, json, timeout):
        requests_payloads.append(json)
        return FakeResponse(
            body={"choices": [{"message": {"content": '{"0":"哈达马托击败了 Harimoto。"}'}}]}
        )

    monkeypatch.setattr("sentence_api.ai_tutor.requests.post", fake_post)
    tutor = AITutor("https://ai.test/v1", "secret", "test-model", 1)

    translated = tutor.translate_sentences(
        [(0, "Harimoto won.")],
        course_title="Table tennis",
        terms=[("Harimoto", "哈瑞么托"), ("哈达马托", "哈瑞么托")],
    )

    assert translated == {0: "哈瑞么托击败了 哈瑞么托。"}
    user_prompt = requests_payloads[0]["messages"][1]["content"]
    assert "Harimoto => 哈瑞么托" in user_prompt
    assert "哈达马托 => 哈瑞么托" in user_prompt


def test_apply_translation_terms_replaces_english_and_wrong_chinese():
    assert _apply_translation_terms(
        "Harimoto won.",
        "哈达马托击败了 Harimoto。",
        [("Harimoto", "哈瑞么托"), ("哈达马托", "哈瑞么托")],
    ) == "哈瑞么托击败了 哈瑞么托。"


def test_parse_translation_payload_accepts_fenced_json_and_numbered_lines():
    assert _parse_translation_payload('```json\n["一", "二"]\n```', 2) == {0: "一", 1: "二"}
    assert _parse_translation_payload("1. 一\n2. 二", 2) == {0: "一", 1: "二"}
    assert _parse_translation_payload('prefix {"12":"甲","13":"乙"} suffix', [12, 13]) == {
        12: "甲",
        13: "乙",
    }