Tests are failing (#534)

* Fix tests, update to new model parameter usage
This commit is contained in:
cybermaggedon 2025-09-25 21:32:19 +01:00 committed by GitHub
parent 6f4f7ce6b4
commit b0a3716b0e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 66 additions and 57 deletions

View file

@ -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(

View file

@ -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__')

View file

@ -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__':

View file

@ -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__')

View file

@ -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__')

View file

@ -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__')

View file

@ -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__':

View file

@ -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')

View file

@ -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')

View file

@ -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