From 76cc5b19cb4fc7787bb7fb567c5bad2d86c4b34c Mon Sep 17 00:00:00 2001
From: Roland Tannous <rolandtannous@gravityq.ai>
Date: Tue, 10 Mar 2026 16:01:57 +0000
Subject: [PATCH] fix: show generated templates in UI, make system prompt
 optional
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

- System prompt is now optional — LLM only generates one when the task
  is ambiguous from the data alone (persona, domain, format constraints)
- Sanitize system_prompt extraction (handle literal "null" string)
- Show system prompt, user template, and assistant template in the
  advisor notification banner so user can see exactly what was generated
- Templates displayed in monospace with labeled sections
---
 studio/backend/utils/datasets/llm_assist.py   | 15 ++++++--
 .../dataset-preview-dialog-mapping.tsx        | 36 +++++++++++++++++--
 .../sections/dataset-preview-dialog.tsx       |  7 ++++
 3 files changed, 52 insertions(+), 6 deletions(-)

diff --git a/studio/backend/utils/datasets/llm_assist.py b/studio/backend/utils/datasets/llm_assist.py
index 34ecd5eb55..ded4c49a00 100644
--- a/studio/backend/utils/datasets/llm_assist.py
+++ b/studio/backend/utils/datasets/llm_assist.py
@@ -562,7 +562,8 @@ def _run_multi_pass_advisor(
                     {samples_text}
 
                     Design a conversion strategy to turn this into conversation format for fine-tuning.
-                    The strategy should create a system prompt, a user message template, and an assistant message template.
+                    The strategy should create a user message template and an assistant message template.
+                    Optionally include a system prompt ONLY if the task is ambiguous from the data alone.
 
                     RULES:
                     - Use {{column_name}} placeholders in templates to reference column values.
@@ -574,10 +575,13 @@ def _run_multi_pass_advisor(
                     - The assistant template should produce the expected model output.
                     - column_roles: mark columns used in the user template as "user",
                       columns used in the assistant template as "assistant".
+                    - system_prompt: set to null if the user/assistant templates alone
+                      make the task clear. Only provide one when extra context is needed
+                      (e.g. persona, domain expertise, output format constraints).
 
                     Respond with a JSON object:
                     {{
-                        "system_prompt": "<system prompt describing what the model should do>",
+                        "system_prompt": "<system prompt or null if not needed>",
                         "user_template": "<template for user message using {{column}} placeholders>",
                         "assistant_template": "<template for assistant response using {{column_name}} placeholders>",
                         "column_roles": {{
@@ -599,7 +603,12 @@ def _run_multi_pass_advisor(
             return None
 
         # ── Extract conversion strategy from Pass 2 ──
-        sys_prompt = pass2.get("system_prompt", "")
+        raw_sys = pass2.get("system_prompt")
+        # LLM may return the literal string "null" or None
+        sys_prompt = (
+            raw_sys if isinstance(raw_sys, str) and raw_sys.lower() not in ("null", "none", "")
+            else ""
+        )
         user_tpl = pass2.get("user_template", "")
         asst_tpl = pass2.get("assistant_template", "")
         label_map = pass2.get("label_mapping", {})
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
index 4e416ed1c6..153230df7d 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog-mapping.tsx
@@ -101,6 +101,9 @@ export function DatasetMappingCard({
   isAiLoading = false,
   aiError,
   advisorNotification,
+  advisorSystemPrompt,
+  advisorUserTemplate,
+  advisorAssistantTemplate,
 }: {
   mapping: Record<string, string>;
   mappingOk: boolean;
@@ -112,6 +115,9 @@ export function DatasetMappingCard({
   isAiLoading?: boolean;
   aiError?: string | null;
   advisorNotification?: string | null;
+  advisorSystemPrompt?: string;
+  advisorUserTemplate?: string;
+  advisorAssistantTemplate?: string;
 }) {
   const entries = Object.entries(mapping);
   const requiredLabel = isAudio
@@ -217,9 +223,33 @@ export function DatasetMappingCard({
             </div>
           )}
           {advisorNotification && (
-            <div className="mt-3 rounded-lg border border-indigo-200 bg-indigo-50 px-3 py-2.5 text-xs text-indigo-700 dark:border-indigo-800 dark:bg-indigo-950 dark:text-indigo-300 flex items-start gap-2">
-              <Sparkles className="size-3.5 shrink-0 mt-0.5" />
-              <span>{advisorNotification}</span>
+            <div className="mt-3 rounded-lg border border-indigo-200 bg-indigo-50 px-3 py-2.5 text-xs text-indigo-700 dark:border-indigo-800 dark:bg-indigo-950 dark:text-indigo-300 space-y-2">
+              <div className="flex items-start gap-2">
+                <Sparkles className="size-3.5 shrink-0 mt-0.5" />
+                <span>{advisorNotification}</span>
+              </div>
+              {(advisorSystemPrompt || advisorUserTemplate || advisorAssistantTemplate) && (
+                <div className="space-y-1.5 pl-5.5 text-[11px] font-mono text-indigo-600/80 dark:text-indigo-400/80">
+                  {advisorSystemPrompt && (
+                    <div>
+                      <span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">System:</span>{" "}
+                      <span className="break-words">{advisorSystemPrompt}</span>
+                    </div>
+                  )}
+                  {advisorUserTemplate && (
+                    <div>
+                      <span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">User:</span>{" "}
+                      <span className="break-words">{advisorUserTemplate}</span>
+                    </div>
+                  )}
+                  {advisorAssistantTemplate && (
+                    <div>
+                      <span className="font-sans font-medium text-indigo-500 dark:text-indigo-400">Assistant:</span>{" "}
+                      <span className="break-words">{advisorAssistantTemplate}</span>
+                    </div>
+                  )}
+                </div>
+              )}
             </div>
           )}
         </div>
diff --git a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
index fcae317b49..446f445a50 100644
--- a/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
+++ b/studio/frontend/src/features/studio/sections/dataset-preview-dialog.tsx
@@ -68,6 +68,7 @@ export function DatasetPreviewDialog({
   const {
     manualMapping, setManualMapping, datasetFormat,
     setDatasetAdvisorFields, datasetAdvisorNotification,
+    datasetSystemPrompt, datasetUserTemplate, datasetAssistantTemplate,
   } = useTrainingConfigStore(
     useShallow((s) => ({
       manualMapping: s.datasetManualMapping,
@@ -75,6 +76,9 @@ export function DatasetPreviewDialog({
       datasetFormat: s.datasetFormat,
       setDatasetAdvisorFields: s.setDatasetAdvisorFields,
       datasetAdvisorNotification: s.datasetAdvisorNotification,
+      datasetSystemPrompt: s.datasetSystemPrompt,
+      datasetUserTemplate: s.datasetUserTemplate,
+      datasetAssistantTemplate: s.datasetAssistantTemplate,
     })),
   );
   const { isStarting, startError, startTrainingRun } = useTrainingActions();
@@ -423,6 +427,9 @@ export function DatasetPreviewDialog({
                   isAiLoading={isAiLoading}
                   aiError={aiError}
                   advisorNotification={datasetAdvisorNotification}
+                  advisorSystemPrompt={datasetSystemPrompt || undefined}
+                  advisorUserTemplate={datasetUserTemplate || undefined}
+                  advisorAssistantTemplate={datasetAssistantTemplate || undefined}
                 />
               )}