codeflash-admin/codeflash-internal
1
0
Fork 0
mirror of https://github.com/codeflash-ai/codeflash-internal.git synced 2026-05-04 18:25:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

async: parallelize endpoint epilogue DB writes (#2490)

Browse source 
## Summary

Parallelize independent DB writes at the end of 4 endpoints using
`asyncio.TaskGroup`. With psycopg3 connection pooling (#2489), each task
gets its own connection from the pool.

### Endpoints optimized

| Endpoint | Before | After |
|----------|--------|-------|
| **Refinement** | `log_features` then `update_optimization_cost` |
`TaskGroup` (concurrent) |
| **Explanations** | `update_optimization_cost` inside inner fn | Moved
to handler, `TaskGroup` with `log_features` |
| **Optimization review** | `update_optimization_cost` inside inner fn |
Moved to handler, `TaskGroup` with `update_optimization_features_review`
|
| **Ranker** | `update_optimization_cost` inside inner fn | Moved to
handler, `TaskGroup` with `log_features` |

Each endpoint saves ~87ms (one DB round-trip) by overlapping two
independent writes.

### Comprehensive audit

All 13 endpoints were audited — no remaining async antipatterns found:
- No blocking calls in async paths
- No `await`-in-loop patterns
- LLM clients already use connection reuse
- All other endpoints have at most 1 DB write in the epilogue

## Test plan

- [x] All 538 tests passing
- [ ] Verify under load in staging

---------

Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-authored-by: Kevin Turcios <KRRT7@users.noreply.github.com>
Co-authored-by: codeflash-ai[bot] <148906541+codeflash-ai[bot]@users.noreply.github.com>
This commit is contained in:
Kevin Turcios 2026-04-01 06:15:16 -05:00 committed by GitHub
parent 2887b34d02
commit 0abc6bf1e3
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 128 additions and 72 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
.github/workflows/claude.yml vendored
View file

@ -53,6 +53,7 @@ jobs:
with:
use_bedrock: "true"
use_sticky_comment: true
allowed_bots: "claude[bot],codeflash-ai[bot]"
prompt: |
REPO: ${{ github.repository }}
PR NUMBER: ${{ github.event.pull_request.number }}

46
django/aiservice/core/languages/python/explanations/explanations.py
View file

@ -1,7 +1,11 @@
from __future__ import annotations
import asyncio
from collections.abc import Coroutine
from pathlib import Path
from typing import Any
import sentry_sdk
from ninja import NinjaAPI
from openai.types.chat import (
ChatCompletionMessageParam,
@ -11,10 +15,11 @@ from openai.types.chat import (
from packaging import version
from aiservice.analytics.posthog import ph
from authapp.auth import AuthenticatedRequest
from aiservice.common.markdown_utils import wrap_code_in_markdown
from aiservice.common.xml_utils import extract_xml_tag
from aiservice.common_utils import validate_trace_id
from aiservice.env_specific import debug_log_sensitive_data
from aiservice.env_specific import IS_PRODUCTION, debug_log_sensitive_data
from aiservice.llm import llm_client
from aiservice.llm_models import EXPLANATIONS_MODEL, LLM
from core.languages.python.explanations.models import (
@ -26,6 +31,8 @@ from core.log_features.log_event import update_optimization_cost
from core.log_features.log_features import log_features
from core.shared.jinja_utils import create_prompt_env
_PARSED_0182 = version.parse("0.18.2")
explanations_api = NinjaAPI(urls_namespace="explanations")
_PROMPT_DIR = Path(__file__).parent / "prompts"
@ -37,9 +44,14 @@ USER_PROMPT_TEMPLATE = _jinja_env.get_template("user_prompt.md.j2")
async def explain_optimizations(
user_id: str, data: ExplanationsSchema, explanations_model: LLM = EXPLANATIONS_MODEL
) -> ExplanationsResponseSchema | ExplanationsErrorResponseSchema:
debug_log_sensitive_data(f"Generating an explanation for {user_id}:\n{data.optimized_code}")
if version.parse(data.codeflash_version) <= version.parse("0.18.2") and data.annotated_tests:
) -> tuple[ExplanationsResponseSchema, float] | ExplanationsErrorResponseSchema:
# Avoid building potentially very large debug strings when logging is disabled.
if not IS_PRODUCTION:
debug_log_sensitive_data(f"Generating an explanation for {user_id}:\n{data.optimized_code}")
# Parse the incoming version once per call (faster than parsing the constant each time).
parsed_codeflash_version = version.parse(data.codeflash_version)
if parsed_codeflash_version <= _PARSED_0182 and data.annotated_tests:
data.annotated_tests = wrap_code_in_markdown(data.annotated_tests)
include_throughput = data.original_throughput is not None and data.optimized_throughput is not None
@ -81,7 +93,8 @@ async def explain_optimizations(
system_message = ChatCompletionSystemMessageParam(role="system", content=system_prompt)
user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)
debug_log_sensitive_data(f"{system_prompt}{user_prompt}")
if not IS_PRODUCTION:
debug_log_sensitive_data(f"{system_prompt}{user_prompt}")
messages: list[ChatCompletionMessageParam] = [system_message, user_message]
obs_context: dict[str, str | float | int] = {"optimization_id": data.optimization_id, "speedup": data.speedup}
@ -97,10 +110,10 @@ async def explain_optimizations(
user_id=user_id,
context=obs_context,
)
await update_optimization_cost(trace_id=data.trace_id, cost=output.cost, user_id=user_id)
except Exception as e:
return ExplanationsErrorResponseSchema(error=str(e))
debug_log_sensitive_data(f"AIClient optimization response:\n{output.content}")
if not IS_PRODUCTION:
debug_log_sensitive_data(f"AIClient optimization response:\n{output.content}")
if output.usage is not None:
ph(
user_id,
@ -111,7 +124,7 @@ async def explain_optimizations(
"usage": {"input_tokens": output.usage.input_tokens, "output_tokens": output.usage.output_tokens},
},
)
return ExplanationsResponseSchema(explanation=output.content)
return ExplanationsResponseSchema(explanation=output.content), output.cost
@explanations_api.post(
@ -123,23 +136,30 @@ async def explain_optimizations(
},
)
async def explain(
request, # noqa: ANN001
data: ExplanationsSchema,
request: AuthenticatedRequest, data: ExplanationsSchema
) -> tuple[int, ExplanationsResponseSchema | ExplanationsErrorResponseSchema]:
ph(request.user, "aiservice-explain-called")
if not validate_trace_id(data.trace_id):
return 400, ExplanationsErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
explanation_response = await explain_optimizations(request.user, data)
if isinstance(explanation_response, ExplanationsErrorResponseSchema):
result = await explain_optimizations(request.user, data)
if isinstance(result, ExplanationsErrorResponseSchema):
ph(request.user, "Explanation not generated, revert to old explanation")
debug_log_sensitive_data("No explanation was generated")
return 500, ExplanationsErrorResponseSchema(error="Error generating optimizations. Internal server error.")
explanation_response, llm_cost = result
ph(request.user, "explanation generated", properties={"explanation": explanation_response})
# parse xml tag for explanation
explanation = extract_xml_tag(explanation_response.explanation, "explain")
if not explanation:
return 500, ExplanationsErrorResponseSchema(error="Failed to parse explanation from LLM response.")
coros: list[Coroutine[Any, Any, Any]] = [
update_optimization_cost(trace_id=data.trace_id, cost=llm_cost, user_id=request.user)
]
if hasattr(request, "should_log_features") and request.should_log_features:
await log_features(trace_id=data.trace_id, user_id=request.user, final_explanation=explanation)
coros.append(log_features(trace_id=data.trace_id, user_id=request.user, final_explanation=explanation))
results = await asyncio.gather(*coros, return_exceptions=True)
for result in results:
if isinstance(result, BaseException):
sentry_sdk.capture_exception(result)
response = ExplanationsResponseSchema(explanation=explanation)
return 200, response

45
django/aiservice/core/languages/python/optimization_review/optimization_review.py
View file

@ -1,5 +1,6 @@
from __future__ import annotations
import asyncio
import json
import logging
from enum import Enum
@ -176,8 +177,11 @@ async def get_optimization_review(
request: AuthenticatedRequest,
data: OptimizationReviewSchema,
optimization_review_model: LLM = OPTIMIZATION_REVIEW_MODEL,
) -> tuple[int, OptimizationReviewResponseSchema | OptimizationReviewErrorSchema]:
"""Compute optimization review via Claude."""
) -> tuple[int, OptimizationReviewResponseSchema | OptimizationReviewErrorSchema, float]:
"""Compute optimization review via Claude.
Returns (status_code, response, llm_cost).
"""
ph(request.user, "aiservice-optimization-review-called")
try:
@ -199,7 +203,6 @@ async def get_optimization_review(
)
cost = response.cost
await update_optimization_cost(data.trace_id, cost, user_id=request.user)
review_text = response.content.strip()
if result := extract_code_block_with_context(review_text, language="json"):
@ -247,16 +250,16 @@ async def get_optimization_review(
logging.exception("Invalid optimization review response")
sentry_sdk.capture_exception(e)
debug_log_sensitive_data(f"Invalid response : {e}")
return 500, OptimizationReviewErrorSchema(error="Invalid response")
return 500, OptimizationReviewErrorSchema(error="Invalid response"), cost
else:
ph(request.user, "aiservice-optimization-review-successful")
return 200, review
return 200, review, cost
else:
return 500, OptimizationReviewErrorSchema(error="Invalid response")
return 500, OptimizationReviewErrorSchema(error="Invalid response"), cost
except Exception as e:
logging.exception("Error in optimization_review")
sentry_sdk.capture_exception(e)
return 500, OptimizationReviewErrorSchema(error="Internal server error")
return 500, OptimizationReviewErrorSchema(error="Internal server error"), 0.0
@optimization_review_api.post(
@ -270,20 +273,24 @@ async def get_optimization_review(
async def optimization_review(
request: AuthenticatedRequest, data: OptimizationReviewSchema
) -> tuple[int, OptimizationReviewResponseSchema | OptimizationReviewErrorSchema]:
response_code, output = await get_optimization_review(request, data)
try:
if isinstance(output, OptimizationReviewResponseSchema):
review_event = output.review.value
review_explanation = output.review_explanation
else:
review_event = output.error
review_explanation = ""
await update_optimization_features_review(
response_code, output, llm_cost = await get_optimization_review(request, data)
if isinstance(output, OptimizationReviewResponseSchema):
review_event = output.review.value
review_explanation = output.review_explanation
else:
review_event = output.error
review_explanation = ""
results = await asyncio.gather(
update_optimization_cost(data.trace_id, llm_cost, user_id=request.user),
update_optimization_features_review(
trace_id=data.trace_id,
review_quality=review_event,
review_explanation=review_explanation,
calling_fn_details=data.calling_fn_details,
)
except Exception as e: # noqa: BLE001
debug_log_sensitive_data(f"event logging failed for optimization review {e}")
),
return_exceptions=True,
)
for result in results:
if isinstance(result, BaseException):
sentry_sdk.capture_exception(result)
return response_code, output

55
django/aiservice/core/languages/python/optimizer/refinement.py
View file

@ -2,8 +2,9 @@ from __future__ import annotations
import asyncio
import uuid
from collections.abc import Coroutine
from pathlib import Path
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any
import libcst as cst
import sentry_sdk
@ -81,7 +82,7 @@ async def refinement( # noqa: D417
messages: list[ChatCompletionMessageParam] = [system_message, user_message]
debug_log_sensitive_data(f"This was the user prompt\n {user_prompt}\n")
obs_context: dict = {"optimization_id": optimization_id, "speedup": ctx.data.speedup}
obs_context: dict[str, object] = {"optimization_id": optimization_id, "speedup": ctx.data.speedup}
if call_sequence is not None:
obs_context["call_sequence"] = call_sequence
@ -247,28 +248,38 @@ async def refine(
if (elem.source_code.strip() not in source_code_set) and elem.source_code != "":
source_code_set.add(elem.source_code.strip())
filtered_refined_optimizations.append(elem)
# Parallel DB writes: log_features + update_optimization_cost run concurrently.
# gather(return_exceptions=True) ensures both complete even if one fails.
coros: list[Coroutine[Any, Any, Any]] = [
update_optimization_cost(trace_id=trace_id, cost=total_llm_cost, user_id=request.user)
]
if hasattr(request, "should_log_features") and request.should_log_features:
await log_features(
trace_id=trace_id,
user_id=request.user,
optimizations_raw={
cei.optimization_id: cei.source_code
for cei in refinement_data
if not isinstance(cei, OptimizeErrorResponseSchema)
},
optimizations_post={cei.optimization_id: cei.source_code for cei in filtered_refined_optimizations},
explanations_raw={
cei.optimization_id: cei.explanation
for cei in refinement_data
if not isinstance(cei, OptimizeErrorResponseSchema)
},
explanations_post={cei.optimization_id: cei.explanation for cei in filtered_refined_optimizations},
optimizations_origin={
cei.optimization_id: {"source": OptimizedCandidateSource.REFINE, "parent": cei.parent_id}
for cei in filtered_refined_optimizations
},
coros.append(
log_features(
trace_id=trace_id,
user_id=request.user,
optimizations_raw={
cei.optimization_id: cei.source_code
for cei in refinement_data
if not isinstance(cei, OptimizeErrorResponseSchema)
},
optimizations_post={cei.optimization_id: cei.source_code for cei in filtered_refined_optimizations},
explanations_raw={
cei.optimization_id: cei.explanation
for cei in refinement_data
if not isinstance(cei, OptimizeErrorResponseSchema)
},
explanations_post={cei.optimization_id: cei.explanation for cei in filtered_refined_optimizations},
optimizations_origin={
cei.optimization_id: {"source": OptimizedCandidateSource.REFINE, "parent": cei.parent_id}
for cei in filtered_refined_optimizations
},
)
)
await update_optimization_cost(trace_id=trace_id, cost=total_llm_cost, user_id=request.user)
results = await asyncio.gather(*coros, return_exceptions=True)
for result in results:
if isinstance(result, BaseException):
sentry_sdk.capture_exception(result)
return 200, Refinementschema(
refinements=[
OptimizeResponseItemSchema(

53
django/aiservice/core/shared/ranker/ranker.py
View file

@ -1,11 +1,14 @@
from __future__ import annotations
import asyncio
import json
import logging
import re
from collections.abc import Coroutine
from dataclasses import dataclass
from typing import TYPE_CHECKING
from typing import TYPE_CHECKING, Any, cast
import sentry_sdk
from ninja import NinjaAPI, Schema
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
@ -145,7 +148,7 @@ class ParsedRankingResponse:
explanation: str
def _extract_json_from_response(content: str) -> dict | None:
def _extract_json_from_response(content: str) -> dict[str, Any] | None:
"""Extract JSON object from LLM response.
Handles:
@ -155,7 +158,7 @@ def _extract_json_from_response(content: str) -> dict | None:
"""
# Try direct JSON parse first
try:
return json.loads(content.strip())
return cast(dict[str, Any], json.loads(content.strip()))
except json.JSONDecodeError:
pass
@ -165,7 +168,7 @@ def _extract_json_from_response(content: str) -> dict | None:
json_str = match.group(1) or match.group(2)
if json_str:
try:
return json.loads(json_str.strip())
return cast(dict[str, Any], json.loads(json_str.strip()))
except json.JSONDecodeError:
pass
@ -183,7 +186,7 @@ def _extract_json_from_response(content: str) -> dict | None:
depth -= 1
if depth == 0:
json_str = content[start : i + 1]
return json.loads(json_str)
return cast(dict[str, Any], json.loads(json_str))
except json.JSONDecodeError:
pass
@ -318,7 +321,7 @@ def _scores_to_ranking(scores: CandidateScores) -> list[int]:
async def rank_optimizations( # noqa: D417
user_id: str, data: RankInputSchema, rank_model: LLM = RANKING_MODEL
) -> RankResponseSchema | RankErrorResponseSchema:
) -> tuple[RankResponseSchema | RankErrorResponseSchema, float]:
"""Rank optimization candidates using multi-dimensional scoring.
Parameters
@ -360,10 +363,9 @@ async def rank_optimizations( # noqa: D417
"python_version": data.python_version,
},
)
await update_optimization_cost(trace_id=data.trace_id, cost=output.cost, user_id=user_id)
except Exception:
logging.exception("Ranking failed for trace_id=%s", data.trace_id)
return RankErrorResponseSchema(error="Failed to rank optimizations. Please try again.")
return RankErrorResponseSchema(error="Failed to rank optimizations. Please try again."), 0.0
debug_log_sensitive_data(f"AIClient optimization response:\n{output}")
if output.raw_response.usage is not None:
@ -379,10 +381,16 @@ async def rank_optimizations( # noqa: D417
json_response = _parse_json_response(output.content, num_candidates)
if json_response is not None:
logging.info("Successfully parsed JSON response")
ranking = (
_scores_to_ranking(json_response.scores) if json_response.scores is not None else json_response.ranking
return (
RankResponseSchema(
ranking=_scores_to_ranking(json_response.scores)
if json_response.scores is not None
else json_response.ranking,
explanation=json_response.explanation,
scores=json_response.scores,
),
output.cost,
)
return RankResponseSchema(ranking=ranking, explanation=json_response.explanation, scores=json_response.scores)
# Fall back to regex parsing (legacy XML-tag format)
logging.info("JSON parsing failed, falling back to regex")
@ -424,9 +432,9 @@ async def rank_optimizations( # noqa: D417
logging.info("Derived ranking from scores")
else:
logging.warning("No valid ranking found")
return RankErrorResponseSchema(error="No ranking found")
return RankErrorResponseSchema(error="No ranking found"), output.cost
return RankResponseSchema(ranking=ranking, explanation=explanation, scores=scores)
return RankResponseSchema(ranking=ranking, explanation=explanation, scores=scores), output.cost
class RankInputSchema(Schema):
@ -455,20 +463,29 @@ async def rank(
ph(request.user, "aiservice-rank-called")
if not validate_trace_id(data.trace_id):
return 400, RankErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
ranking_response = await rank_optimizations(request.user, data)
ranking_response, llm_cost = await rank_optimizations(request.user, data)
if isinstance(ranking_response, RankErrorResponseSchema):
ph(request.user, "Invalid Ranking, fallback to default")
debug_log_sensitive_data("No valid ranking was generated")
return 500, RankErrorResponseSchema(error="Error generating ranking. Internal server error.")
ph(request.user, "ranking generated", properties={"ranking": ranking_response})
ranking_0_idx = [x - 1 for x in ranking_response.ranking]
coros: list[Coroutine[Any, Any, Any]] = [
update_optimization_cost(trace_id=data.trace_id, cost=llm_cost, user_id=request.user)
]
if hasattr(request, "should_log_features") and request.should_log_features:
ranked_opt_ids = [data.optimization_ids[i] for i in ranking_0_idx]
await log_features(
trace_id=data.trace_id,
user_id=request.user,
ranking={"ranking": ranked_opt_ids, "explanation": ranking_response.explanation},
coros.append(
log_features(
trace_id=data.trace_id,
user_id=request.user,
ranking={"ranking": ranked_opt_ids, "explanation": ranking_response.explanation},
)
)
results = await asyncio.gather(*coros, return_exceptions=True)
for result in results:
if isinstance(result, BaseException):
sentry_sdk.capture_exception(result)
return 200, RankResponseSchema(
ranking=ranking_0_idx, explanation=ranking_response.explanation, scores=ranking_response.scores
) # we don't really use `explanation` and `score` but still returning it for future use in CLI