apunkt/trustgraph
1
0
Fork 0
mirror of https://github.com/trustgraph-ai/trustgraph.git synced 2026-07-24 12:41:02 +02:00
Code Issues Projects Releases Packages Wiki Activity Actions 2

Fix some tests, update to new model parameter usage

Browse source
This commit is contained in:
Cyber MacGeddon 2025-09-25 19:18:39 +01:00
parent 6f4f7ce6b4
commit c169d8f0a9
10 changed files with 38 additions and 38 deletions
Download patch file Download diff file Expand all files Collapse all files

6
tests/unit/test_text_completion/test_azure_openai_processor.py
View file

@ -44,7 +44,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gpt-4' assert processor.default_model == 'gpt-4'
assert processor.temperature == 0.0 assert processor.temperature == 0.0
assert processor.max_output == 4192 assert processor.max_output == 4192
assert hasattr(processor, 'openai') assert hasattr(processor, 'openai')
@ -254,7 +254,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gpt-35-turbo' assert processor.default_model == 'gpt-35-turbo'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 2048 assert processor.max_output == 2048
mock_azure_openai_class.assert_called_once_with( mock_azure_openai_class.assert_called_once_with(
@ -289,7 +289,7 @@ class TestAzureOpenAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gpt-4' assert processor.default_model == 'gpt-4'
assert processor.temperature == 0.0 # default_temperature assert processor.temperature == 0.0 # default_temperature
assert processor.max_output == 4192 # default_max_output assert processor.max_output == 4192 # default_max_output
mock_azure_openai_class.assert_called_once_with( mock_azure_openai_class.assert_called_once_with(

6
tests/unit/test_text_completion/test_azure_processor.py
View file

@ -43,7 +43,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase):
assert processor.token == 'test-token' assert processor.token == 'test-token'
assert processor.temperature == 0.0 assert processor.temperature == 0.0
assert processor.max_output == 4192 assert processor.max_output == 4192
assert processor.model == 'AzureAI' assert processor.default_model == 'AzureAI'
@patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.model.text_completion.azure.llm.requests')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@ -261,7 +261,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase):
assert processor.token == 'custom-token' assert processor.token == 'custom-token'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 2048 assert processor.max_output == 2048
assert processor.model == 'AzureAI' assert processor.default_model == 'AzureAI'
@patch('trustgraph.model.text_completion.azure.llm.requests') @patch('trustgraph.model.text_completion.azure.llm.requests')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@ -289,7 +289,7 @@ class TestAzureProcessorSimple(IsolatedAsyncioTestCase):
assert processor.token == 'test-token' assert processor.token == 'test-token'
assert processor.temperature == 0.0 # default_temperature assert processor.temperature == 0.0 # default_temperature
assert processor.max_output == 4192 # default_max_output 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.model.text_completion.azure.llm.requests')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__')

8
tests/unit/test_text_completion/test_claude_processor.py
View file

@ -42,7 +42,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.temperature == 0.0
assert processor.max_output == 8192 assert processor.max_output == 8192
assert hasattr(processor, 'claude') assert hasattr(processor, 'claude')
@ -217,7 +217,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'claude-3-haiku-20240307' assert processor.default_model == 'claude-3-haiku-20240307'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 4096 assert processor.max_output == 4096
mock_anthropic_class.assert_called_once_with(api_key='custom-api-key') mock_anthropic_class.assert_called_once_with(api_key='custom-api-key')
@ -246,7 +246,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.temperature == 0.0 # default_temperature
assert processor.max_output == 8192 # default_max_output assert processor.max_output == 8192 # default_max_output
mock_anthropic_class.assert_called_once_with(api_key='test-api-key') mock_anthropic_class.assert_called_once_with(api_key='test-api-key')
@ -433,7 +433,7 @@ class TestClaudeProcessorSimple(IsolatedAsyncioTestCase):
# Verify processor has the client # Verify processor has the client
assert processor.claude == mock_claude_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__': if __name__ == '__main__':

8
tests/unit/test_text_completion/test_cohere_processor.py
View file

@ -41,7 +41,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'c4ai-aya-23-8b' assert processor.default_model == 'c4ai-aya-23-8b'
assert processor.temperature == 0.0 assert processor.temperature == 0.0
assert hasattr(processor, 'cohere') assert hasattr(processor, 'cohere')
mock_cohere_class.assert_called_once_with(api_key='test-api-key') mock_cohere_class.assert_called_once_with(api_key='test-api-key')
@ -201,7 +201,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'command-light' assert processor.default_model == 'command-light'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
mock_cohere_class.assert_called_once_with(api_key='custom-api-key') mock_cohere_class.assert_called_once_with(api_key='custom-api-key')
@ -229,7 +229,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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 assert processor.temperature == 0.0 # default_temperature
mock_cohere_class.assert_called_once_with(api_key='test-api-key') mock_cohere_class.assert_called_once_with(api_key='test-api-key')
@ -395,7 +395,7 @@ class TestCohereProcessorSimple(IsolatedAsyncioTestCase):
# Verify processor has the client # Verify processor has the client
assert processor.cohere == mock_cohere_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.model.text_completion.cohere.llm.cohere.Client')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__')

8
tests/unit/test_text_completion/test_googleaistudio_processor.py
View file

@ -42,7 +42,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.temperature == 0.0
assert processor.max_output == 8192 assert processor.max_output == 8192
assert hasattr(processor, 'client') assert hasattr(processor, 'client')
@ -205,7 +205,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gemini-1.5-pro' assert processor.default_model == 'gemini-1.5-pro'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 4096 assert processor.max_output == 4096
mock_genai_class.assert_called_once_with(api_key='custom-api-key') mock_genai_class.assert_called_once_with(api_key='custom-api-key')
@ -234,7 +234,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.temperature == 0.0 # default_temperature
assert processor.max_output == 8192 # default_max_output assert processor.max_output == 8192 # default_max_output
mock_genai_class.assert_called_once_with(api_key='test-api-key') mock_genai_class.assert_called_once_with(api_key='test-api-key')
@ -431,7 +431,7 @@ class TestGoogleAIStudioProcessorSimple(IsolatedAsyncioTestCase):
# Verify processor has the client # Verify processor has the client
assert processor.client == mock_genai_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.model.text_completion.googleaistudio.llm.genai.Client')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__')

14
tests/unit/test_text_completion/test_llamafile_processor.py
View file

@ -42,7 +42,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'LLaMA_CPP' assert processor.default_model == 'LLaMA_CPP'
assert processor.llamafile == 'http://localhost:8080/v1' assert processor.llamafile == 'http://localhost:8080/v1'
assert processor.temperature == 0.0 assert processor.temperature == 0.0
assert processor.max_output == 4096 assert processor.max_output == 4096
@ -91,7 +91,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase):
assert result.text == "Generated response from LlamaFile" assert result.text == "Generated response from LlamaFile"
assert result.in_token == 20 assert result.in_token == 20
assert result.out_token == 12 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 # Verify the OpenAI API call structure
mock_openai_client.chat.completions.create.assert_called_once_with( mock_openai_client.chat.completions.create.assert_called_once_with(
@ -157,7 +157,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'custom-llama' assert processor.default_model == 'custom-llama'
assert processor.llamafile == 'http://custom-host:8080/v1' assert processor.llamafile == 'http://custom-host:8080/v1'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 2048 assert processor.max_output == 2048
@ -189,7 +189,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.llamafile == 'http://localhost:8080/v1' # default_llamafile
assert processor.temperature == 0.0 # default_temperature assert processor.temperature == 0.0 # default_temperature
assert processor.max_output == 4096 # default_max_output assert processor.max_output == 4096 # default_max_output
@ -237,7 +237,7 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase):
assert result.text == "Default response" assert result.text == "Default response"
assert result.in_token == 2 assert result.in_token == 2
assert result.out_token == 3 assert result.out_token == 3
assert result.model == 'llama.cpp' assert result.model == 'LLaMA_CPP'
# Verify the combined prompt is sent correctly # Verify the combined prompt is sent correctly
call_args = mock_openai_client.chat.completions.create.call_args call_args = mock_openai_client.chat.completions.create.call_args
@ -408,8 +408,8 @@ class TestLlamaFileProcessorSimple(IsolatedAsyncioTestCase):
result = await processor.generate_content("System", "User") result = await processor.generate_content("System", "User")
# Assert # Assert
assert result.model == 'llama.cpp' # Should always be 'llama.cpp', not 'custom-model-name' assert result.model == 'custom-model-name' # Uses the actual model name passed to generate_content
assert processor.model == 'custom-model-name' # But processor.model should still be custom assert processor.default_model == 'custom-model-name' # But processor.model should still be custom
@patch('trustgraph.model.text_completion.llamafile.llm.OpenAI') @patch('trustgraph.model.text_completion.llamafile.llm.OpenAI')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__') @patch('trustgraph.base.async_processor.AsyncProcessor.__init__')

6
tests/unit/test_text_completion/test_ollama_processor.py
View file

@ -40,7 +40,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'llama2' assert processor.default_model == 'llama2'
assert hasattr(processor, 'llm') assert hasattr(processor, 'llm')
mock_client_class.assert_called_once_with(host='http://localhost:11434') mock_client_class.assert_called_once_with(host='http://localhost:11434')
@ -134,7 +134,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'mistral' assert processor.default_model == 'mistral'
mock_client_class.assert_called_once_with(host='http://192.168.1.100:11434') mock_client_class.assert_called_once_with(host='http://192.168.1.100:11434')
@patch('trustgraph.model.text_completion.ollama.llm.Client') @patch('trustgraph.model.text_completion.ollama.llm.Client')
@ -160,7 +160,7 @@ class TestOllamaProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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) # Should use default_ollama (http://localhost:11434 or from OLLAMA_HOST env)
mock_client_class.assert_called_once() mock_client_class.assert_called_once()

6
tests/unit/test_text_completion/test_openai_processor.py
View file

@ -43,7 +43,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gpt-3.5-turbo' assert processor.default_model == 'gpt-3.5-turbo'
assert processor.temperature == 0.0 assert processor.temperature == 0.0
assert processor.max_output == 4096 assert processor.max_output == 4096
assert hasattr(processor, 'openai') assert hasattr(processor, 'openai')
@ -222,7 +222,7 @@ class TestOpenAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gpt-4' assert processor.default_model == 'gpt-4'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 2048 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') 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) processor = Processor(**config)
# Assert # 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.temperature == 0.0 # default_temperature
assert processor.max_output == 4096 # default_max_output 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') mock_openai_class.assert_called_once_with(base_url='https://api.openai.com/v1', api_key='test-api-key')

8
tests/unit/test_text_completion/test_vertexai_processor.py
View file

@ -47,7 +47,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'gemini-2.0-flash-001' # It's stored as 'model', not 'model_name' assert processor.default_model == 'gemini-2.0-flash-001' # It's stored as 'model', not 'model_name'
assert hasattr(processor, 'generation_config') assert hasattr(processor, 'generation_config')
assert hasattr(processor, 'safety_settings') assert hasattr(processor, 'safety_settings')
assert hasattr(processor, 'llm') assert hasattr(processor, 'llm')
@ -223,7 +223,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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_auth_default.assert_called_once()
mock_vertexai.init.assert_called_once_with( mock_vertexai.init.assert_called_once_with(
location='us-central1', location='us-central1',
@ -296,7 +296,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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) # Verify that generation_config object exists (can't easily check internal values)
assert hasattr(processor, 'generation_config') assert hasattr(processor, 'generation_config')
@ -440,7 +440,7 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'claude-3-sonnet@20240229' assert processor.default_model == 'claude-3-sonnet@20240229'
assert processor.is_anthropic == True assert processor.is_anthropic == True
# Verify service account was called with private key # Verify service account was called with private key

6
tests/unit/test_text_completion/test_vllm_processor.py
View file

@ -42,7 +42,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.base_url == 'http://vllm-service:8899/v1'
assert processor.temperature == 0.0 assert processor.temperature == 0.0
assert processor.max_output == 2048 assert processor.max_output == 2048
@ -199,7 +199,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # Assert
assert processor.model == 'custom-model' assert processor.default_model == 'custom-model'
assert processor.base_url == 'http://custom-vllm:8080/v1' assert processor.base_url == 'http://custom-vllm:8080/v1'
assert processor.temperature == 0.7 assert processor.temperature == 0.7
assert processor.max_output == 1024 assert processor.max_output == 1024
@ -228,7 +228,7 @@ class TestVLLMProcessorSimple(IsolatedAsyncioTestCase):
processor = Processor(**config) processor = Processor(**config)
# Assert # 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.base_url == 'http://vllm-service:8899/v1' # default_base_url
assert processor.temperature == 0.0 # default_temperature assert processor.temperature == 0.0 # default_temperature
assert processor.max_output == 2048 # default_max_output assert processor.max_output == 2048 # default_max_output