codeflash-internal/django/aiservice/optimizer/optimizer_line_profiler.py

from __future__ import annotations

import logging
from pathlib import Path
from typing import TYPE_CHECKING

import sentry_sdk
from aiservice.analytics.posthog import ph
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, llm_clients
from aiservice.models.aimodels import OPTIMIZE_MODEL, calculate_llm_cost
from log_features.log_event import update_optimization_cost
from log_features.log_features import log_features
from ninja import NinjaAPI, Schema
from openai.types.chat import ChatCompletionSystemMessageParam, ChatCompletionUserMessageParam

from optimizer.context_utils.optimizer_context import (
    BaseOptimizerContext,
    OptimizeErrorResponseSchema,
    OptimizeResponseSchema,
)

if TYPE_CHECKING:
    from aiservice.models.aimodels import LLM
    from openai.types.chat import (
        ChatCompletionAssistantMessageParam,
        ChatCompletionFunctionMessageParam,
        ChatCompletionToolMessageParam,
    )

    from optimizer.context_utils.optimizer_context import OptimizeResponseItemSchema


optimize_line_profiler_api = NinjaAPI(urls_namespace="optimize-line-profiler")


# Get the directory of the current file
current_dir = Path(__file__).parent
SYSTEM_PROMPT = (current_dir / "system_prompt.md").read_text()
USER_PROMPT = (current_dir / "user_prompt.md").read_text()


async def optimize_python_code_line_profiler(  # noqa: D417
    user_id: str,
    trace_id: str,
    line_profiler_results: str,
    ctx: BaseOptimizerContext,
    dependency_code: str | None = None,
    n: int = 1,
    optimize_model: LLM = OPTIMIZE_MODEL,
    lsp_mode: bool = False,  # noqa: FBT001, FBT002
    python_version: tuple[int, int, int] = (3, 12, 9),
) -> list[OptimizeResponseItemSchema]:
    """Optimize the given python code for performance using OpenAI's GPT-4o model.

    Parameters
    ----------
    - source_code (str): The python code to optimize.
    - n (int): Number of optimization variants to generate. Default is 1.

    Returns: - List[Tuple[Union[str, None], Union[str, None]]]: A list of tuples where the first element is the
    optimized code and the second is the explanation.

    """
    logging.info("/optimize: Optimizing python code line profile.")
    debug_log_sensitive_data(f"Optimizing python code for user {user_id}:\n{ctx.source_code}")
    if user_id in ["github|1235813", "github|1100399"] or lsp_mode:
        # for Galileo and LSP mode, we only generate 5 LP optimizations
        n = 5

    python_version_str = ".".join(str(x) for x in python_version)

    # TODO: Experiment with iterative approaches to optimization. Take the learnings from the testing phase into the
    #  next optimization iteration
    # TODO: Experiment with iterative chain-of-thought generation. ask what is the
    #  function doing and then ask it to describe how to speed it up and then generate optimization
    system_prompt = ctx.get_system_prompt(python_version_str=python_version_str)
    user_prompt = ctx.get_user_prompt(dependency_code, line_profiler_results)

    system_message = ChatCompletionSystemMessageParam(role="system", content=system_prompt)
    user_message = ChatCompletionUserMessageParam(role="user", content=user_prompt)
    messages: list[
        ChatCompletionSystemMessageParam
        | ChatCompletionUserMessageParam
        | ChatCompletionAssistantMessageParam
        | ChatCompletionToolMessageParam
        | ChatCompletionFunctionMessageParam
    ] = [system_message, user_message]
    debug_log_sensitive_data(f"This was the user prompt\n {user_prompt}\n")
    # TODO: Verify if the context window length is within the model capability
    llm_client = llm_clients[optimize_model.model_type]
    try:
        output = await llm_client.with_options(max_retries=3).chat.completions.create(
            model=optimize_model.name, messages=messages, n=n
        )
        await update_optimization_cost(trace_id=trace_id, cost=calculate_llm_cost(output, optimize_model))
    except Exception as e:
        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 []

    debug_log_sensitive_data(f"OpenAIClient optimization response:\n{output.model_dump_json(indent=2)}")

    if output.usage is not None:
        ph(
            user_id,
            "aiservice-optimize-line-profiler-openai-usage",
            properties={"model": optimize_model.name, "n": n, "usage": output.usage.json()},
        )
    results = [content for op in output.choices if (content := op.message.content)]
    optimization_response_items: list[OptimizeResponseItemSchema] = []
    for result in results:
        ctx.extract_code_and_explanation_from_llm_res(result)
        res = ctx.parse_and_generate_candidate_schema()
        if res is not None and ctx.is_valid_code():
            optimization_response_items.append(res)

        ctx.extracted_code_and_expl = None
        ctx.parsed_code_and_explanation = None

    return optimization_response_items


class OptimizeSchemaLP(Schema):
    source_code: str
    dependency_code: str | None
    line_profiler_results: str | None
    trace_id: str
    python_version: str
    experiment_metadata: dict[str, str] | None = None
    codeflash_version: str | None = None
    lsp_mode: bool = False
    n_candidates_lp: int | None = 6


@optimize_line_profiler_api.post(
    "/", response={200: OptimizeResponseSchema, 400: OptimizeErrorResponseSchema, 500: OptimizeErrorResponseSchema}
)
async def optimize(request, data: OptimizeSchemaLP) -> tuple[int, OptimizeResponseSchema | OptimizeErrorResponseSchema]:  # noqa: ANN001
    ph(request.user, "aiservice-optimize-called")
    ctx: BaseOptimizerContext = BaseOptimizerContext.get_dynamic_context(SYSTEM_PROMPT, USER_PROMPT, data.source_code)
    try:
        python_version: tuple[int, int, int] = parse_python_version(data.python_version)
    except:  # noqa: E722
        return 400, OptimizeErrorResponseSchema(
            error="Invalid Python version, it should look like 3.x.x. We only support Python 3.9 and above."
        )
    try:
        ctx.validate_and_parse_source_code(code=data.source_code, feature_version=python_version[:2])
    except SyntaxError:
        return 400, OptimizeErrorResponseSchema(
            error="Invalid source code. It is not valid Python code. Please check syntax of your code."
        )
    if not validate_trace_id(data.trace_id):
        return 400, OptimizeErrorResponseSchema(error="Invalid trace ID. Please provide a valid UUIDv4.")
    optimization_response_items = await optimize_python_code_line_profiler(
        user_id=request.user,
        trace_id=data.trace_id,
        ctx=ctx,
        dependency_code=data.dependency_code,
        line_profiler_results=data.line_profiler_results,
        n=min(data.n_candidates_lp or 6, 8),
        lsp_mode=data.lsp_mode,
        python_version=python_version,
    )
    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}")
        return 500, OptimizeErrorResponseSchema(error="Error generating optimizations. Internal server error.")
    ph(
        request.user,
        "aiservice-optimize-optimizations-found",
        properties={"num_optimizations": len(optimization_response_items)},
    )

    if hasattr(request, "should_log_features") and request.should_log_features:
        await log_features(
            trace_id=data.trace_id,
            user_id=request.user,
            original_code=data.source_code,
            dependency_code=data.dependency_code,
            line_profiler_results=data.line_profiler_results,
            optimizations_raw={
                op_id: cei.code for op_id, cei in ctx.code_and_explanation_before_post_processing.items()
            },
            optimizations_post={cei.optimization_id: cei.source_code for cei in optimization_response_items},
            explanations_raw={
                op_id: cei.explanation for op_id, cei in ctx.code_and_explanation_before_post_processing.items()
            },
            explanations_post={cei.optimization_id: cei.explanation for cei in optimization_response_items},
            experiment_metadata=data.experiment_metadata if data.experiment_metadata else None,
        )

    response = OptimizeResponseSchema(optimizations=optimization_response_items)

    def log_response() -> None:
        debug_log_sensitive_data(f"Response:\n{response.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-optimize-successful")
    return 200, response