diff --git a/tests/integration/test_text_completion_integration.py b/tests/integration/test_text_completion_integration.py index 1a8e5e1b..5ff6ee18 100644 --- a/tests/integration/test_text_completion_integration.py +++ b/tests/integration/test_text_completion_integration.py @@ -60,13 +60,13 @@ class TestTextCompletionIntegration: """Create text completion processor with test configuration""" # Create a minimal processor instance for testing generate_content processor = MagicMock() - processor.model = processor_config["model"] + processor.default_model = processor_config["model"] processor.temperature = processor_config["temperature"] processor.max_output = processor_config["max_output"] - + # Add the actual generate_content method from Processor class processor.generate_content = Processor.generate_content.__get__(processor, Processor) - + return processor @pytest.mark.asyncio @@ -112,11 +112,11 @@ class TestTextCompletionIntegration: for config in test_configs: # Arrange - Create minimal processor mock processor = MagicMock() - processor.model = config['model'] + processor.default_model = config['model'] processor.temperature = config['temperature'] processor.max_output = config['max_output'] processor.openai = mock_openai_client - + # Add the actual generate_content method processor.generate_content = Processor.generate_content.__get__(processor, Processor) @@ -242,7 +242,7 @@ class TestTextCompletionIntegration: processors = [] for i in range(5): processor = MagicMock() - processor.model = processor_config["model"] + processor.default_model = processor_config["model"] processor.temperature = processor_config["temperature"] processor.max_output = processor_config["max_output"] processor.openai = mock_openai_client @@ -348,7 +348,7 @@ class TestTextCompletionIntegration: """Test that model parameters are correctly passed to OpenAI API""" # Arrange processor = MagicMock() - processor.model = "gpt-4" + processor.default_model = "gpt-4" processor.temperature = 0.8 processor.max_output = 2048 processor.openai = mock_openai_client diff --git a/tests/unit/test_text_completion/test_azure_openai_processor.py b/tests/unit/test_text_completion/test_azure_openai_processor.py index b5669907..967a3893 100644 --- a/tests/unit/test_text_completion/test_azure_openai_processor.py +++ b/tests/unit/test_text_completion/test_azure_openai_processor.py @@ -44,7 +44,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-4' + assert processor.default_model == 'gpt-4' assert processor.temperature == 0.0 assert processor.max_output == 4192 assert hasattr(processor, 'openai') @@ -254,7 +254,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-35-turbo' + assert processor.default_model == 'gpt-35-turbo' assert processor.temperature == 0.7 assert processor.max_output == 2048 mock_azure_openai_class.assert_called_once_with( @@ -289,7 +289,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-4' + assert processor.default_model == 'gpt-4' assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4192 # default_max_output mock_azure_openai_class.assert_called_once_with( diff --git a/tests/unit/test_text_completion/test_azure_processor.py b/tests/unit/test_text_completion/test_azure_processor.py index 6ef78a2c..a1e2ba75 100644 --- a/tests/unit/test_text_completion/test_azure_processor.py +++ b/tests/unit/test_text_completion/test_azure_processor.py @@ -43,7 +43,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): assert processor.token == 'test-token' assert processor.temperature == 0.0 assert processor.max_output == 4192 - assert processor.model == 'AzureAI' + assert processor.default_model == 'AzureAI' @patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @@ -261,7 +261,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): assert processor.token == 'custom-token' assert processor.temperature == 0.7 assert processor.max_output == 2048 - assert processor.model == 'AzureAI' + assert processor.default_model == 'AzureAI' @patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @@ -289,7 +289,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase): assert processor.token == 'test-token' assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4192 # default_max_output - assert processor.model == 'AzureAI' # default_model + assert processor.default_model == 'AzureAI' # default_model @patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_claude_processor.py b/tests/unit/test_text_completion/test_claude_processor.py index 27a18b93..12951019 100644 --- a/tests/unit/test_text_completion/test_claude_processor.py +++ b/tests/unit/test_text_completion/test_claude_processor.py @@ -42,7 +42,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-5-sonnet-20240620' + assert processor.default_model == 'claude-3-5-sonnet-20240620' assert processor.temperature == 0.0 assert processor.max_output == 8192 assert hasattr(processor, 'claude') @@ -217,7 +217,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-haiku-20240307' + assert processor.default_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') @@ -246,7 +246,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-5-sonnet-20240620' # default_model + assert processor.default_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') @@ -433,7 +433,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.claude == mock_claude_client - assert processor.model == 'claude-3-opus-20240229' + assert processor.default_model == 'claude-3-opus-20240229' if __name__ == '__main__': diff --git a/tests/unit/test_text_completion/test_cohere_processor.py b/tests/unit/test_text_completion/test_cohere_processor.py index ebb6b626..9e4397bc 100644 --- a/tests/unit/test_text_completion/test_cohere_processor.py +++ b/tests/unit/test_text_completion/test_cohere_processor.py @@ -41,7 +41,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'c4ai-aya-23-8b' + assert processor.default_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') @@ -201,7 +201,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'command-light' + assert processor.default_model == 'command-light' assert processor.temperature == 0.7 mock_cohere_class.assert_called_once_with(api_key='custom-api-key') @@ -229,7 +229,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'c4ai-aya-23-8b' # default_model + assert processor.default_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') @@ -395,7 +395,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.cohere == mock_cohere_client - assert processor.model == 'command-r' + assert processor.default_model == 'command-r' @patch('trustgraph.model.text_completion.cohere.llm.cohere.Client') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_googleaistudio_processor.py b/tests/unit/test_text_completion/test_googleaistudio_processor.py index a3ca0057..f31715d2 100644 --- a/tests/unit/test_text_completion/test_googleaistudio_processor.py +++ b/tests/unit/test_text_completion/test_googleaistudio_processor.py @@ -42,7 +42,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' + assert processor.default_model == 'gemini-2.0-flash-001' assert processor.temperature == 0.0 assert processor.max_output == 8192 assert hasattr(processor, 'client') @@ -205,7 +205,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-1.5-pro' + assert processor.default_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') @@ -234,7 +234,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' # default_model + assert processor.default_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') @@ -431,7 +431,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase): # Verify processor has the client assert processor.client == mock_genai_client - assert processor.model == 'gemini-1.5-flash' + assert processor.default_model == 'gemini-1.5-flash' @patch('trustgraph.model.text_completion.googleaistudio.llm.genai.Client') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_llamafile_processor.py b/tests/unit/test_text_completion/test_llamafile_processor.py index bae1a4bb..425cfce8 100644 --- a/tests/unit/test_text_completion/test_llamafile_processor.py +++ b/tests/unit/test_text_completion/test_llamafile_processor.py @@ -42,7 +42,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'LLaMA_CPP' + assert processor.default_model == 'LLaMA_CPP' assert processor.llamafile == 'http://localhost:8080/v1' assert processor.temperature == 0.0 assert processor.max_output == 4096 @@ -91,7 +91,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): assert result.text == "Generated response from LlamaFile" assert result.in_token == 20 assert result.out_token == 12 - assert result.model == 'llama.cpp' # Note: model in result is hardcoded to 'llama.cpp' + assert result.model == 'LLaMA_CPP' # Uses the default model name # Verify the OpenAI API call structure mock_openai_client.chat.completions.create.assert_called_once_with( @@ -99,7 +99,15 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): messages=[{ "role": "user", "content": "System prompt\n\nUser prompt" - }] + }], + temperature=0.0, + max_tokens=4096, + top_p=1, + frequency_penalty=0, + presence_penalty=0, + response_format={ + "type": "text" + } ) @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') @@ -157,7 +165,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'custom-llama' + assert processor.default_model == 'custom-llama' assert processor.llamafile == 'http://custom-host:8080/v1' assert processor.temperature == 0.7 assert processor.max_output == 2048 @@ -189,7 +197,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'LLaMA_CPP' # default_model + assert processor.default_model == 'LLaMA_CPP' # default_model assert processor.llamafile == 'http://localhost:8080/v1' # default_llamafile assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4096 # default_max_output @@ -237,7 +245,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): assert result.text == "Default response" assert result.in_token == 2 assert result.out_token == 3 - assert result.model == 'llama.cpp' + assert result.model == 'LLaMA_CPP' # Verify the combined prompt is sent correctly call_args = mock_openai_client.chat.completions.create.call_args @@ -408,8 +416,8 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase): result = await processor.generate_content("System", "User") # Assert - assert result.model == 'llama.cpp' # Should always be 'llama.cpp', not 'custom-model-name' - assert processor.model == 'custom-model-name' # But processor.model should still be custom + assert result.model == 'custom-model-name' # Uses the actual model name passed to generate_content + assert processor.default_model == 'custom-model-name' # But processor.model should still be custom @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') diff --git a/tests/unit/test_text_completion/test_ollama_processor.py b/tests/unit/test_text_completion/test_ollama_processor.py index e846ec12..138a8598 100644 --- a/tests/unit/test_text_completion/test_ollama_processor.py +++ b/tests/unit/test_text_completion/test_ollama_processor.py @@ -40,7 +40,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'llama2' + assert processor.default_model == 'llama2' assert hasattr(processor, 'llm') mock_client_class.assert_called_once_with(host='http://localhost:11434') @@ -81,7 +81,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert result.in_token == 15 assert result.out_token == 8 assert result.model == 'llama2' - mock_client.generate.assert_called_once_with('llama2', "System prompt\n\nUser prompt") + mock_client.generate.assert_called_once_with('llama2', "System prompt\n\nUser prompt", options={'temperature': 0.0}) @patch('trustgraph.model.text_completion.ollama.llm.Client') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @@ -134,7 +134,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'mistral' + assert processor.default_model == 'mistral' mock_client_class.assert_called_once_with(host='http://192.168.1.100:11434') @patch('trustgraph.model.text_completion.ollama.llm.Client') @@ -160,7 +160,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemma2:9b' # default_model + assert processor.default_model == 'gemma2:9b' # default_model # Should use default_ollama (http://localhost:11434 or from OLLAMA_HOST env) mock_client_class.assert_called_once() @@ -203,7 +203,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert result.model == 'llama2' # The prompt should be "" + "\n\n" + "" = "\n\n" - mock_client.generate.assert_called_once_with('llama2', "\n\n") + mock_client.generate.assert_called_once_with('llama2', "\n\n", options={'temperature': 0.0}) @patch('trustgraph.model.text_completion.ollama.llm.Client') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @@ -310,7 +310,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase): assert result.out_token == 15 # Verify the combined prompt - mock_client.generate.assert_called_once_with('llama2', "You are a helpful assistant\n\nWhat is AI?") + mock_client.generate.assert_called_once_with('llama2', "You are a helpful assistant\n\nWhat is AI?", options={'temperature': 0.0}) if __name__ == '__main__': diff --git a/tests/unit/test_text_completion/test_openai_processor.py b/tests/unit/test_text_completion/test_openai_processor.py index 504dad50..56425285 100644 --- a/tests/unit/test_text_completion/test_openai_processor.py +++ b/tests/unit/test_text_completion/test_openai_processor.py @@ -43,7 +43,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-3.5-turbo' + assert processor.default_model == 'gpt-3.5-turbo' assert processor.temperature == 0.0 assert processor.max_output == 4096 assert hasattr(processor, 'openai') @@ -222,7 +222,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-4' + assert processor.default_model == 'gpt-4' assert processor.temperature == 0.7 assert processor.max_output == 2048 mock_openai_class.assert_called_once_with(base_url='https://custom-openai-url.com/v1', api_key='custom-api-key') @@ -251,7 +251,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gpt-3.5-turbo' # default_model + assert processor.default_model == 'gpt-3.5-turbo' # default_model assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 4096 # default_max_output mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key') diff --git a/tests/unit/test_text_completion/test_vertexai_processor.py b/tests/unit/test_text_completion/test_vertexai_processor.py index 3910a30c..48c70d15 100644 --- a/tests/unit/test_text_completion/test_vertexai_processor.py +++ b/tests/unit/test_text_completion/test_vertexai_processor.py @@ -47,10 +47,10 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' # It's stored as 'model', not 'model_name' - assert hasattr(processor, 'generation_config') + assert processor.default_model == 'gemini-2.0-flash-001' # It's stored as 'model', not 'model_name' + assert hasattr(processor, 'generation_configs') # Now a cache dictionary assert hasattr(processor, 'safety_settings') - assert hasattr(processor, 'llm') + assert hasattr(processor, 'model_clients') # LLM clients are now cached mock_service_account.Credentials.from_service_account_file.assert_called_once_with('private.json') mock_vertexai.init.assert_called_once() @@ -102,7 +102,8 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): mock_model.generate_content.assert_called_once() # Verify the call was made with the expected parameters call_args = mock_model.generate_content.call_args - assert call_args[1]['generation_config'] == processor.generation_config + # Generation config is now created dynamically per model + assert 'generation_config' in call_args[1] assert call_args[1]['safety_settings'] == processor.safety_settings @patch('trustgraph.model.text_completion.vertexai.llm.service_account') @@ -223,7 +224,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-2.0-flash-001' + assert processor.default_model == 'gemini-2.0-flash-001' mock_auth_default.assert_called_once() mock_vertexai.init.assert_called_once_with( location='us-central1', @@ -296,11 +297,11 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'gemini-1.5-pro' + assert processor.default_model == 'gemini-1.5-pro' # Verify that generation_config object exists (can't easily check internal values) - assert hasattr(processor, 'generation_config') - assert processor.generation_config is not None + assert hasattr(processor, 'generation_configs') # Now a cache dictionary + assert processor.generation_configs == {} # Empty cache initially # Verify that safety settings are configured assert len(processor.safety_settings) == 4 @@ -353,8 +354,8 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): project='test-project-123' ) - # Verify GenerativeModel was created with the right model name - mock_generative_model.assert_called_once_with('gemini-2.0-flash-001') + # GenerativeModel is now created lazily on first use, not at initialization + mock_generative_model.assert_not_called() @patch('trustgraph.model.text_completion.vertexai.llm.service_account') @patch('trustgraph.model.text_completion.vertexai.llm.vertexai') @@ -440,8 +441,8 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'claude-3-sonnet@20240229' - assert processor.is_anthropic == True + assert processor.default_model == 'claude-3-sonnet@20240229' + # is_anthropic logic is now determined dynamically per request # Verify service account was called with private key mock_service_account.Credentials.from_service_account_file.assert_called_once_with('anthropic-key.json') diff --git a/tests/unit/test_text_completion/test_vllm_processor.py b/tests/unit/test_text_completion/test_vllm_processor.py index 7d30cf74..7124c229 100644 --- a/tests/unit/test_text_completion/test_vllm_processor.py +++ b/tests/unit/test_text_completion/test_vllm_processor.py @@ -42,7 +42,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'TheBloke/Mistral-7B-v0.1-AWQ' + assert processor.default_model == 'TheBloke/Mistral-7B-v0.1-AWQ' assert processor.base_url == 'http://vllm-service:8899/v1' assert processor.temperature == 0.0 assert processor.max_output == 2048 @@ -199,7 +199,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'custom-model' + assert processor.default_model == 'custom-model' assert processor.base_url == 'http://custom-vllm:8080/v1' assert processor.temperature == 0.7 assert processor.max_output == 1024 @@ -228,7 +228,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase): processor = Processor(**config) # Assert - assert processor.model == 'TheBloke/Mistral-7B-v0.1-AWQ' # default_model + assert processor.default_model == 'TheBloke/Mistral-7B-v0.1-AWQ' # default_model assert processor.base_url == 'http://vllm-service:8899/v1' # default_base_url assert processor.temperature == 0.0 # default_temperature assert processor.max_output == 2048 # default_max_output