diff --git a/tests/unit/test_text_completion/test_cohere_processor.py b/tests/unit/test_text_completion/test_cohere_processor.py new file mode 100644 index 00000000..ebb6b626 --- /dev/null +++ b/tests/unit/test_text_completion/test_cohere_processor.py @@ -0,0 +1,447 @@ +""" +Unit tests for trustgraph.model.text_completion.cohere +Following the same successful pattern as previous tests +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from unittest import IsolatedAsyncioTestCase + +# Import the service under test +from trustgraph.model.text_completion.cohere.llm import Processor +from trustgraph.base import LlmResult +from trustgraph.exceptions import TooManyRequests + + +class TestCohereProcessorSimple(IsolatedAsyncioTestCase): + """Test Cohere processor functionality""" + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_basic(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test basic processor initialization""" + # Arrange + mock_cohere_client = MagicMock() + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert processor.model == 'c4ai-aya-23-8b' + assert processor.temperature == 0.0 + assert hasattr(processor, 'cohere') + mock_cohere_class.assert_called_once_with(api_key='test-api-key') + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_success(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test successful content generation""" + # Arrange + mock_cohere_client = MagicMock() + mock_output = MagicMock() + mock_output.text = "Generated response from Cohere" + mock_output.meta.billed_units.input_tokens = 25 + mock_output.meta.billed_units.output_tokens = 15 + + mock_cohere_client.chat.return_value = mock_output + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("System prompt", "User prompt") + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Generated response from Cohere" + assert result.in_token == 25 + assert result.out_token == 15 + assert result.model == 'c4ai-aya-23-8b' + + # Verify the Cohere API call + mock_cohere_client.chat.assert_called_once_with( + model='c4ai-aya-23-8b', + message="User prompt", + preamble="System prompt", + temperature=0.0, + chat_history=[], + prompt_truncation='auto', + connectors=[] + ) + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_rate_limit_error(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test rate limit error handling""" + # Arrange + import cohere + + mock_cohere_client = MagicMock() + mock_cohere_client.chat.side_effect = cohere.TooManyRequestsError("Rate limit exceeded") + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act & Assert + with pytest.raises(TooManyRequests): + await processor.generate_content("System prompt", "User prompt") + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_generic_exception(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test handling of generic exceptions""" + # Arrange + mock_cohere_client = MagicMock() + mock_cohere_client.chat.side_effect = Exception("API connection error") + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act & Assert + with pytest.raises(Exception, match="API connection error"): + await processor.generate_content("System prompt", "User prompt") + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_without_api_key(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test processor initialization without API key (should fail)""" + # Arrange + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': None, # No API key provided + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act & Assert + with pytest.raises(RuntimeError, match="Cohere API key not specified"): + processor = Processor(**config) + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_with_custom_parameters(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test processor initialization with custom parameters""" + # Arrange + mock_cohere_client = MagicMock() + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'command-light', + 'api_key': 'custom-api-key', + 'temperature': 0.7, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert processor.model == 'command-light' + assert processor.temperature == 0.7 + mock_cohere_class.assert_called_once_with(api_key='custom-api-key') + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_with_defaults(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test processor initialization with default values""" + # Arrange + mock_cohere_client = MagicMock() + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + # Only provide required fields, should use defaults + config = { + 'api_key': 'test-api-key', + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert processor.model == 'c4ai-aya-23-8b' # default_model + assert processor.temperature == 0.0 # default_temperature + mock_cohere_class.assert_called_once_with(api_key='test-api-key') + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_empty_prompts(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test content generation with empty prompts""" + # Arrange + mock_cohere_client = MagicMock() + mock_output = MagicMock() + mock_output.text = "Default response" + mock_output.meta.billed_units.input_tokens = 2 + mock_output.meta.billed_units.output_tokens = 3 + + mock_cohere_client.chat.return_value = mock_output + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("", "") + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Default response" + assert result.in_token == 2 + assert result.out_token == 3 + assert result.model == 'c4ai-aya-23-8b' + + # Verify the preamble and message are handled correctly + call_args = mock_cohere_client.chat.call_args + assert call_args[1]['preamble'] == "" + assert call_args[1]['message'] == "" + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_chat_structure(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test that Cohere chat is structured correctly""" + # Arrange + mock_cohere_client = MagicMock() + mock_output = MagicMock() + mock_output.text = "Response with proper structure" + mock_output.meta.billed_units.input_tokens = 30 + mock_output.meta.billed_units.output_tokens = 20 + + mock_cohere_client.chat.return_value = mock_output + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.5, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("You are a helpful assistant", "What is AI?") + + # Assert + assert result.text == "Response with proper structure" + assert result.in_token == 30 + assert result.out_token == 20 + + # Verify the chat structure matches Cohere API format + call_args = mock_cohere_client.chat.call_args + + # Check parameters + assert call_args[1]['model'] == 'c4ai-aya-23-8b' + assert call_args[1]['message'] == "What is AI?" + assert call_args[1]['preamble'] == "You are a helpful assistant" + assert call_args[1]['temperature'] == 0.5 + assert call_args[1]['chat_history'] == [] + assert call_args[1]['prompt_truncation'] == 'auto' + assert call_args[1]['connectors'] == [] + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_token_parsing(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test token parsing from Cohere response""" + # Arrange + mock_cohere_client = MagicMock() + mock_output = MagicMock() + mock_output.text = "Token parsing test" + mock_output.meta.billed_units.input_tokens = 50 + mock_output.meta.billed_units.output_tokens = 25 + + mock_cohere_client.chat.return_value = mock_output + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("System", "User query") + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Token parsing test" + assert result.in_token == 50 + assert result.out_token == 25 + assert result.model == 'c4ai-aya-23-8b' + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_cohere_client_initialization(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test that Cohere client is initialized correctly""" + # Arrange + mock_cohere_client = MagicMock() + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'command-r', + 'api_key': 'co-test-key', + 'temperature': 0.0, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + # Verify Cohere client was called with correct API key + mock_cohere_class.assert_called_once_with(api_key='co-test-key') + + # Verify processor has the client + assert processor.cohere == mock_cohere_client + assert processor.model == 'command-r' + + @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_chat_parameters(self, mock_llm_init, mock_async_init, mock_cohere_class): + """Test that all chat parameters are passed correctly""" + # Arrange + mock_cohere_client = MagicMock() + mock_output = MagicMock() + mock_output.text = "Chat parameter test" + mock_output.meta.billed_units.input_tokens = 20 + mock_output.meta.billed_units.output_tokens = 10 + + mock_cohere_client.chat.return_value = mock_output + mock_cohere_class.return_value = mock_cohere_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'c4ai-aya-23-8b', + 'api_key': 'test-api-key', + 'temperature': 0.3, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("System instructions", "User question") + + # Assert + assert result.text == "Chat parameter test" + + # Verify all parameters are passed correctly + call_args = mock_cohere_client.chat.call_args + assert call_args[1]['model'] == 'c4ai-aya-23-8b' + assert call_args[1]['message'] == "User question" + assert call_args[1]['preamble'] == "System instructions" + assert call_args[1]['temperature'] == 0.3 + assert call_args[1]['chat_history'] == [] + assert call_args[1]['prompt_truncation'] == 'auto' + assert call_args[1]['connectors'] == [] + + +if __name__ == '__main__': + pytest.main([__file__]) \ No newline at end of file diff --git a/tests/unit/test_text_completion/test_googleaistudio_processor.py b/tests/unit/test_text_completion/test_googleaistudio_processor.py new file mode 100644 index 00000000..a3ca0057 --- /dev/null +++ b/tests/unit/test_text_completion/test_googleaistudio_processor.py @@ -0,0 +1,482 @@ +""" +Unit tests for trustgraph.model.text_completion.googleaistudio +Following the same successful pattern as previous tests +""" + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch +from unittest import IsolatedAsyncioTestCase + +# Import the service under test +from trustgraph.model.text_completion.googleaistudio.llm import Processor +from trustgraph.base import LlmResult +from trustgraph.exceptions import TooManyRequests + + +class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): + """Test Google AI Studio processor functionality""" + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_basic(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test basic processor initialization""" + # Arrange + mock_genai_client = MagicMock() + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert processor.model == 'gemini-2.0-flash-001' + assert processor.temperature == 0.0 + assert processor.max_output == 8192 + assert hasattr(processor, 'client') + assert hasattr(processor, 'safety_settings') + assert len(processor.safety_settings) == 4 # 4 safety categories + mock_genai_class.assert_called_once_with(api_key='test-api-key') + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_success(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test successful content generation""" + # Arrange + mock_genai_client = MagicMock() + mock_response = MagicMock() + mock_response.text = "Generated response from Google AI Studio" + mock_response.usage_metadata.prompt_token_count = 25 + mock_response.usage_metadata.candidates_token_count = 15 + + mock_genai_client.models.generate_content.return_value = mock_response + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("System prompt", "User prompt") + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Generated response from Google AI Studio" + assert result.in_token == 25 + assert result.out_token == 15 + assert result.model == 'gemini-2.0-flash-001' + + # Verify the Google AI Studio API call + mock_genai_client.models.generate_content.assert_called_once() + call_args = mock_genai_client.models.generate_content.call_args + assert call_args[1]['model'] == 'gemini-2.0-flash-001' + assert call_args[1]['contents'] == "User prompt" + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_rate_limit_error(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test rate limit error handling""" + # Arrange + from google.api_core.exceptions import ResourceExhausted + + mock_genai_client = MagicMock() + mock_genai_client.models.generate_content.side_effect = ResourceExhausted("Rate limit exceeded") + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act & Assert + with pytest.raises(TooManyRequests): + await processor.generate_content("System prompt", "User prompt") + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_generic_exception(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test handling of generic exceptions""" + # Arrange + mock_genai_client = MagicMock() + mock_genai_client.models.generate_content.side_effect = Exception("API connection error") + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act & Assert + with pytest.raises(Exception, match="API connection error"): + await processor.generate_content("System prompt", "User prompt") + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_without_api_key(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test processor initialization without API key (should fail)""" + # Arrange + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': None, # No API key provided + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act & Assert + with pytest.raises(RuntimeError, match="Google AI Studio API key not specified"): + processor = Processor(**config) + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_with_custom_parameters(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test processor initialization with custom parameters""" + # Arrange + mock_genai_client = MagicMock() + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-1.5-pro', + 'api_key': 'custom-api-key', + 'temperature': 0.7, + 'max_output': 4096, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert processor.model == 'gemini-1.5-pro' + assert processor.temperature == 0.7 + assert processor.max_output == 4096 + mock_genai_class.assert_called_once_with(api_key='custom-api-key') + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_processor_initialization_with_defaults(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test processor initialization with default values""" + # Arrange + mock_genai_client = MagicMock() + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + # Only provide required fields, should use defaults + config = { + 'api_key': 'test-api-key', + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert processor.model == 'gemini-2.0-flash-001' # default_model + assert processor.temperature == 0.0 # default_temperature + assert processor.max_output == 8192 # default_max_output + mock_genai_class.assert_called_once_with(api_key='test-api-key') + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_empty_prompts(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test content generation with empty prompts""" + # Arrange + mock_genai_client = MagicMock() + mock_response = MagicMock() + mock_response.text = "Default response" + mock_response.usage_metadata.prompt_token_count = 2 + mock_response.usage_metadata.candidates_token_count = 3 + + mock_genai_client.models.generate_content.return_value = mock_response + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("", "") + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Default response" + assert result.in_token == 2 + assert result.out_token == 3 + assert result.model == 'gemini-2.0-flash-001' + + # Verify the system instruction and content are handled correctly + call_args = mock_genai_client.models.generate_content.call_args + assert call_args[1]['contents'] == "" + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_configuration_structure(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test that generation configuration is structured correctly""" + # Arrange + mock_genai_client = MagicMock() + mock_response = MagicMock() + mock_response.text = "Response with proper structure" + mock_response.usage_metadata.prompt_token_count = 30 + mock_response.usage_metadata.candidates_token_count = 20 + + mock_genai_client.models.generate_content.return_value = mock_response + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.5, + 'max_output': 1024, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("You are a helpful assistant", "What is AI?") + + # Assert + assert result.text == "Response with proper structure" + assert result.in_token == 30 + assert result.out_token == 20 + + # Verify the generation configuration + call_args = mock_genai_client.models.generate_content.call_args + config_arg = call_args[1]['config'] + + # Check that the configuration has the right structure + assert call_args[1]['model'] == 'gemini-2.0-flash-001' + assert call_args[1]['contents'] == "What is AI?" + # Config should be a GenerateContentConfig object + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_safety_settings_initialization(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test that safety settings are initialized correctly""" + # Arrange + mock_genai_client = MagicMock() + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + assert hasattr(processor, 'safety_settings') + assert len(processor.safety_settings) == 4 + # Should have 4 safety categories: hate speech, harassment, sexually explicit, dangerous content + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_token_parsing(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test token parsing from Google AI Studio response""" + # Arrange + mock_genai_client = MagicMock() + mock_response = MagicMock() + mock_response.text = "Token parsing test" + mock_response.usage_metadata.prompt_token_count = 50 + mock_response.usage_metadata.candidates_token_count = 25 + + mock_genai_client.models.generate_content.return_value = mock_response + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("System", "User query") + + # Assert + assert isinstance(result, LlmResult) + assert result.text == "Token parsing test" + assert result.in_token == 50 + assert result.out_token == 25 + assert result.model == 'gemini-2.0-flash-001' + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_genai_client_initialization(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test that Google AI Studio client is initialized correctly""" + # Arrange + mock_genai_client = MagicMock() + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-1.5-flash', + 'api_key': 'gai-test-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + # Verify Google AI Studio client was called with correct API key + mock_genai_class.assert_called_once_with(api_key='gai-test-key') + + # Verify processor has the client + assert processor.client == mock_genai_client + assert processor.model == 'gemini-1.5-flash' + + @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_system_instruction(self, mock_llm_init, mock_async_init, mock_genai_class): + """Test that system instruction is handled correctly""" + # Arrange + mock_genai_client = MagicMock() + mock_response = MagicMock() + mock_response.text = "System instruction test" + mock_response.usage_metadata.prompt_token_count = 35 + mock_response.usage_metadata.candidates_token_count = 25 + + mock_genai_client.models.generate_content.return_value = mock_response + mock_genai_class.return_value = mock_genai_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'gemini-2.0-flash-001', + 'api_key': 'test-api-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + processor = Processor(**config) + + # Act + result = await processor.generate_content("Be helpful and concise", "Explain quantum computing") + + # Assert + assert result.text == "System instruction test" + assert result.in_token == 35 + assert result.out_token == 25 + + # Verify the system instruction is passed in the config + call_args = mock_genai_client.models.generate_content.call_args + config_arg = call_args[1]['config'] + # The system instruction should be in the config object + assert call_args[1]['contents'] == "Explain quantum computing" + + +if __name__ == '__main__': + pytest.main([__file__]) \ No newline at end of file