For diversity in testgen

This commit is contained in:
Sarthak Agarwal 2025-12-29 22:05:09 +05:30
parent 238ed71576
commit af7f00cdae

View file

@ -17,7 +17,7 @@ from openai import OpenAIError
from aiservice.analytics.posthog import ph
from aiservice.common_utils import parse_python_version, safe_isort, should_hack_for_demo, validate_trace_id
from aiservice.env_specific import IS_PRODUCTION, debug_log_sensitive_data
from aiservice.llm import EXECUTE_MODEL, LLMResponse, calculate_llm_cost, call_llm
from aiservice.llm import EXECUTE_MODEL, LLMResponse, OpenAI_GPT_4_1, Anthropic_Claude_Sonnet_4_5, calculate_llm_cost, call_llm, llm_clients
from aiservice.observability.decorators import observe_llm_call
from log_features.log_event import update_optimization_cost
from log_features.log_features import log_features
@ -417,6 +417,22 @@ async def testgen(
try:
debug_log_sensitive_data(f"Generating tests for function {data.function_to_optimize.function_name}")
# Use different LLMs for different test_index values to get more diverse tests
# test_index 0 uses OpenAI model, test_index 1 uses Anthropic model (or vice versa)
test_index = data.test_index if data.test_index is not None else 0
# Select model based on test_index: even indices use OpenAI, odd indices use Anthropic
if test_index % 2 == 0:
# Prefer OpenAI if available, fall back to Anthropic
execute_model = OpenAI_GPT_4_1() if llm_clients.get("openai") else Anthropic_Claude_Sonnet_4_5()
model_source = "OpenAI"
else:
# Prefer Anthropic if available, fall back to OpenAI
execute_model = Anthropic_Claude_Sonnet_4_5() if llm_clients.get("anthropic") else OpenAI_GPT_4_1()
model_source = "Anthropic"
logging.info(f"Using {model_source} model ({execute_model.name}) for test_index {test_index} to generate diverse tests")
(
generated_test_source,
instrumented_behavior_tests,
@ -431,6 +447,7 @@ async def testgen(
is_async=data.is_async,
trace_id=data.trace_id,
call_sequence=data.call_sequence,
execute_model=execute_model,
)
ph(request.user, "aiservice-testgen-tests-generated")