More tests

This commit is contained in:
Cyber MacGeddon 2025-09-25 22:11:35 +01:00
parent eee8de016c
commit 5cb51e282d
5 changed files with 1121 additions and 0 deletions

View file

@ -0,0 +1,276 @@
"""
Integration tests for Dynamic LLM Parameters
Testing end-to-end flow of runtime parameter changes in LLM processors
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from openai.types.chat import ChatCompletion, ChatCompletionMessage
from openai.types.chat.chat_completion import Choice
from openai.types.completion_usage import CompletionUsage
from trustgraph.model.text_completion.openai.llm import Processor as OpenAIProcessor
from trustgraph.base import LlmResult
@pytest.mark.integration
class TestDynamicLlmParameters:
"""Integration tests for dynamic parameter configuration"""
@pytest.fixture
def mock_openai_client(self):
"""Mock OpenAI client that returns realistic responses"""
client = MagicMock()
# Default mock response
usage = CompletionUsage(prompt_tokens=25, completion_tokens=15, total_tokens=40)
message = ChatCompletionMessage(role="assistant", content="Dynamic parameter test response")
choice = Choice(index=0, message=message, finish_reason="stop")
completion = ChatCompletion(
id="chatcmpl-test-dynamic",
choices=[choice],
created=1234567890,
model="gpt-4", # Will be overridden based on test
object="chat.completion",
usage=usage
)
client.chat.completions.create.return_value = completion
return client
@pytest.fixture
def base_processor_config(self):
"""Base configuration for test processors"""
return {
"api_key": "test-api-key",
"url": "https://api.openai.com/v1",
"temperature": 0.0, # Default temperature
"max_output": 1024,
}
@patch('trustgraph.model.text_completion.openai.llm.OpenAI')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_runtime_temperature_override(self, mock_llm_init, mock_async_init,
mock_openai_class, mock_openai_client, base_processor_config):
"""Test that temperature can be overridden at runtime"""
# Arrange
mock_openai_class.return_value = mock_openai_client
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = base_processor_config | {
"model": "gpt-3.5-turbo",
"concurrency": 1,
"taskgroup": AsyncMock(),
"id": "test-processor"
}
processor = OpenAIProcessor(**config)
# Act - Call with different temperature than configured default (0.0)
result = await processor.generate_content(
"System prompt",
"User prompt",
model=None, # Use default model
temperature=0.9 # Override temperature
)
# Assert
assert isinstance(result, LlmResult)
assert result.text == "Dynamic parameter test response"
# Verify the OpenAI API was called with the overridden temperature
mock_openai_client.chat.completions.create.assert_called_once()
call_args = mock_openai_client.chat.completions.create.call_args
assert call_args.kwargs['temperature'] == 0.9 # Should use runtime parameter
assert call_args.kwargs['model'] == "gpt-3.5-turbo" # Should use processor default
@patch('trustgraph.model.text_completion.openai.llm.OpenAI')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_runtime_model_override(self, mock_llm_init, mock_async_init,
mock_openai_class, mock_openai_client, base_processor_config):
"""Test that model can be overridden at runtime"""
# Arrange
mock_openai_class.return_value = mock_openai_client
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = base_processor_config | {
"model": "gpt-3.5-turbo", # Default model
"concurrency": 1,
"taskgroup": AsyncMock(),
"id": "test-processor"
}
processor = OpenAIProcessor(**config)
# Act - Call with different model than configured default
result = await processor.generate_content(
"System prompt",
"User prompt",
model="gpt-4", # Override model
temperature=None # Use default temperature
)
# Assert
assert isinstance(result, LlmResult)
# Verify the OpenAI API was called with the overridden model
mock_openai_client.chat.completions.create.assert_called_once()
call_args = mock_openai_client.chat.completions.create.call_args
assert call_args.kwargs['model'] == "gpt-4" # Should use runtime parameter
assert call_args.kwargs['temperature'] == 0.0 # Should use processor default
@patch('trustgraph.model.text_completion.openai.llm.OpenAI')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_both_parameters_override(self, mock_llm_init, mock_async_init,
mock_openai_class, mock_openai_client, base_processor_config):
"""Test that both model and temperature can be overridden simultaneously"""
# Arrange
mock_openai_class.return_value = mock_openai_client
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = base_processor_config | {
"model": "gpt-3.5-turbo", # Default model
"concurrency": 1,
"taskgroup": AsyncMock(),
"id": "test-processor"
}
processor = OpenAIProcessor(**config)
# Act - Override both parameters
result = await processor.generate_content(
"System prompt",
"User prompt",
model="gpt-4", # Override model
temperature=0.5 # Override temperature
)
# Assert
assert isinstance(result, LlmResult)
# Verify both parameters were overridden
mock_openai_client.chat.completions.create.assert_called_once()
call_args = mock_openai_client.chat.completions.create.call_args
assert call_args.kwargs['model'] == "gpt-4" # Should use runtime parameter
assert call_args.kwargs['temperature'] == 0.5 # Should use runtime parameter
@patch('trustgraph.model.text_completion.openai.llm.OpenAI')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_fallback_to_defaults_when_no_override(self, mock_llm_init, mock_async_init,
mock_openai_class, mock_openai_client, base_processor_config):
"""Test that processor falls back to configured defaults when no parameters are provided"""
# Arrange
mock_openai_class.return_value = mock_openai_client
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = base_processor_config | {
"model": "gpt-3.5-turbo", # Default model
"temperature": 0.2, # Default temperature
"concurrency": 1,
"taskgroup": AsyncMock(),
"id": "test-processor"
}
processor = OpenAIProcessor(**config)
# Act - Call with no parameter overrides
result = await processor.generate_content(
"System prompt",
"User prompt",
model=None, # Use default
temperature=None # Use default
)
# Assert
assert isinstance(result, LlmResult)
# Verify defaults were used
mock_openai_client.chat.completions.create.assert_called_once()
call_args = mock_openai_client.chat.completions.create.call_args
assert call_args.kwargs['model'] == "gpt-3.5-turbo" # Should use processor default
assert call_args.kwargs['temperature'] == 0.2 # Should use processor default
@patch('trustgraph.model.text_completion.openai.llm.OpenAI')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_multiple_concurrent_calls_different_parameters(self, mock_llm_init, mock_async_init,
mock_openai_class, mock_openai_client, base_processor_config):
"""Test multiple concurrent calls with different parameters don't interfere"""
# Arrange
mock_openai_class.return_value = mock_openai_client
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = base_processor_config | {
"model": "gpt-3.5-turbo",
"concurrency": 1,
"taskgroup": AsyncMock(),
"id": "test-processor"
}
processor = OpenAIProcessor(**config)
# Reset the mock to track multiple calls
mock_openai_client.reset_mock()
# Act - Make multiple calls with different parameters concurrently
import asyncio
tasks = [
processor.generate_content("System 1", "Prompt 1", model="gpt-3.5-turbo", temperature=0.1),
processor.generate_content("System 2", "Prompt 2", model="gpt-4", temperature=0.8),
processor.generate_content("System 3", "Prompt 3", model="gpt-3.5-turbo", temperature=0.5)
]
results = await asyncio.gather(*tasks)
# Assert
assert len(results) == 3
for result in results:
assert isinstance(result, LlmResult)
# Verify all calls were made with correct parameters
assert mock_openai_client.chat.completions.create.call_count == 3
# Get all call arguments
call_args_list = mock_openai_client.chat.completions.create.call_args_list
# Verify each call had the expected parameters
expected_params = [
("gpt-3.5-turbo", 0.1),
("gpt-4", 0.8),
("gpt-3.5-turbo", 0.5)
]
for i, (expected_model, expected_temp) in enumerate(expected_params):
call_kwargs = call_args_list[i].kwargs
assert call_kwargs['model'] == expected_model
assert call_kwargs['temperature'] == expected_temp
async def test_parameter_boundary_values(self, mock_openai_client, base_processor_config):
"""Test parameter boundary values (edge cases)"""
# This would test extreme values like temperature=0.0, temperature=2.0, etc.
# Implementation depends on specific validation requirements
pass
async def test_invalid_parameter_types_handling(self, mock_openai_client, base_processor_config):
"""Test handling of invalid parameter types"""
# This would test what happens with invalid temperature values, non-existent models, etc.
# Implementation depends on error handling requirements
pass
if __name__ == '__main__':
pytest.main([__file__])

View file

@ -0,0 +1,237 @@
"""
Unit tests for Flow Parameter Specification functionality
Testing parameter specification registration and handling in flow processors
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from unittest import IsolatedAsyncioTestCase
from trustgraph.base.flow_processor import FlowProcessor
from trustgraph.base import ParameterSpec, ConsumerSpec, ProducerSpec
class TestFlowParameterSpecs(IsolatedAsyncioTestCase):
"""Test flow processor parameter specification functionality"""
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_parameter_spec_registration(self, mock_async_init):
"""Test that parameter specs can be registered with flow processors"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-flow-processor',
'concurrency': 1
}
processor = FlowProcessor(**config)
# Create test parameter specs
model_spec = ParameterSpec(name="model")
temperature_spec = ParameterSpec(name="temperature")
# Act
processor.register_specification(model_spec)
processor.register_specification(temperature_spec)
# Assert
assert len(processor.specifications) >= 2
param_specs = [spec for spec in processor.specifications
if isinstance(spec, ParameterSpec)]
assert len(param_specs) >= 2
param_names = [spec.name for spec in param_specs]
assert "model" in param_names
assert "temperature" in param_names
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_mixed_specification_types(self, mock_async_init):
"""Test registration of mixed specification types (parameters, consumers, producers)"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-flow-processor',
'concurrency': 1
}
processor = FlowProcessor(**config)
# Create different spec types
param_spec = ParameterSpec(name="model")
consumer_spec = ConsumerSpec(name="input")
producer_spec = ProducerSpec(name="output")
# Act
processor.register_specification(param_spec)
processor.register_specification(consumer_spec)
processor.register_specification(producer_spec)
# Assert
assert len(processor.specifications) == 3
# Count each type
param_specs = [s for s in processor.specifications if isinstance(s, ParameterSpec)]
consumer_specs = [s for s in processor.specifications if isinstance(s, ConsumerSpec)]
producer_specs = [s for s in processor.specifications if isinstance(s, ProducerSpec)]
assert len(param_specs) == 1
assert len(consumer_specs) == 1
assert len(producer_specs) == 1
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_parameter_spec_metadata(self, mock_async_init):
"""Test parameter specification metadata handling"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-flow-processor',
'concurrency': 1
}
processor = FlowProcessor(**config)
# Create parameter specs with metadata (if supported)
model_spec = ParameterSpec(name="model")
temperature_spec = ParameterSpec(name="temperature")
# Act
processor.register_specification(model_spec)
processor.register_specification(temperature_spec)
# Assert
param_specs = [spec for spec in processor.specifications
if isinstance(spec, ParameterSpec)]
model_spec_registered = next((s for s in param_specs if s.name == "model"), None)
temperature_spec_registered = next((s for s in param_specs if s.name == "temperature"), None)
assert model_spec_registered is not None
assert temperature_spec_registered is not None
assert model_spec_registered.name == "model"
assert temperature_spec_registered.name == "temperature"
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_duplicate_parameter_spec_handling(self, mock_async_init):
"""Test handling of duplicate parameter spec registration"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-flow-processor',
'concurrency': 1
}
processor = FlowProcessor(**config)
# Create duplicate parameter specs
model_spec1 = ParameterSpec(name="model")
model_spec2 = ParameterSpec(name="model")
# Act
processor.register_specification(model_spec1)
processor.register_specification(model_spec2)
# Assert - Should allow duplicates (or handle appropriately)
param_specs = [spec for spec in processor.specifications
if isinstance(spec, ParameterSpec) and spec.name == "model"]
# Either should have 2 duplicates or the system should handle deduplication
assert len(param_specs) >= 1 # At least one should be registered
@patch('trustgraph.base.flow_processor.Flow')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
async def test_parameter_specs_available_to_flows(self, mock_async_init, mock_flow_class):
"""Test that parameter specs are available when flows are created"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-flow-processor',
'concurrency': 1
}
processor = FlowProcessor(**config)
processor.id = 'test-processor'
# Register parameter specs
model_spec = ParameterSpec(name="model")
temperature_spec = ParameterSpec(name="temperature")
processor.register_specification(model_spec)
processor.register_specification(temperature_spec)
mock_flow = AsyncMock()
mock_flow_class.return_value = mock_flow
flow_name = 'test-flow'
flow_defn = {'config': 'test-config'}
# Act
await processor.start_flow(flow_name, flow_defn)
# Assert - Flow should be created with access to processor specifications
mock_flow_class.assert_called_once_with('test-processor', flow_name, processor, flow_defn)
# The flow should have access to the processor's specifications
# (The exact mechanism depends on Flow implementation)
assert len(processor.specifications) >= 2
class TestParameterSpecValidation(IsolatedAsyncioTestCase):
"""Test parameter specification validation functionality"""
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_parameter_spec_name_validation(self, mock_async_init):
"""Test parameter spec name validation"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-flow-processor',
'concurrency': 1
}
processor = FlowProcessor(**config)
# Act & Assert - Valid parameter names
valid_specs = [
ParameterSpec(name="model"),
ParameterSpec(name="temperature"),
ParameterSpec(name="max_tokens"),
ParameterSpec(name="api_key")
]
for spec in valid_specs:
# Should not raise any exceptions
processor.register_specification(spec)
assert len([s for s in processor.specifications if isinstance(s, ParameterSpec)]) >= 4
def test_parameter_spec_creation_validation(self):
"""Test parameter spec creation with various inputs"""
# Test valid parameter spec creation
valid_specs = [
ParameterSpec(name="model"),
ParameterSpec(name="temperature"),
ParameterSpec(name="max_output"),
]
for spec in valid_specs:
assert spec.name is not None
assert isinstance(spec.name, str)
# Test edge cases (if parameter specs have validation)
# This depends on the actual ParameterSpec implementation
try:
empty_name_spec = ParameterSpec(name="")
# May or may not be valid depending on implementation
except Exception:
# If validation exists, it should catch invalid names
pass
if __name__ == '__main__':
pytest.main([__file__])

View file

@ -0,0 +1,248 @@
"""
Unit tests for LLM Service Parameter Specifications
Testing the new parameter-aware functionality added to the LLM base service
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from unittest import IsolatedAsyncioTestCase
from trustgraph.base.llm_service import LlmService, LlmResult
from trustgraph.base import ParameterSpec
class TestLlmServiceParameters(IsolatedAsyncioTestCase):
"""Test LLM service parameter specification functionality"""
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_parameter_specs_registration(self, mock_async_init):
"""Test that LLM service registers model and temperature parameter specs"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-llm-service',
'concurrency': 1
}
# Act
service = LlmService(**config)
# Assert
param_specs = {spec.name: spec for spec in service.specifications
if isinstance(spec, ParameterSpec)}
assert "model" in param_specs
assert "temperature" in param_specs
assert len(param_specs) >= 2 # May have other parameter specs
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_model_parameter_spec_properties(self, mock_async_init):
"""Test that model parameter spec has correct properties"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-llm-service',
'concurrency': 1
}
# Act
service = LlmService(**config)
# Assert
model_spec = None
for spec in service.specifications:
if isinstance(spec, ParameterSpec) and spec.name == "model":
model_spec = spec
break
assert model_spec is not None
assert model_spec.name == "model"
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
def test_temperature_parameter_spec_properties(self, mock_async_init):
"""Test that temperature parameter spec has correct properties"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-llm-service',
'concurrency': 1
}
# Act
service = LlmService(**config)
# Assert
temperature_spec = None
for spec in service.specifications:
if isinstance(spec, ParameterSpec) and spec.name == "temperature":
temperature_spec = spec
break
assert temperature_spec is not None
assert temperature_spec.name == "temperature"
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
async def test_on_request_extracts_parameters_from_flow(self, mock_async_init):
"""Test that on_request method extracts model and temperature from flow"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-llm-service',
'concurrency': 1
}
service = LlmService(**config)
# Mock the generate_content method to capture parameters
service.generate_content = AsyncMock(return_value=LlmResult(
text="test response",
in_token=10,
out_token=5,
model="gpt-4"
))
# Mock message and flow
mock_message = MagicMock()
mock_message.value.return_value = MagicMock()
mock_message.value.return_value.system = "system prompt"
mock_message.value.return_value.prompt = "user prompt"
mock_message.properties.return_value = {"id": "test-id"}
mock_consumer = MagicMock()
mock_consumer.name = "request"
mock_flow = MagicMock()
mock_flow.name = "test-flow"
mock_flow.return_value = "test-model" # flow("model") returns this
mock_flow.side_effect = lambda param: {
"model": "gpt-4",
"temperature": 0.7
}.get(param, f"mock-{param}")
mock_producer = AsyncMock()
mock_flow.return_value = mock_producer # flow("response") returns producer
# Act
await service.on_request(mock_message, mock_consumer, mock_flow)
# Assert
# Verify that generate_content was called with parameters from flow
service.generate_content.assert_called_once()
call_args = service.generate_content.call_args
assert call_args[0][0] == "system prompt" # system
assert call_args[0][1] == "user prompt" # prompt
assert call_args[0][2] == "gpt-4" # model
assert call_args[0][3] == 0.7 # temperature
# Verify flow was queried for both parameters
mock_flow.assert_any_call("model")
mock_flow.assert_any_call("temperature")
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
async def test_on_request_handles_missing_parameters_gracefully(self, mock_async_init):
"""Test that on_request handles missing parameters gracefully"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-llm-service',
'concurrency': 1
}
service = LlmService(**config)
# Mock the generate_content method
service.generate_content = AsyncMock(return_value=LlmResult(
text="test response",
in_token=10,
out_token=5,
model="default-model"
))
# Mock message and flow where flow returns None for parameters
mock_message = MagicMock()
mock_message.value.return_value = MagicMock()
mock_message.value.return_value.system = "system prompt"
mock_message.value.return_value.prompt = "user prompt"
mock_message.properties.return_value = {"id": "test-id"}
mock_consumer = MagicMock()
mock_consumer.name = "request"
mock_flow = MagicMock()
mock_flow.name = "test-flow"
mock_flow.return_value = None # Both parameters return None
mock_producer = AsyncMock()
mock_flow.return_value = mock_producer
# Act
await service.on_request(mock_message, mock_consumer, mock_flow)
# Assert
# Should still call generate_content, with None values that will use processor defaults
service.generate_content.assert_called_once()
call_args = service.generate_content.call_args
assert call_args[0][0] == "system prompt" # system
assert call_args[0][1] == "user prompt" # prompt
assert call_args[0][2] is None # model (will use processor default)
assert call_args[0][3] is None # temperature (will use processor default)
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
async def test_on_request_error_handling_preserves_behavior(self, mock_async_init):
"""Test that parameter extraction doesn't break existing error handling"""
# Arrange
mock_async_init.return_value = None
config = {
'id': 'test-llm-service',
'concurrency': 1
}
service = LlmService(**config)
# Mock generate_content to raise an exception
service.generate_content = AsyncMock(side_effect=Exception("Test error"))
# Mock message and flow
mock_message = MagicMock()
mock_message.value.return_value = MagicMock()
mock_message.value.return_value.system = "system prompt"
mock_message.value.return_value.prompt = "user prompt"
mock_message.properties.return_value = {"id": "test-id"}
mock_consumer = MagicMock()
mock_consumer.name = "request"
mock_flow = MagicMock()
mock_flow.name = "test-flow"
mock_flow.side_effect = lambda param: {
"model": "gpt-4",
"temperature": 0.7
}.get(param, f"mock-{param}")
mock_producer = AsyncMock()
mock_flow.producer = {"response": mock_producer}
# Act
await service.on_request(mock_message, mock_consumer, mock_flow)
# Assert
# Should have sent error response
mock_producer.send.assert_called_once()
error_response = mock_producer.send.call_args[0][0]
assert error_response.error is not None
assert error_response.error.type == "llm-error"
assert "Test error" in error_response.error.message
assert error_response.response is None
if __name__ == '__main__':
pytest.main([__file__])

View file

@ -0,0 +1,186 @@
"""
Unit tests for Parameter-Based Caching in LLM Processors
Testing processors that cache based on temperature parameters (Bedrock, GoogleAIStudio)
"""
import pytest
from unittest.mock import AsyncMock, MagicMock, patch
from unittest import IsolatedAsyncioTestCase
from trustgraph.model.text_completion.googleaistudio.llm import Processor as GoogleAIProcessor
from trustgraph.base import LlmResult
class TestParameterCaching(IsolatedAsyncioTestCase):
"""Test parameter-based caching functionality"""
@patch('trustgraph.model.text_completion.googleaistudio.llm.genai')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_googleai_temperature_cache_keys(self, mock_llm_init, mock_async_init, mock_genai):
"""Test that GoogleAI processor creates separate cache entries for different temperatures"""
# Arrange
mock_client = MagicMock()
mock_genai.Client.return_value = mock_client
mock_response = MagicMock()
mock_response.text = "Generated response"
mock_response.usage_metadata.prompt_token_count = 10
mock_response.usage_metadata.candidates_token_count = 5
mock_client.models.generate_content.return_value = mock_response
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, # Default temperature
'max_output': 1024,
'concurrency': 1,
'taskgroup': AsyncMock(),
'id': 'test-processor'
}
processor = GoogleAIProcessor(**config)
# Act - Call with different temperatures
await processor.generate_content("System", "Prompt 1", model="gemini-2.0-flash-001", temperature=0.0)
await processor.generate_content("System", "Prompt 2", model="gemini-2.0-flash-001", temperature=0.5)
await processor.generate_content("System", "Prompt 3", model="gemini-2.0-flash-001", temperature=1.0)
# Assert - Should have 3 different cache entries
cache_keys = list(processor.generation_configs.keys())
assert len(cache_keys) == 3
assert "gemini-2.0-flash-001:0.0" in cache_keys
assert "gemini-2.0-flash-001:0.5" in cache_keys
assert "gemini-2.0-flash-001:1.0" in cache_keys
# Verify each cached config has the correct temperature
assert processor.generation_configs["gemini-2.0-flash-001:0.0"].temperature == 0.0
assert processor.generation_configs["gemini-2.0-flash-001:0.5"].temperature == 0.5
assert processor.generation_configs["gemini-2.0-flash-001:1.0"].temperature == 1.0
@patch('trustgraph.model.text_completion.googleaistudio.llm.genai')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_googleai_cache_reuse_same_parameters(self, mock_llm_init, mock_async_init, mock_genai):
"""Test that GoogleAI processor reuses cache for identical model+temperature combinations"""
# Arrange
mock_client = MagicMock()
mock_genai.Client.return_value = mock_client
mock_response = MagicMock()
mock_response.text = "Generated response"
mock_response.usage_metadata.prompt_token_count = 10
mock_response.usage_metadata.candidates_token_count = 5
mock_client.models.generate_content.return_value = mock_response
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': 1024,
'concurrency': 1,
'taskgroup': AsyncMock(),
'id': 'test-processor'
}
processor = GoogleAIProcessor(**config)
# Act - Call multiple times with same parameters
await processor.generate_content("System", "Prompt 1", model="gemini-2.0-flash-001", temperature=0.7)
await processor.generate_content("System", "Prompt 2", model="gemini-2.0-flash-001", temperature=0.7)
await processor.generate_content("System", "Prompt 3", model="gemini-2.0-flash-001", temperature=0.7)
# Assert - Should have only 1 cache entry for the repeated parameters
cache_keys = list(processor.generation_configs.keys())
assert len(cache_keys) == 1
assert "gemini-2.0-flash-001:0.7" in cache_keys
# The same config object should be reused
config_obj = processor.generation_configs["gemini-2.0-flash-001:0.7"]
assert config_obj.temperature == 0.7
@patch('trustgraph.model.text_completion.googleaistudio.llm.genai')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_googleai_different_models_separate_caches(self, mock_llm_init, mock_async_init, mock_genai):
"""Test that different models create separate cache entries even with same temperature"""
# Arrange
mock_client = MagicMock()
mock_genai.Client.return_value = mock_client
mock_response = MagicMock()
mock_response.text = "Generated response"
mock_response.usage_metadata.prompt_token_count = 10
mock_response.usage_metadata.candidates_token_count = 5
mock_client.models.generate_content.return_value = mock_response
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': 1024,
'concurrency': 1,
'taskgroup': AsyncMock(),
'id': 'test-processor'
}
processor = GoogleAIProcessor(**config)
# Act - Call with different models, same temperature
await processor.generate_content("System", "Prompt 1", model="gemini-2.0-flash-001", temperature=0.5)
await processor.generate_content("System", "Prompt 2", model="gemini-1.5-flash-001", temperature=0.5)
# Assert - Should have separate cache entries for different models
cache_keys = list(processor.generation_configs.keys())
assert len(cache_keys) == 2
assert "gemini-2.0-flash-001:0.5" in cache_keys
assert "gemini-1.5-flash-001:0.5" in cache_keys
# Note: Bedrock tests would be similar but testing the Bedrock processor's caching behavior
# The Bedrock processor caches model variants with temperature in the cache key
async def test_bedrock_temperature_cache_keys(self):
"""Test Bedrock processor temperature-aware caching"""
# This would test the Bedrock processor's _get_or_create_variant method
# with different temperature values to ensure proper cache key generation
# Implementation would follow similar pattern to GoogleAI tests above
# but using the Bedrock processor and testing model_variants cache
pass
async def test_bedrock_cache_isolation_different_temperatures(self):
"""Test that Bedrock processor isolates cache entries by temperature"""
pass
async def test_cache_memory_efficiency(self):
"""Test that caches don't grow unbounded with many different parameter combinations"""
# This could test cache size limits or cleanup behavior if implemented
pass
class TestCachePerformance(IsolatedAsyncioTestCase):
"""Test caching performance characteristics"""
async def test_cache_hit_performance(self):
"""Test that cache hits are faster than cache misses"""
# This would measure timing differences between cache hits and misses
pass
async def test_concurrent_cache_access(self):
"""Test concurrent access to cached configurations"""
# This would test thread-safety of cache access
pass
if __name__ == '__main__':
pytest.main([__file__])

View file

@ -460,6 +460,180 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase):
assert processor.api_params["top_p"] == 1.0 assert processor.api_params["top_p"] == 1.0
assert processor.api_params["top_k"] == 32 assert processor.api_params["top_k"] == 32
@patch('trustgraph.model.text_completion.vertexai.llm.service_account')
@patch('trustgraph.model.text_completion.vertexai.llm.vertexai')
@patch('trustgraph.model.text_completion.vertexai.llm.GenerativeModel')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_generate_content_temperature_override(self, mock_llm_init, mock_async_init, mock_generative_model, mock_vertexai, mock_service_account):
"""Test temperature parameter override functionality"""
# Arrange
mock_credentials = MagicMock()
mock_service_account.Credentials.from_service_account_file.return_value = mock_credentials
mock_model = MagicMock()
mock_response = MagicMock()
mock_response.text = "Response with custom temperature"
mock_response.usage_metadata.prompt_token_count = 20
mock_response.usage_metadata.candidates_token_count = 12
mock_model.generate_content.return_value = mock_response
mock_generative_model.return_value = mock_model
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = {
'region': 'us-central1',
'model': 'gemini-2.0-flash-001',
'temperature': 0.0, # Default temperature
'max_output': 8192,
'private_key': 'private.json',
'concurrency': 1,
'taskgroup': AsyncMock(),
'id': 'test-processor'
}
processor = Processor(**config)
# Act - Override temperature at runtime
result = await processor.generate_content(
"System prompt",
"User prompt",
model=None, # Use default model
temperature=0.8 # Override temperature
)
# Assert
assert isinstance(result, LlmResult)
assert result.text == "Response with custom temperature"
# Verify Gemini API was called with overridden temperature
mock_model.generate_content.assert_called_once()
call_args = mock_model.generate_content.call_args
# Check that generation_config has the overridden temperature
generation_config = call_args.kwargs['generation_config']
assert generation_config.temperature == 0.8 # Should use runtime override
@patch('trustgraph.model.text_completion.vertexai.llm.service_account')
@patch('trustgraph.model.text_completion.vertexai.llm.vertexai')
@patch('trustgraph.model.text_completion.vertexai.llm.GenerativeModel')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_generate_content_model_override(self, mock_llm_init, mock_async_init, mock_generative_model, mock_vertexai, mock_service_account):
"""Test model parameter override functionality"""
# Arrange
mock_credentials = MagicMock()
mock_service_account.Credentials.from_service_account_file.return_value = mock_credentials
# Mock different models
mock_model_default = MagicMock()
mock_model_override = MagicMock()
mock_response = MagicMock()
mock_response.text = "Response with custom model"
mock_response.usage_metadata.prompt_token_count = 18
mock_response.usage_metadata.candidates_token_count = 14
mock_model_override.generate_content.return_value = mock_response
# GenerativeModel should return different models based on input
def model_factory(model_name):
if model_name == 'gemini-1.5-pro':
return mock_model_override
return mock_model_default
mock_generative_model.side_effect = model_factory
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = {
'region': 'us-central1',
'model': 'gemini-2.0-flash-001', # Default model
'temperature': 0.2, # Default temperature
'max_output': 8192,
'private_key': 'private.json',
'concurrency': 1,
'taskgroup': AsyncMock(),
'id': 'test-processor'
}
processor = Processor(**config)
# Act - Override model at runtime
result = await processor.generate_content(
"System prompt",
"User prompt",
model="gemini-1.5-pro", # Override model
temperature=None # Use default temperature
)
# Assert
assert isinstance(result, LlmResult)
assert result.text == "Response with custom model"
# Verify the overridden model was used
mock_model_override.generate_content.assert_called_once()
# Verify GenerativeModel was called with the override model
mock_generative_model.assert_called_with('gemini-1.5-pro')
@patch('trustgraph.model.text_completion.vertexai.llm.service_account')
@patch('trustgraph.model.text_completion.vertexai.llm.vertexai')
@patch('trustgraph.model.text_completion.vertexai.llm.GenerativeModel')
@patch('trustgraph.base.async_processor.AsyncProcessor.__init__')
@patch('trustgraph.base.llm_service.LlmService.__init__')
async def test_generate_content_both_parameters_override(self, mock_llm_init, mock_async_init, mock_generative_model, mock_vertexai, mock_service_account):
"""Test overriding both model and temperature parameters simultaneously"""
# Arrange
mock_credentials = MagicMock()
mock_service_account.Credentials.from_service_account_file.return_value = mock_credentials
mock_model = MagicMock()
mock_response = MagicMock()
mock_response.text = "Response with both overrides"
mock_response.usage_metadata.prompt_token_count = 22
mock_response.usage_metadata.candidates_token_count = 16
mock_model.generate_content.return_value = mock_response
mock_generative_model.return_value = mock_model
mock_async_init.return_value = None
mock_llm_init.return_value = None
config = {
'region': 'us-central1',
'model': 'gemini-2.0-flash-001', # Default model
'temperature': 0.0, # Default temperature
'max_output': 8192,
'private_key': 'private.json',
'concurrency': 1,
'taskgroup': AsyncMock(),
'id': 'test-processor'
}
processor = Processor(**config)
# Act - Override both parameters at runtime
result = await processor.generate_content(
"System prompt",
"User prompt",
model="gemini-1.5-flash-001", # Override model
temperature=0.9 # Override temperature
)
# Assert
assert isinstance(result, LlmResult)
assert result.text == "Response with both overrides"
# Verify both overrides were used
mock_model.generate_content.assert_called_once()
call_args = mock_model.generate_content.call_args
# Verify model override
mock_generative_model.assert_called_with('gemini-1.5-flash-001') # Should use runtime override
# Verify temperature override
generation_config = call_args.kwargs['generation_config']
assert generation_config.temperature == 0.9 # Should use runtime override
if __name__ == '__main__': if __name__ == '__main__':
pytest.main([__file__]) pytest.main([__file__])