Merge pull request #1183 from optuna/feat/edit-graph-generated-by-llm

Introduce edit functionality to plotly graph generation feature

Author: c-bata

jupyterlab/src/apiClient.ts

@@ -9,6 +9,8 @@ import {
   ParamImportancesResponse,
   PlotResponse,
   PlotType,
+  ReGeneratePlotlyGraphQueryRequest,
+  ReGeneratePlotlyGraphQueryResponse,
   RenameStudyResponse,
   StudyDetail,
   StudyDetailResponse,
@@ -318,4 +320,11 @@ export class JupyterlabAPIClient extends APIClient {
       "Generate plotly graph query is not implemented in JupyterLab API client."
     )
   }
+  callReGeneratePlotlyGraphQuery(
+    request: ReGeneratePlotlyGraphQueryRequest
+  ): Promise<ReGeneratePlotlyGraphQueryResponse> {
+    throw new Error(
+      "Re-generate plotly graph query is not implemented in JupyterLab API client."
+    )
+  }
 }

optuna_dashboard/llm/_api_views.py

@@ -10,6 +10,7 @@
 from optuna_dashboard.llm.prompt_templates._generate_plotly_graph import (
     get_generate_plotly_graph_prompt,
     get_generate_plotly_graph_title_prompt,
+    get_re_generate_plotly_graph_prompt,
 )
 from .._bottle_util import json_api_view
 from .prompt_templates._trial_filter import get_trial_filtering_prompt
@@ -103,3 +104,42 @@ def get_generate_plotly_graph_func_str() -> dict[str, str]:
             "generate_plotly_graph_func_str": generate_plotly_graph_func_str,
             "generate_plotly_graph_title": generated_plotly_graph_title,
         }
+
+    @app.post("/api/llm/re_generate_plotly_graph_query")
+    @json_api_view
+    def get_re_generate_plotly_graph_func_str() -> dict[str, str]:
+        if llm_provider is None:
+            response.status = 400  # Bad Request
+            return {"reason": "No access to the LLM provider."}
+
+        modification_request_query = request.json.get("modification_request_query", "")
+        previous_function = request.json.get("previous_function", "")
+        if not modification_request_query or not previous_function:
+            response.status = 400  # Bad Request
+            return {"reason": "Insufficient data provided."}
+
+        func_str = request.json.get("last_response", {}).get("func_str")
+        err_msg = request.json.get("last_response", {}).get("error_message")
+
+        try:
+            prompt = get_re_generate_plotly_graph_prompt(
+                previous_function,
+                modification_request_query,
+                func_str,
+                err_msg,
+            )
+            re_generated_plotly_graph_func_str = llm_provider.call(prompt)
+        except RateLimitExceeded as e:
+            response.status = 429  # Too Many Requests
+            reason = f"Rate limit exceeded. Try again later. The actual error: {str(e)}"
+            return {"reason": reason}
+        except InvalidAuthentication as e:
+            response.status = 401  # Unauthorized
+            reason = f"Invalid authentication. Check your API key. The actual error: {str(e)}"
+            return {"reason": reason}
+        except Exception as e:
+            response.status = 500
+            return {"reason": str(e)}
+
+        response.status = 200
+        return {"re_generated_plotly_graph_func_str": re_generated_plotly_graph_func_str}

optuna_dashboard/llm/prompt_templates/_generate_plotly_graph.py

@@ -159,6 +159,44 @@
 Return ONLY the final title string.
 """  # noqa: E501
 
+_RE_GENERATE_PLOTLY_GRAPH_PROMPT_TEMPLATE = """Please rewrite a JavaScript function that, given the Study object named `study`, returns an array of Plotly trace objects (PlotData[]) to visualize the user's request. Apply the user's modification request to the previously generated function while preserving all unaffected behavior.
+
+Requirements:
+1. Input: a variable `study` of type Study (see definitions below). Do not redefine `study`.
+2. Output: an array (e.g. `[trace1, trace2, ...]`) of Plotly trace objects (each object corresponds to a Plotly `data` item). Do NOT include the layout object, only data traces.
+3. The function must be a pure function: it must not modify `study` or any trials inside it.
+4. No network requests, no DOM access, no external I/O, no eval, no Function constructor.
+5. Handle missing or incomplete trials safely: skip trials where required values are undefined (e.g. failed trials without `values`).
+6. Ensure arrays `x` and `y` are the same length for each trace. Filter out undefined entries first.
+7. Avoid extremely large custom objects; keep each trace under ~5000 points (if more, sample evenly).
+8. Return only JavaScript function code; no surrounding code fences or commentary.
+
+Study / Trial type definitions (for reference):
+{study_definition_in_typescript}
+
+====== Instructions Finished ======
+
+Context:
+
+Previously generated function:
+------
+{previous_function}
+------
+
+Your task:
+Using the context above, return a valid JavaScript function code as is (note: do not wrap your output using code blocks such as ```javascript``` ).
+At the very top of the function, include a brief comment summarizing how you applied the modification and any reasonable assumptions you had to make.
+If the modification request is ambiguous, make a reasonable assumption and note it in a code comment at the top of the function.
+
+Here is the user's modification request:
+------
+{modification_request_query}
+------
+
+====== User Query Finished ======
+{re_generate_plotly_graph_failure_message}
+"""  # noqa: E501
+
 
 def get_generate_plotly_graph_prompt(
     user_query: str, last_func_str: str | None = None, last_error_msg: str | None = None
@@ -182,3 +220,23 @@ def get_generate_plotly_graph_title_prompt(user_query: str, generated_function:
         generated_function=generated_function,
         study_definition_in_typescript=_STUDY_DEFINITION_IN_TYPESCRIPT,
     )
+
+
+def get_re_generate_plotly_graph_prompt(
+    previous_function: str,
+    modification_request_query: str,
+    last_func_str: str | None = None,
+    last_error_msg: str | None = None,
+) -> str:
+    failure_msg = ""
+    if last_func_str is not None:
+        failure_msg = GENERATE_PLOTLY_GRAPH_FAILURE_MESSAGE_TEMPLATE.format(
+            last_func_str=last_func_str,
+            error_message=last_error_msg or "No Error Message Provided.",
+        )
+    return _RE_GENERATE_PLOTLY_GRAPH_PROMPT_TEMPLATE.format(
+        previous_function=previous_function,
+        modification_request_query=modification_request_query,
+        re_generate_plotly_graph_failure_message=failure_msg,
+        study_definition_in_typescript=_STUDY_DEFINITION_IN_TYPESCRIPT,
+    )

optuna_dashboard/ts/apiClient.ts

@@ -164,6 +164,19 @@ export type GeneratePlotlyGraphQueryResponse = {
   generate_plotly_graph_title: string
 }
 
+export type ReGeneratePlotlyGraphQueryRequest = {
+  modification_request_query: string
+  previous_function: string
+  last_response?: {
+    func_str: string
+    error_message: string
+  }
+}
+
+export type ReGeneratePlotlyGraphQueryResponse = {
+  re_generated_plotly_graph_func_str: string
+}
+
 export abstract class APIClient {
   constructor() {}
 
@@ -291,4 +304,7 @@ export abstract class APIClient {
   abstract callGeneratePlotlyGraphQuery(
     request: GeneratePlotlyGraphQueryRequest
   ): Promise<GeneratePlotlyGraphQueryResponse>
+  abstract callReGeneratePlotlyGraphQuery(
+    request: ReGeneratePlotlyGraphQueryRequest
+  ): Promise<ReGeneratePlotlyGraphQueryResponse>
 }

optuna_dashboard/ts/axiosClient.ts

@@ -10,6 +10,8 @@ import {
   ParamImportancesResponse,
   PlotResponse,
   PlotType,
+  ReGeneratePlotlyGraphQueryRequest,
+  ReGeneratePlotlyGraphQueryResponse,
   RenameStudyResponse,
   StudyDetailResponse,
   StudySummariesResponse,
@@ -321,4 +323,13 @@ export class AxiosClient extends APIClient {
         request
       )
       .then<GeneratePlotlyGraphQueryResponse>((res) => res.data)
+  callReGeneratePlotlyGraphQuery = (
+    request: ReGeneratePlotlyGraphQueryRequest
+  ): Promise<ReGeneratePlotlyGraphQueryResponse> =>
+    this.axiosInstance
+      .post<ReGeneratePlotlyGraphQueryResponse>(
+        "/api/llm/re_generate_plotly_graph_query",
+        request
+      )
+      .then<ReGeneratePlotlyGraphQueryResponse>((res) => res.data)
 }