| import unittest | |
| from app.web import GemmaPromptWrapper | |
| class FakeRunnable: | |
| def __init__(self): | |
| self.last_input = None | |
| def invoke(self, input_value, *args, **kwargs): | |
| self.last_input = input_value | |
| return {"ok": True} | |
| class FakeLLM: | |
| def __init__(self): | |
| self.last_input = None | |
| self.structured_runnable = FakeRunnable() | |
| def invoke(self, input_value, *args, **kwargs): | |
| self.last_input = input_value | |
| return {"ok": True} | |
| def with_structured_output(self, *args, **kwargs): | |
| return self.structured_runnable | |
| class GemmaPromptWrapperTests(unittest.TestCase): | |
| def test_prepends_thought_off_for_gemma_string_prompts(self): | |
| llm = FakeLLM() | |
| wrapped = GemmaPromptWrapper(llm, "gemma-3-27b-it") | |
| wrapped.invoke("hello world") | |
| self.assertEqual(llm.last_input, "<thought off>\nhello world") | |
| def test_does_not_duplicate_thought_off(self): | |
| llm = FakeLLM() | |
| wrapped = GemmaPromptWrapper(llm, "gemma-3-27b-it") | |
| wrapped.invoke("<thought off>\nhello world") | |
| self.assertEqual(llm.last_input, "<thought off>\nhello world") | |
| def test_leaves_non_gemma_models_unchanged(self): | |
| llm = FakeLLM() | |
| wrapped = GemmaPromptWrapper(llm, "llama-3.1-8b-instant") | |
| wrapped.invoke("hello world") | |
| self.assertEqual(llm.last_input, "hello world") | |
| def test_wraps_structured_output_prompts_too(self): | |
| llm = FakeLLM() | |
| wrapped = GemmaPromptWrapper(llm, "gemma-3-27b-it") | |
| runnable = wrapped.with_structured_output(object) | |
| runnable.invoke("structured prompt") | |
| self.assertEqual(llm.structured_runnable.last_input, "<thought off>\nstructured prompt") | |
| if __name__ == "__main__": | |
| unittest.main() | |