llm call optimization fail error log and small refactoring (#2447)

django/aiservice/aiservice/llm.py

@@ -230,6 +230,15 @@ async def call_llm(
 
     except Exception as e:
         error = e
+        logger.error(
+            "LLM call failed: model=%s, provider=%s, call_type=%s, trace_id=%s, error=%s: %s",
+            llm.name,
+            llm.model_type,
+            call_type,
+            trace_id,
+            type(e).__name__,
+            e,
+        )
         raise
 
     finally:

django/aiservice/core/languages/python/optimizer/optimizer.py

@@ -115,7 +115,7 @@ ASYNC_USER_PROMPT = (current_dir / "async_user_prompt.md").read_text()
 JIT_INSTRUCTIONS = (current_dir / "jit_instructions.md").read_text()
 
 
-async def optimize_python_code_single(
+async def generate_optimization_candidate(
     user_id: str,
     ctx: BaseOptimizerContext,
     trace_id: str,
@@ -153,14 +153,14 @@ async def optimize_python_code_single(
             context=obs_context,
         )
     except Exception as e:
-        logging.exception("OpenAI Code Generation error in optimizer")
+        logging.exception("LLM code generation error in optimizer (model=%s)", optimize_model.name)
         sentry_sdk.capture_exception(e)
         debug_log_sensitive_data(f"Failed to generate code for source:\n{ctx.source_code}")
         return None, None, optimize_model.name
 
     llm_cost = calculate_llm_cost(output.raw_response, optimize_model)
 
-    debug_log_sensitive_data(f"OpenAIClient optimization response:\n{output.raw_response.model_dump_json(indent=2)}")
+    debug_log_sensitive_data(f"LLM optimization response:\n{output.raw_response.model_dump_json(indent=2)}")
 
     if output.raw_response.usage is not None:
         await asyncio.to_thread(
@@ -220,7 +220,7 @@ async def optimize_python_code(
                     ctx.base_system_prompt, ctx.base_user_prompt, original_source_code
                 )
                 task = tg.create_task(
-                    optimize_python_code_single(
+                    generate_optimization_candidate(
                         user_id=user_id,
                         ctx=task_ctx,
                         trace_id=trace_id,

django/aiservice/core/languages/python/testgen/testgen.py

@@ -9,17 +9,17 @@ from typing import TYPE_CHECKING, TypedDict
 
 import sentry_sdk
 import stamina
-from jinja2 import Environment, FileSystemLoader, StrictUndefined
-from libcst import parse_module
-from ninja.errors import HttpError
-from openai import OpenAIError
-
 from aiservice.analytics.posthog import ph
 from aiservice.common.markdown_utils import extract_code_block
 from aiservice.common_utils import parse_python_version, safe_isort, should_hack_for_demo, validate_trace_id
 from aiservice.env_specific import debug_log_sensitive_data
 from aiservice.llm import EXECUTE_MODEL, HAIKU_MODEL, OPENAI_MODEL, calculate_llm_cost, call_llm
 from aiservice.models.functions_to_optimize import FunctionToOptimize
+from jinja2 import Environment, FileSystemLoader, StrictUndefined
+from libcst import parse_module
+from ninja.errors import HttpError
+from openai import OpenAIError
+
 from core.languages.python.cst_utils import parse_module_to_cst
 from core.languages.python.testgen.instrumentation.edit_generated_test import replace_definition_with_import
 from core.languages.python.testgen.instrumentation.instrument_new_tests import instrument_test_source
@@ -41,10 +41,10 @@ from core.shared.testgen_models import (
 )
 
 if TYPE_CHECKING:
-    from openai.types.chat import ChatCompletionMessageParam
-
     from aiservice.llm import LLM
     from authapp.auth import AuthenticatedRequest
+    from openai.types.chat import ChatCompletionMessageParam
+
     from core.shared.testgen_models import TestGenSchema
 
 
@@ -208,7 +208,6 @@ async def generate_and_validate_test_code(
     ctx: BaseTestGenContext,
     python_version: tuple[int, int, int],
     error_context: str,
-    execute_model: LLM,
     cost_tracker: list[float],
     user_id: str,
     posthog_event_suffix: str,
@@ -252,19 +251,17 @@ async def generate_and_validate_test_code(
         context=obs_context,
     )
 
-    cost = calculate_llm_cost(response.raw_response, execute_model)
+    cost = calculate_llm_cost(response.raw_response, model)
     cost_tracker.append(cost)
 
-    debug_log_sensitive_data(
-        f"OpenAIClient {error_context}execute response:\n{response.raw_response.model_dump_json(indent=2)}"
-    )
+    debug_log_sensitive_data(f"LLM {error_context}execute response:\n{response.raw_response.model_dump_json(indent=2)}")
 
     if response.raw_response.usage:
         await asyncio.to_thread(
             ph,
             user_id,
             f"aiservice-testgen-{posthog_event_suffix}execute-openai-usage",
-            properties={"model": execute_model.name, "usage": response.raw_response.usage.model_dump_json()},
+            properties={"model": model.name, "usage": response.raw_response.usage.model_dump_json()},
         )
 
     raw_llm_content = response.content
@@ -314,7 +311,6 @@ async def generate_regression_tests_from_function(
             ctx=ctx,
             python_version=python_version,
             error_context=error_context,
-            execute_model=execute_model,
             cost_tracker=cost_tracker,
             user_id=user_id,
             posthog_event_suffix=posthog_event_suffix,

js/cf-api/resend/completed_optimization_email_template.html

@@ -210,7 +210,7 @@
       <div class="header">
         <div class="logo-wrapper">
           <img
-            src="https://app.codeflash.ai/images/codeflash_light.svg"
+            src="https://app.codeflash.ai/images/codeflash_light.png"
             alt="Codeflash Logo"
             style="height: 40px; width: auto; vertical-align: middle"
           />
@@ -224,8 +224,8 @@
         <div class="greeting">Dear {{userName}},</div>
 
         <div class="main-message">
-          <h2>Optimization Completed</h2>
-          <p>We are pleased to inform you that your optimization has been successfully completed. A pull request has been created with the optimized code changes.</p>
+          <h2>Optimization Complete</h2>
+          <p>Your code has been optimized and a pull request is ready for review.</p>
         </div>
 
       {{PR}}
@@ -241,4 +241,4 @@
       </div>
     </div>
   </body>
-</html>
+</html>

js/cf-webapp/src/app/observability/llm-calls/page.tsx

@@ -265,8 +265,9 @@ export default async function LLMCallsPage({ searchParams }: { searchParams: Sea
     where.trace_id = { in: filteredTraceIds }
   }
 
-    // Fetch LLM calls with pagination, aggregate stats
-    const [llmCalls, totalCount, aggregateStats, successCount] = await Promise.all([
+    // Fetch LLM calls with pagination and count
+    // Batched into groups of 2 to avoid exhausting the connection pool
+    const [llmCalls, totalCount] = await Promise.all([
       prisma.llm_calls.findMany({
         where,
         orderBy: { created_at: "desc" },
@@ -291,7 +292,10 @@ export default async function LLMCallsPage({ searchParams }: { searchParams: Sea
         },
       }),
       prisma.llm_calls.count({ where }),
-      // Get aggregate stats for all filtered data (not just current page)
+    ])
+
+    // Aggregate stats and success count (second batch)
+    const [aggregateStats, successCount] = await Promise.all([
       prisma.llm_calls.aggregate({
         where,
         _sum: {
@@ -305,7 +309,6 @@ export default async function LLMCallsPage({ searchParams }: { searchParams: Sea
           status: true,
         },
       }),
-      // Get success count for success rate calculation
       prisma.llm_calls.count({
         where: {
           ...where,
@@ -371,16 +374,10 @@ export default async function LLMCallsPage({ searchParams }: { searchParams: Sea
         .map(f => f.trace_id.substring(0, 36)),
     )
 
-    // Get unique call types and models for filters
+    // Get unique call types and models for filters (cached)
     const [callTypes, models] = await Promise.all([
-      prisma.llm_calls.findMany({
-      select: { call_type: true },
-      distinct: ["call_type"],
-    }),
-    prisma.llm_calls.findMany({
-      select: { model_name: true },
-      distinct: ["model_name"],
-    }),
+      getCallTypes(),
+      getModels(),
     ])
 
     const totalPages = Math.ceil(totalCount / pageSize)

js/cf-webapp/src/lib/prisma.ts

@@ -4,6 +4,19 @@ const globalForPrisma = globalThis as unknown as {
   prisma: PrismaClient | undefined
 }
 
-export const prisma = globalForPrisma.prisma ?? new PrismaClient()
+function buildDatabaseUrl() {
+  const baseUrl = process.env.DATABASE_URL ?? ""
+  if (baseUrl.includes("connection_limit")) return baseUrl
+  const separator = baseUrl.includes("?") ? "&" : "?"
+  return `${baseUrl}${separator}connection_limit=10&pool_timeout=20`
+}
+
+export const prisma =
+  globalForPrisma.prisma ??
+  new PrismaClient({
+    datasources: {
+      db: { url: buildDatabaseUrl() },
+    },
+  })
 
 if (process.env.NODE_ENV !== "production") globalForPrisma.prisma = prisma