From eb808573dcad84b6b23954794879d0463df70234 Mon Sep 17 00:00:00 2001 From: Cyber MacGeddon Date: Fri, 11 Jul 2025 17:09:10 +0100 Subject: [PATCH] Added more LLMs --- tests/unit/test_text_completion/conftest.py | 88 +++- .../test_claude_processor.py | 436 ++++++++++++++++++ 2 files changed, 523 insertions(+), 1 deletion(-) create mode 100644 tests/unit/test_text_completion/test_claude_processor.py diff --git a/tests/unit/test_text_completion/conftest.py b/tests/unit/test_text_completion/conftest.py index 2424cbdd..830f025e 100644 --- a/tests/unit/test_text_completion/conftest.py +++ b/tests/unit/test_text_completion/conftest.py @@ -220,4 +220,90 @@ def mock_openai_client(): def mock_openai_rate_limit_error(): """Mock OpenAI rate limit error""" from openai import RateLimitError - return RateLimitError("Rate limit exceeded", response=MagicMock(), body=None) \ No newline at end of file + return RateLimitError("Rate limit exceeded", response=MagicMock(), body=None) + + +# === Azure OpenAI Specific Fixtures === + +@pytest.fixture +def azure_openai_processor_config(base_processor_config): + """Default configuration for Azure OpenAI processor""" + config = base_processor_config.copy() + config.update({ + 'model': 'gpt-4', + 'endpoint': 'https://test.openai.azure.com/', + 'token': 'test-token', + 'api_version': '2024-12-01-preview', + 'temperature': 0.0, + 'max_output': 4192 + }) + return config + + +@pytest.fixture +def mock_azure_openai_client(): + """Mock Azure OpenAI client""" + mock_client = MagicMock() + + # Mock the response structure + mock_response = MagicMock() + mock_response.choices = [MagicMock()] + mock_response.choices[0].message.content = "Test response from Azure OpenAI" + mock_response.usage.prompt_tokens = 20 + mock_response.usage.completion_tokens = 10 + + mock_client.chat.completions.create.return_value = mock_response + return mock_client + + +@pytest.fixture +def mock_azure_openai_rate_limit_error(): + """Mock Azure OpenAI rate limit error""" + from openai import RateLimitError + return RateLimitError("Rate limit exceeded", response=MagicMock(), body=None) + + +# === Azure Specific Fixtures === + +@pytest.fixture +def azure_processor_config(base_processor_config): + """Default configuration for Azure processor""" + config = base_processor_config.copy() + config.update({ + 'endpoint': 'https://test.inference.ai.azure.com/v1/chat/completions', + 'token': 'test-token', + 'temperature': 0.0, + 'max_output': 4192 + }) + return config + + +@pytest.fixture +def mock_azure_requests(): + """Mock requests for Azure processor""" + mock_requests = MagicMock() + + # Mock successful response + mock_response = MagicMock() + mock_response.status_code = 200 + mock_response.json.return_value = { + 'choices': [{ + 'message': { + 'content': 'Test response from Azure' + } + }], + 'usage': { + 'prompt_tokens': 18, + 'completion_tokens': 9 + } + } + mock_requests.post.return_value = mock_response + return mock_requests + + +@pytest.fixture +def mock_azure_rate_limit_response(): + """Mock Azure rate limit response""" + mock_response = MagicMock() + mock_response.status_code = 429 + return mock_response \ No newline at end of file diff --git a/tests/unit/test_text_completion/test_claude_processor.py b/tests/unit/test_text_completion/test_claude_processor.py new file mode 100644 index 00000000..ffa633a2 --- /dev/null +++ b/tests/unit/test_text_completion/test_claude_processor.py @@ -0,0 +1,436 @@ +""" +Unit tests for trustgraph.model.text_completion.claude +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.claude.llm import Processor +from trustgraph.base import LlmResult +from trustgraph.exceptions import TooManyRequests + + +class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): + """Test Claude processor functionality""" + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test basic processor initialization""" + # Arrange + mock_claude_client = MagicMock() + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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 == 'claude-3-5-sonnet-20240620' + assert processor.temperature == 0.0 + assert processor.max_output == 8192 + assert hasattr(processor, 'claude') + mock_anthropic_class.assert_called_once_with(api_key='test-api-key') + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test successful content generation""" + # Arrange + mock_claude_client = MagicMock() + mock_response = MagicMock() + mock_response.content = [MagicMock()] + mock_response.content[0].text = "Generated response from Claude" + mock_response.usage.input_tokens = 25 + mock_response.usage.output_tokens = 15 + + mock_claude_client.messages.create.return_value = mock_response + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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 Claude" + assert result.in_token == 25 + assert result.out_token == 15 + assert result.model == 'claude-3-5-sonnet-20240620' + + # Verify the Claude API call + mock_claude_client.messages.create.assert_called_once_with( + model='claude-3-5-sonnet-20240620', + max_tokens=8192, + temperature=0.0, + system="System prompt", + messages=[{ + "role": "user", + "content": [{ + "type": "text", + "text": "User prompt" + }] + }] + ) + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test rate limit error handling""" + # Arrange + import anthropic + + mock_claude_client = MagicMock() + mock_claude_client.messages.create.side_effect = anthropic.RateLimitError("Rate limit exceeded") + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test handling of generic exceptions""" + # Arrange + mock_claude_client = MagicMock() + mock_claude_client.messages.create.side_effect = Exception("API connection error") + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test processor initialization without API key (should fail)""" + # Arrange + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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="Claude API key not specified"): + processor = Processor(**config) + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test processor initialization with custom parameters""" + # Arrange + mock_claude_client = MagicMock() + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-haiku-20240307', + '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 == 'claude-3-haiku-20240307' + assert processor.temperature == 0.7 + assert processor.max_output == 4096 + mock_anthropic_class.assert_called_once_with(api_key='custom-api-key') + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test processor initialization with default values""" + # Arrange + mock_claude_client = MagicMock() + mock_anthropic_class.return_value = mock_claude_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 == 'claude-3-5-sonnet-20240620' # default_model + assert processor.temperature == 0.0 # default_temperature + assert processor.max_output == 8192 # default_max_output + mock_anthropic_class.assert_called_once_with(api_key='test-api-key') + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @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_anthropic_class): + """Test content generation with empty prompts""" + # Arrange + mock_claude_client = MagicMock() + mock_response = MagicMock() + mock_response.content = [MagicMock()] + mock_response.content[0].text = "Default response" + mock_response.usage.input_tokens = 2 + mock_response.usage.output_tokens = 3 + + mock_claude_client.messages.create.return_value = mock_response + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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 == 'claude-3-5-sonnet-20240620' + + # Verify the system prompt and user content are handled correctly + call_args = mock_claude_client.messages.create.call_args + assert call_args[1]['system'] == "" + assert call_args[1]['messages'][0]['content'][0]['text'] == "" + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_message_structure(self, mock_llm_init, mock_async_init, mock_anthropic_class): + """Test that Claude messages are structured correctly""" + # Arrange + mock_claude_client = MagicMock() + mock_response = MagicMock() + mock_response.content = [MagicMock()] + mock_response.content[0].text = "Response with proper structure" + mock_response.usage.input_tokens = 30 + mock_response.usage.output_tokens = 20 + + mock_claude_client.messages.create.return_value = mock_response + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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 message structure matches Claude API format + call_args = mock_claude_client.messages.create.call_args + + # Check system prompt + assert call_args[1]['system'] == "You are a helpful assistant" + + # Check user message structure + messages = call_args[1]['messages'] + assert len(messages) == 1 + assert messages[0]['role'] == 'user' + assert messages[0]['content'][0]['type'] == 'text' + assert messages[0]['content'][0]['text'] == "What is AI?" + + # Verify other parameters + assert call_args[1]['model'] == 'claude-3-5-sonnet-20240620' + assert call_args[1]['temperature'] == 0.5 + assert call_args[1]['max_tokens'] == 1024 + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_generate_content_multiple_content_blocks(self, mock_llm_init, mock_async_init, mock_anthropic_class): + """Test handling of multiple content blocks in response""" + # Arrange + mock_claude_client = MagicMock() + mock_response = MagicMock() + + # Mock multiple content blocks (Claude can return multiple) + mock_content_1 = MagicMock() + mock_content_1.text = "First part of response" + mock_content_2 = MagicMock() + mock_content_2.text = "Second part of response" + mock_response.content = [mock_content_1, mock_content_2] + + mock_response.usage.input_tokens = 40 + mock_response.usage.output_tokens = 30 + + mock_claude_client.messages.create.return_value = mock_response + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-5-sonnet-20240620', + '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) + # Should take the first content block + assert result.text == "First part of response" + assert result.in_token == 40 + assert result.out_token == 30 + assert result.model == 'claude-3-5-sonnet-20240620' + + @patch('trustgraph.model.text_completion.claude.llm.anthropic.Anthropic') + @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') + @patch('trustgraph.base.llm_service.LlmService.__init__') + async def test_claude_client_initialization(self, mock_llm_init, mock_async_init, mock_anthropic_class): + """Test that Claude client is initialized correctly""" + # Arrange + mock_claude_client = MagicMock() + mock_anthropic_class.return_value = mock_claude_client + + mock_async_init.return_value = None + mock_llm_init.return_value = None + + config = { + 'model': 'claude-3-opus-20240229', + 'api_key': 'sk-ant-test-key', + 'temperature': 0.0, + 'max_output': 8192, + 'concurrency': 1, + 'taskgroup': AsyncMock(), + 'id': 'test-processor' + } + + # Act + processor = Processor(**config) + + # Assert + # Verify Anthropic client was called with correct API key + mock_anthropic_class.assert_called_once_with(api_key='sk-ant-test-key') + + # Verify processor has the client + assert processor.claude == mock_claude_client + assert processor.model == 'claude-3-opus-20240229' + + +if __name__ == '__main__': + pytest.main([__file__]) \ No newline at end of file