Update optimizer.py

django/aiservice/optimizer/optimizer.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
+import asyncio
 import logging
-import sentry_sdk
-import time
 import uuid
 from pathlib import Path
 from typing import TYPE_CHECKING
 
 import libcst as cst
+import sentry_sdk
 from ninja import NinjaAPI, Schema
-from openai import OpenAIError
 from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam
 from pydantic import ValidationError
 
@@ -66,6 +65,17 @@ optimizations_json = [
     },
 ]
 
+async def hack_for_demo(ctx: BaseOptimizerContext) -> list[OptimizeResponseSchema]:
+    response_list: list[OptimizeResponseItemSchema] = [
+        OptimizeResponseItemSchema(
+            explanation=optimization["explanation"],
+            optimization_id=optimization["optimization_id"],
+            source_code=group_code({ctx.file_name: optimization["source_code"]}),
+        )
+        for optimization in optimizations_json
+    ]
+    await asyncio.sleep(5)
+    return OptimizeResponseSchema(optimizations=response_list)
 
 optimize_api = NinjaAPI(urls_namespace="optimize")
 
@@ -160,6 +170,7 @@ async def optimize_python_code(
             ctx.extracted_code_and_expl = None
             ctx.parsed_code_and_explanation = None
         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}")
             continue
@@ -205,11 +216,19 @@ async def optimize(
         )
     if not validate_trace_id(data.trace_id):
         return 400, OptimizeErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
+
     if should_hack_for_demo(request.user, ctx):
         return hack_for_demo(ctx)
-    optimization_response_items, llm_cost = await optimize_python_code(
-        request.user, ctx, data.dependency_code, n=min(data.n_candidates or 5, 5), python_version=python_version
-    )
+
+    try:
+        optimization_response_items, llm_cost = await optimize_python_code(
+            request.user, ctx, data.dependency_code, n=min(data.n_candidates or 5, 5), python_version=python_version
+        )
+    except Exception as e:
+        logging.exception("Error in optimize endpoint")
+        sentry_sdk.capture_exception(e)
+        return 500, OptimizeErrorResponseSchema(error="Error generating optimizations. Internal server error.")
+
     if len(optimization_response_items) == 0:
         ph(request.user, "aiservice-optimize-no-optimizations-found")
         debug_log_sensitive_data(f"No optimizations found for source:\n{data.source_code}")
@@ -278,14 +297,3 @@ def should_hack_for_demo(user_id: str, ctx: BaseOptimizerContext) -> bool:
     return bool(is_codeflash_employee(user_id) and "def find_common_tags(articles" in ctx.source_code)
 
 
-def hack_for_demo(ctx: BaseOptimizerContext) -> list[OptimizeResponseSchema]:
-    response_list: list[OptimizeResponseItemSchema] = [
-        OptimizeResponseItemSchema(
-            explanation=optimization["explanation"],
-            optimization_id=optimization["optimization_id"],
-            source_code=group_code({ctx.file_name: optimization["source_code"]}),
-        )
-        for optimization in optimizations_json
-    ]
-    time.sleep(5)
-    return OptimizeResponseSchema(optimizations=response_list)