diff --git a/tests/unit/test_text_completion/test_ollama_processor.py b/tests/unit/test_text_completion/test_ollama_processor.py index 138a8598..0bf5e0ab 100644 --- a/tests/unit/test_text_completion/test_ollama_processor.py +++ b/tests/unit/test_text_completion/test_ollama_processor.py @@ -312,6 +312,150 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): # Verify the combined prompt mock_client.generate.assert_called_once_with('llama2', "You are a helpful assistant\n\nWhat is AI?", options={'temperature': 0.0}) + @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_temperature_override(self, mock_llm_init, mock_async_init, mock_client_class): + """Test temperature parameter override functionality""" + # Arrange + mock_client = MagicMock() + mock_response = { + 'response': 'Response with custom temperature', + 'prompt_eval_count': 20, + 'eval_count': 12 + } + mock_client.generate.return_value = mock_response + mock_client_class.return_value = mock_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'llama2', + 'ollama': 'http://localhost:11434', + 'temperature': 0.0, # Default temperature + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override temperature at runtime + result = await processor.generate_content( + "System prompt", + "User prompt", + model=None, # Use default model + temperature=0.8 # Override temperature + ) + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Response with custom temperature" + + # Verify Ollama API was called with overridden temperature + mock_client.generate.assert_called_once_with( + 'llama2', + "System prompt\n\nUser prompt", + options={'temperature': 0.8} # Should use runtime override + ) + + @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_model_override(self, mock_llm_init, mock_async_init, mock_client_class): + """Test model parameter override functionality""" + # Arrange + mock_client = MagicMock() + mock_response = { + 'response': 'Response with custom model', + 'prompt_eval_count': 18, + 'eval_count': 14 + } + mock_client.generate.return_value = mock_response + mock_client_class.return_value = mock_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'llama2', # Default model + 'ollama': 'http://localhost:11434', + 'temperature': 0.1, # Default temperature + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override model at runtime + result = await processor.generate_content( + "System prompt", + "User prompt", + model="mistral", # Override model + temperature=None # Use default temperature + ) + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Response with custom model" + + # Verify Ollama API was called with overridden model + mock_client.generate.assert_called_once_with( + 'mistral', # Should use runtime override + "System prompt\n\nUser prompt", + options={'temperature': 0.1} # Should use processor default + ) + + @patch('trustgraph.model.text_completion.ollama.llm.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_both_parameters_override(self, mock_llm_init, mock_async_init, mock_client_class): + """Test overriding both model and temperature parameters simultaneously""" + # Arrange + mock_client = MagicMock() + mock_response = { + 'response': 'Response with both overrides', + 'prompt_eval_count': 22, + 'eval_count': 16 + } + mock_client.generate.return_value = mock_response + mock_client_class.return_value = mock_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'llama2', # Default model + 'ollama': 'http://localhost:11434', + 'temperature': 0.0, # Default temperature + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act - Override both parameters at runtime + result = await processor.generate_content( + "System prompt", + "User prompt", + model="codellama", # Override model + temperature=0.9 # Override temperature + ) + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Response with both overrides" + + # Verify Ollama API was called with both overrides + mock_client.generate.assert_called_once_with( + 'codellama', # Should use runtime override + "System prompt\n\nUser prompt", + options={'temperature': 0.9} # Should use runtime override + ) + if __name__ == '__main__': pytest.main([__file__]) \ No newline at end of file