From 23a9a4190b1027f024a0bfc537d0d1a5cee533a4 Mon Sep 17 00:00:00 2001
From: yuneng-jiang <yuneng.jiang@gmail.com>
Date: Mon, 23 Mar 2026 22:31:58 -0700
Subject: [PATCH] Address greptile review feedback (greploop iteration 2)

- Memoize promptsForSource with useMemo for consistency
- Extract getTestSourceOptions to pipeline_utils.tsx, removing
  direct getFrameworks import from pipeline_test_drawer.tsx

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
---
 .../src/components/policies/pipeline_test_drawer.tsx   | 10 +++-------
 .../src/components/policies/pipeline_utils.tsx         |  8 ++++++++
 2 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx
index f54b2a39056..8f7691fbb3e 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_test_drawer.tsx
@@ -4,12 +4,12 @@ import { Button } from "@tremor/react";
 import { GuardrailPipeline, PipelineStep, PipelineTestResult } from "./types";
 import { testPipelineCall } from "../networking";
 import type { CompliancePrompt } from "../../data/compliancePrompts";
-import { getFrameworks } from "../../data/compliancePrompts";
 import {
   TEST_SOURCE_QUICK,
   TEST_SOURCE_ALL,
   ACTION_LABELS,
   getPromptsForTestSource,
+  getTestSourceOptions,
   complianceMatchExpected,
 } from "./pipeline_utils";
 
@@ -62,14 +62,10 @@ export const PipelineTestPanel: React.FC<PipelineTestPanelProps> = ({
   const [error, setError] = useState<string | null>(null);
   const [complianceResults, setComplianceResults] = useState<ComplianceRunEntry[]>([]);
 
-  const testSourceOptions = useMemo(() => [
-    { value: TEST_SOURCE_QUICK, label: "Quick chat (custom message)" },
-    ...getFrameworks().map((f) => ({ value: f.name, label: f.name })),
-    { value: TEST_SOURCE_ALL, label: "All compliance datasets" },
-  ], []);
+  const testSourceOptions = useMemo(() => getTestSourceOptions(), []);
 
   const isQuickChat = testSource === TEST_SOURCE_QUICK;
-  const promptsForSource = getPromptsForTestSource(testSource);
+  const promptsForSource = useMemo(() => getPromptsForTestSource(testSource), [testSource]);
   const isDataset = promptsForSource.length > 0;
 
   const handleRunTest = async () => {
diff --git a/ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx b/ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx
index 44c59579bbc..006decf1878 100644
--- a/ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx
+++ b/ui/litellm-dashboard/src/components/policies/pipeline_utils.tsx
@@ -37,6 +37,14 @@ export function getPromptsForTestSource(source: string): CompliancePrompt[] {
   return fw ? fw.categories.flatMap((c) => c.prompts) : [];
 }
 
+export function getTestSourceOptions(): { value: string; label: string }[] {
+  return [
+    { value: TEST_SOURCE_QUICK, label: "Quick chat (custom message)" },
+    ...getFrameworks().map((f) => ({ value: f.name, label: f.name })),
+    { value: TEST_SOURCE_ALL, label: "All compliance datasets" },
+  ];
+}
+
 export function createDefaultStep(): PipelineStep {
   return {
     guardrail: "",