cleaning up

This commit is contained in:
aseembits93 2026-01-14 19:01:43 -08:00
parent 9d9bf41858
commit 9c4c7126a4

View file

@ -14,7 +14,7 @@ from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUs
from pydantic import ValidationError
from aiservice.analytics.posthog import ph
from aiservice.common_utils import parse_python_version, should_hack_for_demo, validate_trace_id
from aiservice.common_utils import parse_python_version, validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data, debug_log_sensitive_data_from_callable
from aiservice.llm import LLM, OPTIMIZE_MODEL, calculate_llm_cost, call_llm
from authapp.auth import AuthenticatedRequest
@ -30,7 +30,7 @@ from optimizer.context_utils.optimizer_context import (
OptimizeResponseSchema,
)
from optimizer.diff_patches_utils.diff import DiffMethod
from optimizer.models import OptimizedCandidateSource, OptimizeSchema, get_model_from_name
from optimizer.models import OptimizedCandidateSource, OptimizeSchema
if TYPE_CHECKING:
from openai.types.chat import ChatCompletionMessageParam
@ -127,8 +127,7 @@ async def jit_rewrite_python_code(
return [], 0.0, {}, {}
async with asyncio.TaskGroup() as tg:
for model_name, num_calls in get_model_distribution(n_candidates, MAX_OPTIMIZER_CALLS):
model = get_model_from_name(model_name)
for model, num_calls in get_model_distribution(n_candidates, MAX_OPTIMIZER_CALLS):
for _ in range(num_calls):
# Each call needs its own context instance to avoid shared state issues
# Use original_source_code (markdown format) to preserve file path info
@ -284,11 +283,13 @@ async def jit_rewrite(
for item in jit_rewrite_response_items:
item.optimization_event_id = str(event.id) if event else None
response = OptimizeResponseSchema(optimizations=jit_rewrite_response_items)
def log_response() -> None:
debug_log_sensitive_data(f"Response:\n{response.model_dump_json()}")
for opt in response.optimizations:
debug_log_sensitive_data(f"Optimized source:\n{opt.source_code}")
debug_log_sensitive_data(f"Optimization explanation:\n{opt.explanation}")
debug_log_sensitive_data_from_callable(log_response)
ph(request.user, "aiservice-jit-rewrite-successful")
return 200, response