codeflash-ai/codeflash-internal
1
0
Fork 0
Code Issues Pull requests 49 Projects Releases Packages Wiki Activity Actions 2.9k

include the model name

Browse source
This commit is contained in:
Kevin Turcios 2025-12-26 19:22:55 -05:00
parent 6f45e0d0f1
commit 262c193365
3 changed files with 43 additions and 23 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

27
django/aiservice/optimizer/optimizer.py
View file

@ -123,7 +123,7 @@ async def optimize_python_code_single(
optimize_model: LLM = OPTIMIZE_MODEL,
python_version: tuple[int, int, int] = (3, 12, 9),
call_sequence: int | None = None,
) -> tuple[OptimizeResponseItemSchema | None, float | None]:
) -> tuple[OptimizeResponseItemSchema | None, float | None, str]:
"""Optimize the given python code for performance using LLMs."""
logging.info("/optimize: Optimizing python code.")
debug_log_sensitive_data(f"Optimizing python code for user {user_id}:\n{ctx.source_code}")
@ -156,7 +156,7 @@ async def optimize_python_code_single(
logging.exception("OpenAI Code Generation error in optimizer")
sentry_sdk.capture_exception(e)
debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.source_code}")
return None, None
return None, None, optimize_model.name
llm_cost = calculate_llm_cost(output.raw_response, optimize_model)
@ -173,13 +173,13 @@ async def optimize_python_code_single(
try:
res = ctx.parse_and_generate_candidate_schema()
if res is not None and ctx.is_valid_code():
return res, llm_cost
return res, llm_cost, optimize_model.name
except (ValueError, ValidationError, cst.ParserSyntaxError) as e:
sentry_sdk.capture_message(f"Error parsing optimization result: {e}")
debug_log_sensitive_data(f"error for source:\n{ctx.source_code}")
debug_log_sensitive_data(f"Traceback: {e}")
return None, llm_cost
return None, llm_cost, optimize_model.name
async def optimize_python_code(
@ -190,7 +190,7 @@ async def optimize_python_code(
dependency_code: str | None = None,
python_version: tuple[int, int, int] = (3, 12, 9),
model_distribution: list[tuple[str, int]] | None = None,
) -> tuple[list[OptimizeResponseItemSchema], float, dict[str, dict[str, str]]]:
) -> tuple[list[OptimizeResponseItemSchema], float, dict[str, dict[str, str]], dict[str, str]]:
"""Run parallel optimizations with multiple models based on the distribution config.
Returns:
@ -198,13 +198,14 @@ async def optimize_python_code(
- list of optimization results
- total LLM cost
- dict of raw code/explanations keyed by optimization_id
- dict mapping optimization_id to model name
"""
if model_distribution is None:
model_distribution = MODEL_DISTRIBUTION
# Create tasks for each model call
tasks: list[tuple[asyncio.Task[tuple[OptimizeResponseItemSchema | None, float | None]], BaseOptimizerContext]] = []
tasks: list[tuple[asyncio.Task[tuple[OptimizeResponseItemSchema | None, float | None, str]], BaseOptimizerContext]] = []
call_sequence = 1
async with asyncio.TaskGroup() as tg:
@ -234,18 +235,20 @@ async def optimize_python_code(
optimization_results: list[OptimizeResponseItemSchema] = []
total_cost = 0.0
code_and_explanations: dict[str, dict[str, str]] = {}
optimization_models: dict[str, str] = {}
for task, task_ctx in tasks:
result, cost = task.result()
result, cost, model_name = task.result()
if cost:
total_cost += cost
if result is not None:
optimization_results.append(result)
optimization_models[result.optimization_id] = model_name
# Collect raw code/explanations for logging
for op_id, cei in task_ctx.code_and_explanation_before_post_processing.items():
code_and_explanations[op_id] = {"code": cei.code, "explanation": cei.explanation}
return optimization_results, total_cost, code_and_explanations
return optimization_results, total_cost, code_and_explanations, optimization_models
def validate_request_data(data: OptimizeSchema, ctx: BaseOptimizerContext) -> tuple[int, int, int]:
@ -342,7 +345,7 @@ async def optimize(
sentry_sdk.capture_exception(e)
return 500, OptimizeErrorResponseSchema(error="Error generating optimizations. Internal server error.")
optimization_response_items, llm_cost, code_and_explanations = optimize_task.result()
optimization_response_items, llm_cost, code_and_explanations, optimization_models = optimize_task.result()
if user_task:
user = await user_task
if user and user.github_username:
@ -388,7 +391,11 @@ async def optimize(
explanations_post={cei.optimization_id: cei.explanation for cei in optimization_response_items},
experiment_metadata=data.experiment_metadata if data.experiment_metadata else None,
optimizations_origin={
cei.optimization_id: {"source": OptimizedCandidateSource.OPTIMIZE, "parent": None}
cei.optimization_id: {
"source": OptimizedCandidateSource.OPTIMIZE,
"parent": None,
"model": optimization_models.get(cei.optimization_id, "unknown"),
}
for cei in optimization_response_items
},
)

29
django/aiservice/optimizer/optimizer_line_profiler.py
View file

@ -53,7 +53,7 @@ async def optimize_python_code_line_profiler_single(
lsp_mode: bool = False, # noqa: FBT001, FBT002
python_version: tuple[int, int, int] = (3, 12, 9),
call_sequence: int | None = None,
) -> tuple[OptimizeResponseItemSchema | None, float | None]:
) -> tuple[OptimizeResponseItemSchema | None, float | None, str]:
"""Optimize the given python code for performance using LLMs."""
logging.info("/optimize: Optimizing python code line profile.")
debug_log_sensitive_data(f"Optimizing python code for user {user_id}:\n{ctx.source_code}")
@ -97,7 +97,7 @@ async def optimize_python_code_line_profiler_single(
logging.exception("OpenAI Code Generation error in optimizer-line-profiler")
sentry_sdk.capture_exception(e)
debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.source_code}")
return None, None
return None, None, optimize_model.name
llm_cost = calculate_llm_cost(output.raw_response, optimize_model)
@ -113,9 +113,9 @@ async def optimize_python_code_line_profiler_single(
ctx.extract_code_and_explanation_from_llm_res(output.content)
res = ctx.parse_and_generate_candidate_schema()
if res is not None and ctx.is_valid_code():
return res, llm_cost
return res, llm_cost, optimize_model.name
return None, llm_cost
return None, llm_cost, optimize_model.name
async def optimize_python_code_line_profiler(
@ -127,7 +127,7 @@ async def optimize_python_code_line_profiler(
dependency_code: str | None = None,
lsp_mode: bool = False,
python_version: tuple[int, int, int] = (3, 12, 9),
) -> tuple[list[OptimizeResponseItemSchema], float, dict[str, dict]]:
) -> tuple[list[OptimizeResponseItemSchema], float, dict[str, dict], dict[str, str]]:
"""Run parallel line profiler optimizations with multiple models.
Returns:
@ -135,12 +135,13 @@ async def optimize_python_code_line_profiler(
- list of optimization results
- total LLM cost
- dict of raw code/explanations keyed by optimization_id
- dict mapping optimization_id to model name
"""
model_distribution = MODEL_DISTRIBUTION_LP_LSP if lsp_mode else MODEL_DISTRIBUTION_LP
# Create tasks for each model call
tasks: list[tuple[asyncio.Task, BaseOptimizerContext]] = []
tasks: list[tuple[asyncio.Task[tuple[OptimizeResponseItemSchema | None, float | None, str]], BaseOptimizerContext]] = []
call_sequence = 1
async with asyncio.TaskGroup() as tg:
@ -172,18 +173,20 @@ async def optimize_python_code_line_profiler(
optimization_results: list[OptimizeResponseItemSchema] = []
total_cost = 0.0
code_and_explanations: dict[str, dict] = {}
optimization_models: dict[str, str] = {}
for task, task_ctx in tasks:
result, cost = task.result()
result, cost, model_name = task.result()
if cost:
total_cost += cost
if result is not None:
optimization_results.append(result)
optimization_models[result.optimization_id] = model_name
# Collect raw code/explanations for logging
for op_id, cei in task_ctx.code_and_explanation_before_post_processing.items():
code_and_explanations[op_id] = {"code": cei.code, "explanation": cei.explanation}
return optimization_results, total_cost, code_and_explanations
return optimization_results, total_cost, code_and_explanations, optimization_models
class OptimizeSchemaLP(Schema):
@ -220,7 +223,7 @@ async def optimize(request, data: OptimizeSchemaLP) -> tuple[int, OptimizeRespon
if not validate_trace_id(data.trace_id):
return 400, OptimizeErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
optimization_response_items, llm_cost, code_and_explanations = await optimize_python_code_line_profiler(
optimization_response_items, llm_cost, code_and_explanations, optimization_models = await optimize_python_code_line_profiler(
user_id=request.user,
trace_id=data.trace_id,
line_profiler_results=data.line_profiler_results,
@ -257,7 +260,11 @@ async def optimize(request, data: OptimizeSchemaLP) -> tuple[int, OptimizeRespon
explanations_post={cei.optimization_id: cei.explanation for cei in optimization_response_items},
experiment_metadata=data.experiment_metadata if data.experiment_metadata else None,
optimizations_origin={
cei.optimization_id: {"source": OptimizedCandidateSource.OPTIMIZE_LP, "parent": None}
cei.optimization_id: {
"source": OptimizedCandidateSource.OPTIMIZE_LP,
"parent": None,
"model": optimization_models.get(cei.optimization_id, "unknown"),
}
for cei in optimization_response_items
},
)
@ -265,7 +272,7 @@ async def optimize(request, data: OptimizeSchemaLP) -> tuple[int, OptimizeRespon
response = OptimizeResponseSchema(optimizations=optimization_response_items)
def log_response() -> None:
debug_log_sensitive_data(f"Response:\n{response.json()}")
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}")

10
js/cf-webapp/src/app/observability/trace/[trace_id]/page.tsx
View file

@ -43,8 +43,8 @@ export default async function TracePage({ params }: TracePageProps) {
])
// Extract optimization candidates from optimization_features
// Also get the origin of each candidate (OPTIMIZE vs OPTIMIZE_LP)
const optimizationsOrigin = (optimizationFeatures?.optimizations_origin as Record<string, { source: string }>) || {}
// Also get the origin of each candidate (OPTIMIZE vs OPTIMIZE_LP) and model info
const optimizationsOrigin = (optimizationFeatures?.optimizations_origin as Record<string, { source: string; model?: string }>) || {}
const allCandidates = optimizationFeatures?.optimizations_post
? Object.entries(optimizationFeatures.optimizations_post as Record<string, string>).map(
@ -52,6 +52,7 @@ export default async function TracePage({ params }: TracePageProps) {
id,
code: typeof code === "string" ? code : "",
source: optimizationsOrigin[id]?.source || "OPTIMIZE",
model: optimizationsOrigin[id]?.model,
})
)
: []
@ -273,6 +274,11 @@ export default async function TracePage({ params }: TracePageProps) {
<span className="font-mono text-xs text-gray-500 dark:text-gray-400">
{candidate.id.substring(0, 8)}...
</span>
{candidate.model && (
<span className="bg-purple-100 dark:bg-purple-900 text-purple-800 dark:text-purple-200 text-xs px-2 py-0.5 rounded">
{candidate.model}
</span>
)}
</div>
<span className="text-gray-400 text-sm">▼</span>
</summary>