mirror of
https://github.com/trustgraph-ai/trustgraph.git
synced 2026-07-24 20:51:02 +02:00
More tests
This commit is contained in:
parent
eee8de016c
commit
5cb51e282d
5 changed files with 1121 additions and 0 deletions
276
tests/integration/test_dynamic_llm_parameters.py
Normal file
276
tests/integration/test_dynamic_llm_parameters.py
Normal 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__])
|
||||
237
tests/unit/test_base/test_flow_parameter_specs.py
Normal file
237
tests/unit/test_base/test_flow_parameter_specs.py
Normal 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__])
|
||||
248
tests/unit/test_base/test_llm_service_parameters.py
Normal file
248
tests/unit/test_base/test_llm_service_parameters.py
Normal 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__])
|
||||
186
tests/unit/test_text_completion/test_parameter_caching.py
Normal file
186
tests/unit/test_text_completion/test_parameter_caching.py
Normal 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__])
|
||||
|
|
@ -460,6 +460,180 @@ class TestVertexAIProcessorSimple(IsolatedAsyncioTestCase):
|
|||
assert processor.api_params["top_p"] == 1.0
|
||||
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__':
|
||||
pytest.main([__file__])
|
||||
Loading…
Add table
Add a link
Reference in a new issue