85 lines
4.4 KiB
Python
85 lines
4.4 KiB
Python
import tempfile
|
|
import unittest
|
|
from pathlib import Path
|
|
from types import SimpleNamespace
|
|
from unittest.mock import Mock, patch
|
|
|
|
from google.genai import types
|
|
|
|
from copienator import configuration, prompting
|
|
from copienator.commands import correction, gemini_for_enonce, gemini_for_labels
|
|
from copienator.gemini_config import thinking_config
|
|
|
|
|
|
class GeminiConfigTests(unittest.TestCase):
|
|
def assert_request_config(self, config, level):
|
|
data = config.model_dump(mode="json", by_alias=True, exclude_none=True)
|
|
for key in ("temperature", "topP", "topK", "thinkingBudget"):
|
|
self.assertNotIn(key, data)
|
|
self.assertNotIn(key, data.get("thinkingConfig", {}))
|
|
if level is None:
|
|
self.assertNotIn("thinkingConfig", data)
|
|
else:
|
|
self.assertEqual(data["thinkingConfig"], {"thinkingLevel": level.upper()})
|
|
|
|
def test_grading_live_and_batch_settings(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
image = Path(directory) / "answer.jpg"
|
|
image.write_bytes(b"image")
|
|
image.with_suffix(".json").write_text('[["01", 0, 400, 1.0, "Ex 1"]]')
|
|
for level in ("minimal", "low", "medium", "high", None):
|
|
with self.subTest(level=level), patch.object(prompting, "make_prompt", return_value="Grade"):
|
|
_, config = prompting.generate_request(directory, image, "Ex 1", thinking_level=level)
|
|
self.assert_request_config(config, level)
|
|
batch = prompting.batch_generation_config(level)
|
|
self.assert_request_config(types.GenerateContentConfig.model_validate(batch), level)
|
|
self.assertEqual(batch["responseSchema"], prompting.UNROLLED_SCHEMA)
|
|
|
|
def test_labels_and_statement_use_their_own_settings(self):
|
|
with tempfile.TemporaryDirectory() as directory:
|
|
image = Path(directory) / "answer.jpg"
|
|
image.write_bytes(b"image")
|
|
with patch.object(configuration, "THINKING_LEVEL_FOR_LABEL", "low"):
|
|
_, config = gemini_for_labels.generate_request(image, "Ex 1", "Student", [], [])
|
|
self.assert_request_config(config, "low")
|
|
client = Mock()
|
|
with patch.object(configuration, "THINKING_LEVEL_LITE", "high"):
|
|
gemini_for_enonce._request_rubrics(client, "Question")
|
|
self.assert_request_config(client.models.generate_content.call_args.kwargs["config"], "high")
|
|
|
|
def test_correction_helpers_use_flash_level(self):
|
|
with patch.object(configuration, "THINKING_LEVEL_FLASH", "low"), patch.object(
|
|
prompting, "get_single_image_bytes", return_value=b"image"
|
|
):
|
|
requests = [
|
|
prompting.request_for_box_correction("answer.pdf", []),
|
|
prompting.request_for_wrong_label("answer.pdf", "Ex 1", "Question", "Ex 1"),
|
|
prompting.request_for_additional_answer("answer.pdf", "Ex 1", "Question", "Ex 1"),
|
|
]
|
|
for _, config in requests:
|
|
self.assert_request_config(config, "low")
|
|
|
|
def test_quota_fallback_switches_thinking_level_without_mutating_request(self):
|
|
client = Mock()
|
|
client.models.generate_content_stream.side_effect = [
|
|
RuntimeError("429 daily quota exhausted"), [SimpleNamespace(text="result")],
|
|
]
|
|
config = types.GenerateContentConfig(thinking_config=thinking_config("high"))
|
|
with patch.object(correction, "client", client), patch.object(
|
|
correction, "MODEL_ID_pro", "pro-model"
|
|
), patch.object(correction, "pro_quota_exhausted", False), patch.object(
|
|
configuration, "THINKING_LEVEL_FLASH", "low"
|
|
), patch.object(correction, "tprint"), patch.object(correction, "stop_requested") as stop:
|
|
stop.is_set.return_value = False
|
|
result = correction.call_gemini_with_retries("pro-model", [], config, "flash-model")
|
|
self.assertEqual(result, "result")
|
|
calls = client.models.generate_content_stream.call_args_list
|
|
self.assert_request_config(calls[0].kwargs["config"], "high")
|
|
self.assert_request_config(calls[1].kwargs["config"], "low")
|
|
self.assertEqual(calls[1].kwargs["model"], "flash-model")
|
|
self.assert_request_config(config, "high")
|
|
|
|
def test_invalid_level_fails_before_submission(self):
|
|
with self.assertRaisesRegex(ValueError, "Invalid Gemini thinking level"):
|
|
thinking_config("ultra")
|