From 95f9e0ba41c9c7fbfe0bbe62d1926906a7c99044 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 13:07:54 +0100 Subject: [PATCH 1/7] feat(studio): add support for code block actions including copy and download options in markdown blocks --- .../components/assistant-ui/markdown-text.tsx | 131 +++++++++++++++++- 1 file changed, 127 insertions(+), 4 deletions(-) diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index 252faea6c3..c967b33c27 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -7,18 +7,69 @@ import { HugeiconsIcon } from "@hugeicons/react"; import { code } from "@streamdown/code"; import { math } from "@streamdown/math"; import { mermaid } from "@streamdown/mermaid"; +import { DownloadIcon } from "lucide-react"; import { Block, type BlockProps, Streamdown } from "streamdown"; import { useEffect, useRef, useState } from "react"; import "katex/dist/katex.min.css"; import { AudioPlayer } from "./audio-player"; -const { withSmoothContextProvider, useSmoothStatus } = INTERNAL; +const { withSmoothContextProvider } = INTERNAL; function getMermaidSource(blockContent: string): string | null { const source = blockContent.match(/```mermaid\s*([\s\S]*?)```/i)?.[1]?.trim(); return source && source.length > 0 ? source : null; } +function getCodeFence( + blockContent: string, +): { language: string | null; source: string } | null { + const match = blockContent + .trimEnd() + .match(/^```([^\n`]*)\n([\s\S]*?)\n?```$/); + if (!match) return null; + + return { + language: match[1]?.trim() || null, + source: match[2], + }; +} + +function getCodeFilename(language: string | null) { + const extByLanguage: Record = { + bash: "sh", + javascript: "js", + js: "js", + json: "json", + jsx: "jsx", + markdown: "md", + md: "md", + python: "py", + py: "py", + shell: "sh", + sh: "sh", + sql: "sql", + ts: "ts", + tsx: "tsx", + typescript: "ts", + yaml: "yml", + yml: "yml", + }; + + const normalized = language?.toLowerCase(); + const ext = normalized ? extByLanguage[normalized] || normalized : "txt"; + return `snippet.${ext}`; +} + +function downloadTextFile(filename: string, text: string) { + const blob = new Blob([text], { type: "text/plain;charset=utf-8" }); + const url = URL.createObjectURL(blob); + const anchor = document.createElement("a"); + anchor.href = url; + anchor.download = filename; + anchor.click(); + URL.revokeObjectURL(url); +} + const COPY_RESET_MS = 2000; function MermaidCopyButton({ source }: { source: string }) { @@ -55,9 +106,68 @@ function MermaidCopyButton({ source }: { source: string }) { ); } +function CodeBlockActions({ + disabled, + language, + source, +}: { + disabled: boolean; + language: string | null; + source: string; +}) { + const [copied, setCopied] = useState(false); + const resetTimeoutRef = useRef | null>(null); + + useEffect(() => { + return () => { + if (resetTimeoutRef.current) { + clearTimeout(resetTimeoutRef.current); + } + }; + }, []); + + return ( +
+
+ + +
+
+ ); +} + function StreamdownBlock(props: BlockProps) { const hasMermaidFence = props.content.includes("```mermaid"); const mermaidSource = getMermaidSource(props.content); + const codeFence = getCodeFence(props.content); if (props.isIncomplete && hasMermaidFence) { return ( @@ -69,20 +179,32 @@ function StreamdownBlock(props: BlockProps) { if (mermaidSource) { return ( -
+
); } + if (codeFence) { + return ( +
+ + +
+ ); + } + return ; } const AUDIO_PLAYER_RE = //; const MarkdownTextImpl = () => { - const { text } = useMessagePartText(); - const status = useSmoothStatus(); + const { text, status } = useMessagePartText(); const audioMatch = text.match(AUDIO_PLAYER_RE); if (audioMatch) { @@ -96,6 +218,7 @@ const MarkdownTextImpl = () => { isAnimating={status.type === "running"} plugins={{ code, math, mermaid }} controls={{ + code: false, mermaid: { fullscreen: true, download: true, From d66bc2760b5b2e279f30f820ac7b8445989283b2 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 14:47:20 +0100 Subject: [PATCH 2/7] feat(studio): rework chart settings with a new preferences store and revamped settings UI --- .../components/assistant-ui/markdown-text.tsx | 111 ++-- studio/frontend/src/components/ui/chart.tsx | 153 +++--- .../studio/sections/charts-content.tsx | 198 ++++--- .../charts/chart-preferences-store.ts | 93 ++++ .../sections/charts/chart-settings-sheet.tsx | 312 +++++++++++ .../sections/charts/eval-loss-chart-card.tsx | 35 +- .../sections/charts/grad-norm-chart-card.tsx | 69 +-- .../charts/learning-rate-chart-card.tsx | 55 +- .../sections/charts/shared-chart-settings.tsx | 80 --- .../charts/training-loss-chart-card.tsx | 131 ++--- .../features/studio/sections/charts/types.ts | 7 - .../features/studio/sections/charts/utils.ts | 92 +++- .../studio/sections/progress-section.tsx | 503 +++++++++++------- 13 files changed, 1130 insertions(+), 709 deletions(-) create mode 100644 studio/frontend/src/features/studio/sections/charts/chart-preferences-store.ts create mode 100644 studio/frontend/src/features/studio/sections/charts/chart-settings-sheet.tsx delete mode 100644 studio/frontend/src/features/studio/sections/charts/shared-chart-settings.tsx diff --git a/studio/frontend/src/components/assistant-ui/markdown-text.tsx b/studio/frontend/src/components/assistant-ui/markdown-text.tsx index c967b33c27..ecc272815c 100644 --- a/studio/frontend/src/components/assistant-ui/markdown-text.tsx +++ b/studio/frontend/src/components/assistant-ui/markdown-text.tsx @@ -1,32 +1,42 @@ "use client"; -import { INTERNAL, useMessagePartText } from "@assistant-ui/react"; import { copyToClipboard } from "@/lib/copy-to-clipboard"; +import { INTERNAL, useMessagePartText } from "@assistant-ui/react"; import { Copy02Icon, Tick02Icon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import { code } from "@streamdown/code"; import { math } from "@streamdown/math"; import { mermaid } from "@streamdown/mermaid"; import { DownloadIcon } from "lucide-react"; -import { Block, type BlockProps, Streamdown } from "streamdown"; import { useEffect, useRef, useState } from "react"; +import { Block, type BlockProps, Streamdown } from "streamdown"; import "katex/dist/katex.min.css"; import { AudioPlayer } from "./audio-player"; const { withSmoothContextProvider } = INTERNAL; +const COPY_RESET_MS = 2000; +const MERMAID_SOURCE_RE = /```mermaid\s*([\s\S]*?)```/i; +const CODE_FENCE_RE = /^```([^\r\n`]*)\r?\n([\s\S]*?)\r?\n?```$/; +const ACTION_PANEL_CLASS = + "pointer-events-auto flex shrink-0 items-center gap-2 rounded-md border border-sidebar bg-sidebar/80 px-1.5 py-1 supports-[backdrop-filter]:bg-sidebar/70 supports-[backdrop-filter]:backdrop-blur"; +const ACTION_BUTTON_CLASS = + "cursor-pointer p-1 text-muted-foreground transition-all hover:text-foreground disabled:cursor-not-allowed disabled:opacity-50"; + +type CodeFence = { + language: string | null; + source: string; +}; function getMermaidSource(blockContent: string): string | null { - const source = blockContent.match(/```mermaid\s*([\s\S]*?)```/i)?.[1]?.trim(); + const source = blockContent.match(MERMAID_SOURCE_RE)?.[1]?.trim(); return source && source.length > 0 ? source : null; } -function getCodeFence( - blockContent: string, -): { language: string | null; source: string } | null { - const match = blockContent - .trimEnd() - .match(/^```([^\n`]*)\n([\s\S]*?)\n?```$/); - if (!match) return null; +function getCodeFence(blockContent: string): CodeFence | null { + const match = blockContent.trimEnd().match(CODE_FENCE_RE); + if (!match) { + return null; + } return { language: match[1]?.trim() || null, @@ -56,23 +66,26 @@ function getCodeFilename(language: string | null) { }; const normalized = language?.toLowerCase(); - const ext = normalized ? extByLanguage[normalized] || normalized : "txt"; + const fallbackExt = normalized?.replace(/[^a-z0-9]+/g, "-"); + const ext = normalized + ? extByLanguage[normalized] || fallbackExt || "txt" + : "txt"; return `snippet.${ext}`; } -function downloadTextFile(filename: string, text: string) { +function downloadTextFile(filename: string, text: string): void { const blob = new Blob([text], { type: "text/plain;charset=utf-8" }); const url = URL.createObjectURL(blob); const anchor = document.createElement("a"); anchor.href = url; anchor.download = filename; + document.body.appendChild(anchor); anchor.click(); - URL.revokeObjectURL(url); + document.body.removeChild(anchor); + window.setTimeout(() => URL.revokeObjectURL(url), 0); } -const COPY_RESET_MS = 2000; - -function MermaidCopyButton({ source }: { source: string }) { +function useCopiedState() { const [copied, setCopied] = useState(false); const resetTimeoutRef = useRef | null>(null); @@ -84,24 +97,39 @@ function MermaidCopyButton({ source }: { source: string }) { }; }, []); + const showCopied = () => { + setCopied(true); + if (resetTimeoutRef.current) { + clearTimeout(resetTimeoutRef.current); + } + resetTimeoutRef.current = setTimeout(() => { + setCopied(false); + resetTimeoutRef.current = null; + }, COPY_RESET_MS); + }; + + return { copied, showCopied }; +} + +function MermaidCopyButton({ source }: { source: string }) { + const { copied, showCopied } = useCopiedState(); + return ( ); } @@ -115,42 +143,31 @@ function CodeBlockActions({ language: string | null; source: string; }) { - const [copied, setCopied] = useState(false); - const resetTimeoutRef = useRef | null>(null); - - useEffect(() => { - return () => { - if (resetTimeoutRef.current) { - clearTimeout(resetTimeoutRef.current); - } - }; - }, []); + const { copied, showCopied } = useCopiedState(); return (
-
+
+ ))} +
+ ); +} + +function SettingRow({ + label, + description, + control, +}: { + label: string; + description?: string; + control: ReactElement; +}): ReactElement { + return ( +
+
+ + {description ? ( +

{description}

+ ) : null} +
+
{control}
+
+ ); +} + +function ScaleSection({ + title, + scale, + setScale, + outlierMode, + setOutlierMode, +}: { + title: string; + scale: ScaleMode; + setScale: (value: ScaleMode) => void; + outlierMode: OutlierMode; + setOutlierMode: (value: OutlierMode) => void; +}): ReactElement { + return ( +
+
+

{title}

+

Scale and cleanup

+
+ + +
+ ); +} + +export function ChartSettingsSheet(): ReactElement { + const [open, setOpen] = useState(false); + const { + availableSteps, + windowSize, + smoothing, + showRaw, + showSmoothed, + showAvgLine, + lossScale, + lrScale, + gradScale, + lossOutlierMode, + gradOutlierMode, + lrOutlierMode, + setWindowSize, + setSmoothing, + setShowRaw, + setShowSmoothed, + setShowAvgLine, + setLossScale, + setLrScale, + setGradScale, + setLossOutlierMode, + setGradOutlierMode, + setLrOutlierMode, + resetPreferences, + } = useChartPreferencesStore( + useShallow((state) => ({ + availableSteps: state.availableSteps, + windowSize: state.windowSize, + smoothing: state.smoothing, + showRaw: state.showRaw, + showSmoothed: state.showSmoothed, + showAvgLine: state.showAvgLine, + lossScale: state.lossScale, + lrScale: state.lrScale, + gradScale: state.gradScale, + lossOutlierMode: state.lossOutlierMode, + gradOutlierMode: state.gradOutlierMode, + lrOutlierMode: state.lrOutlierMode, + setWindowSize: state.setWindowSize, + setSmoothing: state.setSmoothing, + setShowRaw: state.setShowRaw, + setShowSmoothed: state.setShowSmoothed, + setShowAvgLine: state.setShowAvgLine, + setLossScale: state.setLossScale, + setLrScale: state.setLrScale, + setGradScale: state.setGradScale, + setLossOutlierMode: state.setLossOutlierMode, + setGradOutlierMode: state.setGradOutlierMode, + setLrOutlierMode: state.setLrOutlierMode, + resetPreferences: state.resetPreferences, + })), + ); + + const minWindow = Math.min(10, Math.max(1, availableSteps)); + const effectiveWindowSize = + windowSize == null ? Math.max(availableSteps, 1) : windowSize; + const showingAll = + availableSteps > 0 && + (windowSize == null || effectiveWindowSize >= availableSteps); + const sliderMax = Math.max(minWindow, availableSteps || 1); + + return ( + <> + + + + + Chart Settings + + Tune chart presentation while training keeps running. + + +
+
+
+

View window

+

+ Show latest steps only or the full history. +

+
+
+
+ Window + + {showingAll ? "All" : effectiveWindowSize} + +
+ setWindowSize(value)} + min={minWindow} + max={sliderMax} + step={1} + disabled={availableSteps <= 1} + /> +
+
+ +
+
+

Training loss

+

+ Control overlays and EMA smoothing. +

+
+
+
+ Smoothing + {smoothing.toFixed(2)} +
+ setSmoothing(value)} + min={0} + max={0.9} + step={0.01} + /> +

+ Move right for more smoothing. `0` = raw. +

+
+ + } + /> + + } + /> + + } + /> +
+ + + + + + +
+ + + + + + + + ); +} diff --git a/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx b/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx index 844eb87b4e..a165f62d3a 100644 --- a/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx +++ b/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx @@ -11,7 +11,7 @@ import { ChartAverageIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; import { CartesianGrid, Line, LineChart, XAxis, YAxis } from "recharts"; -import { formatStepTick, placeholderEvalData } from "./utils"; +import { formatMetric, formatStepTick, placeholderEvalData } from "./utils"; const evalLossConfig = { loss: { label: "Eval Loss", color: "#ef4444" }, @@ -33,14 +33,23 @@ export function EvalLossChartCard({ return ( - 0 ? "" : " text-muted-foreground"}`}> + 0 ? "" : " text-muted-foreground"}`} + > Eval Loss {data.length > 0 ? ( - - + + Number(value).toFixed(2)} + width={80} + tickFormatter={(value) => formatMetric(Number(value))} /> `Step ${payload?.[0]?.payload?.step ?? ""}` } + formatter={(_value, _name, item) => [ + formatMetric(Number(item?.payload?.loss)), + "Eval Loss", + ]} /> } /> @@ -91,7 +104,10 @@ export function EvalLossChartCard({ ) : (
- +
- +

{isTraining && evalEnabled ? "Waiting for first evaluation step…" diff --git a/studio/frontend/src/features/studio/sections/charts/grad-norm-chart-card.tsx b/studio/frontend/src/features/studio/sections/charts/grad-norm-chart-card.tsx index 0efa57c3fd..3866837ca4 100644 --- a/studio/frontend/src/features/studio/sections/charts/grad-norm-chart-card.tsx +++ b/studio/frontend/src/features/studio/sections/charts/grad-norm-chart-card.tsx @@ -1,4 +1,4 @@ -import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { ChartContainer, ChartLegend, @@ -7,19 +7,15 @@ import { ChartTooltipContent, } from "@/components/ui/chart"; import type { ChartConfig } from "@/components/ui/chart"; -import { - DropdownMenu, - DropdownMenuContent, - DropdownMenuLabel, - DropdownMenuTrigger, -} from "@/components/ui/dropdown-menu"; -import { Settings02Icon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; import { CartesianGrid, Line, LineChart, XAxis, YAxis } from "recharts"; -import { SharedChartSettings } from "./shared-chart-settings"; -import type { OutlierMode, ScaleMode, ViewSettingsState } from "./types"; -import { CHART_SYNC_ID, formatMetric, formatStepTick, fromLog1p } from "./utils"; +import type { ScaleMode } from "./types"; +import { + CHART_SYNC_ID, + formatMetric, + formatStepTick, + fromLog1p, +} from "./utils"; const gradNormConfig = { displayGradNorm: { label: "Grad Norm", color: "#f97316" }, @@ -37,50 +33,25 @@ export function GradNormChartCard({ visibleStepDomain, xAxisTicks, scale, - setScale, - outlierMode, - setOutlierMode, - viewSettings, }: { data: GradNormPoint[]; domain: [number, number]; visibleStepDomain: [number, number]; xAxisTicks: number[]; scale: ScaleMode; - setScale: (value: ScaleMode) => void; - outlierMode: OutlierMode; - setOutlierMode: (value: OutlierMode) => void; - viewSettings: ViewSettingsState; }): ReactElement { + const showPoint = data.length <= 1 ? { r: 3, strokeWidth: 0 } : false; + return ( - Gradient Norm - - - - - - - Chart Settings - - - - + Gradient Norm - + { const num = Number(value); - if (!Number.isFinite(num)) return "0"; + if (!Number.isFinite(num)) { + return "0"; + } const shown = scale === "log" ? fromLog1p(num) : num; return formatMetric(shown); }} @@ -133,11 +106,11 @@ export function GradNormChartCard({ } /> void; - outlierMode: OutlierMode; - setOutlierMode: (value: OutlierMode) => void; - viewSettings: ViewSettingsState; }): ReactElement { + const showPoint = data.length <= 1 ? { r: 3, strokeWidth: 0 } : false; + return ( - Learning Rate - - - - - - - Chart Settings - - - - + Learning Rate @@ -114,7 +77,9 @@ export function LearningRateChartCard({ width={52} tickFormatter={(value) => { const num = Number(value); - if (!Number.isFinite(num)) return "0e+0"; + if (!Number.isFinite(num)) { + return "0e+0"; + } const shown = scale === "log" ? fromLog1p(num) : num; return shown.toExponential(0); }} @@ -136,11 +101,11 @@ export function LearningRateChartCard({ } /> void; - outlierMode: OutlierMode; - setOutlierMode: (value: OutlierMode) => void; -}): ReactElement { - const showingAll = view.allStepsLength > 0 && view.effectiveWindowSize >= view.allStepsLength; - - return ( - <> - - View -

-
- - - {showingAll ? "All" : view.effectiveWindowSize} - -
- view.setWindowSize(Math.max(1, Math.round(v)))} - min={view.minWindow} - max={Math.max(view.minWindow, view.allStepsLength)} - step={1} - /> - - Always follows latest steps - -
- - Y Scale - checked && setScale("linear")} - > - Linear - - checked && setScale("log")} - > - Log (log1p) - - - Outliers - checked && setOutlierMode("none")} - > - No clipping - - checked && setOutlierMode("p99")} - > - Clip above p99 - - checked && setOutlierMode("p95")} - > - Clip above p95 - - - ); -} diff --git a/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx b/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx index bc548d0c11..1edf2bad55 100644 --- a/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx +++ b/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx @@ -1,4 +1,4 @@ -import { Card, CardAction, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; +import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; import { ChartContainer, ChartLegend, @@ -7,23 +7,22 @@ import { ChartTooltipContent, } from "@/components/ui/chart"; import type { ChartConfig } from "@/components/ui/chart"; -import { - DropdownMenu, - DropdownMenuCheckboxItem, - DropdownMenuContent, - DropdownMenuLabel, - DropdownMenuSeparator, - DropdownMenuTrigger, -} from "@/components/ui/dropdown-menu"; -import { Label } from "@/components/ui/label"; -import { Slider } from "@/components/ui/slider"; -import { Settings02Icon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; -import { CartesianGrid, Line, LineChart, ReferenceLine, XAxis, YAxis } from "recharts"; -import { SharedChartSettings } from "./shared-chart-settings"; -import type { OutlierMode, ScaleMode, ViewSettingsState } from "./types"; -import { CHART_SYNC_ID, formatMetric, formatStepTick, fromLog1p } from "./utils"; +import { + CartesianGrid, + Line, + LineChart, + ReferenceLine, + XAxis, + YAxis, +} from "recharts"; +import type { ScaleMode } from "./types"; +import { + CHART_SYNC_ID, + formatMetric, + formatStepTick, + fromLog1p, +} from "./utils"; const lossConfig = { displayLoss: { label: "Loss", color: "#3b82f6" }, @@ -45,19 +44,10 @@ export function TrainingLossChartCard({ xAxisTicks, avgRaw, avgDisplay, - smoothing, - setSmoothing, showRaw, - setShowRaw, showSmoothed, - setShowSmoothed, showAvgLine, - setShowAvgLine, - viewSettings, scale, - setScale, - outlierMode, - setOutlierMode, }: { data: LossChartPoint[]; domain: [number, number]; @@ -65,81 +55,17 @@ export function TrainingLossChartCard({ xAxisTicks: number[]; avgRaw: number; avgDisplay: number; - smoothing: number; - setSmoothing: (value: number) => void; showRaw: boolean; - setShowRaw: (value: boolean) => void; showSmoothed: boolean; - setShowSmoothed: (value: boolean) => void; showAvgLine: boolean; - setShowAvgLine: (value: boolean) => void; - viewSettings: ViewSettingsState; scale: ScaleMode; - setScale: (value: ScaleMode) => void; - outlierMode: OutlierMode; - setOutlierMode: (value: OutlierMode) => void; }): ReactElement { + const showPoint = data.length <= 1 ? { r: 3, strokeWidth: 0 } : false; + return ( - Training Loss - - - - - - - Chart Settings - -
-
- - - {smoothing.toFixed(2)} - -
- setSmoothing(v)} - min={0} - max={0.99} - step={0.01} - /> -
- - setShowRaw(Boolean(value))} - > - Show raw loss - - setShowSmoothed(Boolean(value))} - > - Show smoothed loss - - setShowAvgLine(Boolean(value))} - > - Show average line - - - - - + Training Loss @@ -173,10 +99,12 @@ export function TrainingLossChartCard({ axisLine={false} tickMargin={4} fontSize={10} - width={52} + width={80} tickFormatter={(value) => { const num = Number(value); - if (!Number.isFinite(num)) return "0"; + if (!Number.isFinite(num)) { + return "0"; + } const shown = scale === "log" ? fromLog1p(num) : num; return formatMetric(shown); }} @@ -189,7 +117,10 @@ export function TrainingLossChartCard({ } formatter={(_value, name, item) => { if (name === "displaySmoothed") { - return [formatMetric(Number(item?.payload?.smoothed)), "Smoothed"]; + return [ + formatMetric(Number(item?.payload?.smoothed)), + "Smoothed", + ]; } return [formatMetric(Number(item?.payload?.loss)), "Loss"]; }} @@ -212,12 +143,12 @@ export function TrainingLossChartCard({ )} {showRaw && ( void; -} diff --git a/studio/frontend/src/features/studio/sections/charts/utils.ts b/studio/frontend/src/features/studio/sections/charts/utils.ts index 482ca00f76..db5b53bd09 100644 --- a/studio/frontend/src/features/studio/sections/charts/utils.ts +++ b/studio/frontend/src/features/studio/sections/charts/utils.ts @@ -3,6 +3,8 @@ import type { LossHistoryItem, OutlierMode, SmoothedLossItem } from "./types"; export const CHART_SYNC_ID = "train-metrics-sync"; export const MAX_RENDER_POINTS = 800; export const DEFAULT_VISIBLE_POINTS = 160; +const TRAILING_ZEROES_RE = /\.?0+$/; +const NEGATIVE_ZERO_RE = /^-0$/; export const placeholderEvalData = [ { step: 0, loss: 2.8 }, @@ -22,11 +24,30 @@ export function fromLog1p(value: number): number { } export function formatMetric(value: number): string { - if (!Number.isFinite(value)) return "0"; - if (value === 0) return "0"; - if (value >= 1000) return value.toFixed(0); - if (value >= 1) return value.toFixed(2); - return value.toExponential(2); + if (!Number.isFinite(value)) { + return "0"; + } + const abs = Math.abs(value); + let decimals = 6; + + if (abs >= 1000) { + decimals = 0; + } else if (abs >= 100) { + decimals = 2; + } else if (abs >= 1) { + decimals = 4; + } else if (abs >= 0.01) { + decimals = 5; + } else if (abs >= 0.0001) { + decimals = 6; + } else { + decimals = 8; + } + + return value + .toFixed(decimals) + .replace(TRAILING_ZEROES_RE, "") + .replace(NEGATIVE_ZERO_RE, "0"); } export function formatStepTick(value: number): string { @@ -54,18 +75,12 @@ export function clamp(value: number, min: number, max: number): number { return Math.min(max, Math.max(min, value)); } -export function getDefaultWindowSize(totalSteps: number): number { - if (totalSteps <= 1) { - return Math.max(totalSteps, 1); - } - if (totalSteps <= DEFAULT_VISIBLE_POINTS) { - return clamp(Math.floor(totalSteps * 0.6), 1, totalSteps); - } - return DEFAULT_VISIBLE_POINTS; -} - -export function buildStepTicks(min: number, max: number, targetCount = 6): number[] { - if (!Number.isFinite(min) || !Number.isFinite(max)) { +export function buildStepTicks( + min: number, + max: number, + targetCount = 6, +): number[] { + if (!(Number.isFinite(min) && Number.isFinite(max))) { return [0, 1]; } if (max <= min) { @@ -104,10 +119,17 @@ export function buildYDomain(values: number[]): [number, number] { return [min - pad, max + pad]; } -function getUpperPercentile(values: number[], mode: OutlierMode): number | null { - if (mode === "none") return null; +function getUpperPercentile( + values: number[], + mode: OutlierMode, +): number | null { + if (mode === "none") { + return null; + } const finiteValues = values.filter((value) => Number.isFinite(value)); - if (finiteValues.length < 3) return null; + if (finiteValues.length < 3) { + return null; + } const sorted = [...finiteValues].sort((a, b) => a - b); const q = mode === "p99" ? 0.99 : 0.95; @@ -120,18 +142,36 @@ function getUpperPercentile(values: number[], mode: OutlierMode): number | null export function applyOutlierCap(values: number[], mode: OutlierMode): number[] { const cap = getUpperPercentile(values, mode); - if (cap == null) return values; + if (cap == null) { + return values; + } return values.map((value) => Math.min(value, cap)); } -export function ema(data: LossHistoryItem[], alpha: number): SmoothedLossItem[] { +export function ema( + data: LossHistoryItem[], + alpha: number, +): SmoothedLossItem[] { if (data.length === 0) { return []; } - let s = data[0].loss; - return data.map((d) => { - s = alpha * d.loss + (1 - alpha) * s; - return { ...d, smoothed: +s.toFixed(4) }; + const values = data.map((point) => point.loss); + const isConstant = values.every((value) => value === values[0]); + + let last = 0; + let count = 0; + + return data.map((point) => { + const next = point.loss; + if (!Number.isFinite(next) || isConstant) { + return { ...point, smoothed: next }; + } + + last = last * alpha + (1 - alpha) * next; + count += 1; + + const debias = alpha === 1 ? 1 : 1 - alpha ** count; + return { ...point, smoothed: last / debias }; }); } diff --git a/studio/frontend/src/features/studio/sections/progress-section.tsx b/studio/frontend/src/features/studio/sections/progress-section.tsx index 10a547b82d..2125c95164 100644 --- a/studio/frontend/src/features/studio/sections/progress-section.tsx +++ b/studio/frontend/src/features/studio/sections/progress-section.tsx @@ -15,11 +15,16 @@ import { PopoverContent, PopoverTrigger, } from "@/components/ui/popover"; +import { Progress } from "@/components/ui/progress"; +import { OPTIMIZER_OPTIONS } from "@/config/training"; +import { setTrainingCompareHandoff } from "@/features/chat"; import { - useTrainingConfigStore, useTrainingActions, + useTrainingConfigStore, useTrainingRuntimeStore, } from "@/features/training"; +import { useGpuUtilization } from "@/hooks"; +import { cn } from "@/lib/utils"; import { ChartAverageIcon, DashboardSpeed01Icon, @@ -30,13 +35,28 @@ import { ZapIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { useEffect, useState, type ReactElement, type ReactNode } from "react"; import { Link, useNavigate } from "@tanstack/react-router"; +import { type ReactElement, type ReactNode, useEffect, useState } from "react"; import { useShallow } from "zustand/react/shallow"; -import { useGpuUtilization } from "@/hooks"; -import { setTrainingCompareHandoff } from "@/features/chat"; -import { OPTIMIZER_OPTIONS } from "@/config/training"; -import { formatDuration, formatNumber, phaseColors, phaseLabel } from "./progress-section-lib"; +import { ChartSettingsSheet } from "./charts/chart-settings-sheet"; +import { + formatDuration, + formatNumber, + phaseColors, + phaseLabel, +} from "./progress-section-lib"; + +type ConfigGroup = { + section: string; + rows: [string, string | number | null | undefined][]; +}; + +function configRow( + label: string, + value: string | number | null | undefined, +): [string, string | number | null | undefined] { + return [label, value]; +} export function ProgressSection(): ReactElement { const navigate = useNavigate(); @@ -94,12 +114,12 @@ export function ProgressSection(): ReactElement { const pct = runtime.totalSteps > 0 ? Math.min( - 100, - Math.max( - 0, - Math.round((runtime.currentStep / runtime.totalSteps) * 100), - ), - ) + 100, + Math.max( + 0, + Math.round((runtime.currentStep / runtime.totalSteps) * 100), + ), + ) : Math.round(runtime.progressPercent); const elapsed = runtime.elapsedSeconds; @@ -110,15 +130,26 @@ export function ProgressSection(): ReactElement { const eta = runtime.etaSeconds ?? derivedEta; const stepsPerSecond = - elapsed != null && elapsed > 0 - ? runtime.currentStep / elapsed - : null; + elapsed != null && elapsed > 0 ? runtime.currentStep / elapsed : null; const showHalfwayHint = runtime.phase === "training" && pct >= 50 && pct < 100; const showCompletedHint = runtime.phase === "completed"; - const handleCompareInChat = () => { + const handleCompareInChat = async () => { setTrainingCompareHandoff(config.selectedModel); - void navigate({ to: "/chat" }); + await navigate({ to: "/chat" }); + }; + const requestStop = async (saveCheckpoint: boolean) => { + setStopRequested(true); + setStopDialogOpen(false); + useTrainingRuntimeStore.getState().setStopRequested(true); + try { + const ok = await stopTrainingRun(saveCheckpoint); + if (!ok) { + setStopRequested(false); + } + } catch { + setStopRequested(false); + } }; const stoppedLoss = getDisplayMetric( @@ -133,37 +164,37 @@ export function ProgressSection(): ReactElement { ); const stoppedGradNorm = runtime.isTrainingRunning ? runtime.currentGradNorm - : lastNonZeroValue(runtime.gradNormHistory) ?? runtime.currentGradNorm; + : (lastNonZeroValue(runtime.gradNormHistory) ?? runtime.currentGradNorm); const optimizerLabel = OPTIMIZER_OPTIONS.find((o) => o.value === config.optimizerType)?.label ?? config.optimizerType; - const configItems = [ + const configItems: ConfigGroup[] = [ { section: "Hyperparams", rows: [ - ["Epochs", config.epochs], - ["Batch size", config.batchSize], - ["Learning rate", config.learningRate], - ["Optimizer", optimizerLabel], - ["Max steps", config.maxSteps], - ["Context length", config.contextLength], - ["Warmup steps", config.warmupSteps], + configRow("Epochs", config.epochs), + configRow("Batch size", config.batchSize), + configRow("Learning rate", config.learningRate), + configRow("Optimizer", optimizerLabel), + configRow("Max steps", config.maxSteps), + configRow("Context length", config.contextLength), + configRow("Warmup steps", config.warmupSteps), ], }, ...(config.trainingMethod !== "full" ? [ - { - section: "LoRA", - rows: [ - ["Rank", config.loraRank], - ["Alpha", config.loraAlpha], - ["Dropout", config.loraDropout], - ["Variant", config.loraVariant], - ], - }, - ] + { + section: "LoRA", + rows: [ + configRow("Rank", config.loraRank), + configRow("Alpha", config.loraAlpha), + configRow("Dropout", config.loraDropout), + configRow("Variant", config.loraVariant), + ], + }, + ] : []), ]; @@ -173,182 +204,76 @@ export function ProgressSection(): ReactElement { title="Training Progress" description={runtime.message || "Live training metrics"} accent="emerald" - className="shadow-border ring-1 ring-border" + className="shadow-border border border-border/60 bg-card/90 ring-0 backdrop-blur-sm" headerAction={ -
- - - - - -
-

Training Config

- {configItems.map((group) => ( -
-

- {group.section} -

- {group.rows.map(([label, value]) => ( -
- - {String(label)} - - - {String(value)} - -
- ))} -
- ))} -
- - - - - - - Stop Training - - Choose how you want to stop the current training run. - - - - Continue Training - { - setStopRequested(true); - setStopDialogOpen(false); - useTrainingRuntimeStore.getState().setStopRequested(true); - void stopTrainingRun(false).then((ok) => { - if (!ok) setStopRequested(false); - }); - }} - > - Cancel Training - - { - setStopRequested(true); - setStopDialogOpen(false); - useTrainingRuntimeStore.getState().setStopRequested(true); - void stopTrainingRun(true).then((ok) => { - if (!ok) setStopRequested(false); - }); - }} - > - Stop and Save - - - - -
+ } > -
+
-
+
{phaseLabel[runtime.phase]} Epoch {runtime.currentEpoch.toFixed(2)} + + {pct}% complete +
-
+
Step {runtime.currentStep} / {runtime.totalSteps || "--"} {pct}%
-
-
-
+
- {(showHalfwayHint || showCompletedHint) && ( -
-

- {showCompletedHint - ? "Training done. Next step: compare base vs fine-tuned outputs." - : "Halfway done. Training is past 50%."} -

- {showCompletedHint && ( -
- - -
- )} -
- )} + {runtime.error && ( -

{runtime.error}

+

+ {runtime.error} +

)} -
-
-

Loss

-

- {stoppedLoss.toFixed(4)} -

-
-
-

LR

-

- {stoppedLr.toExponential(2)} -

-
-
-

Grad Norm

-

- {formatNumber(stoppedGradNorm, 3)} -

-
-
-

Model

-

- {config.selectedModel ?? "--"} -

-
-
-

Method

-

- {config.trainingMethod.toUpperCase()} -

-
+
+ + {stoppedLoss.toFixed(4)} + + {stoppedLr.toExponential(2)} + + {formatNumber(stoppedGradNorm, 3)} + + + {config.selectedModel ?? "--"} + + + {config.trainingMethod.toUpperCase()} +
-
+
Elapsed: {formatDuration(elapsed)} ETA: {formatDuration(eta)} @@ -363,8 +288,13 @@ export function ProgressSection(): ReactElement {
-

GPU Monitor

-
+
+

+ GPU Monitor +

+ Live +
+
} - value={gpu.gpu_utilization_pct != null ? `${gpu.gpu_utilization_pct}%` : "--"} + value={ + gpu.gpu_utilization_pct != null + ? `${gpu.gpu_utilization_pct}%` + : "--" + } pct={gpu.gpu_utilization_pct ?? 0} /> } - value={gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--"} + icon={ + + } + value={ + gpu.temperature_c != null ? `${gpu.temperature_c}°C` : "--" + } pct={gpu.temperature_c ?? 0} max={100} /> } - value={gpu.vram_used_gb != null && gpu.vram_total_gb != null ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB` : "--"} + value={ + gpu.vram_used_gb != null && gpu.vram_total_gb != null + ? `${gpu.vram_used_gb} / ${gpu.vram_total_gb} GB` + : "--" + } pct={gpu.vram_utilization_pct ?? 0} /> } - value={gpu.power_draw_w != null ? (gpu.power_limit_w != null ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W` : `${gpu.power_draw_w} W`) : "--"} + value={ + gpu.power_draw_w != null + ? gpu.power_limit_w != null + ? `${gpu.power_draw_w} / ${gpu.power_limit_w} W` + : `${gpu.power_draw_w} W` + : "--" + } pct={gpu.power_utilization_pct ?? 0} />
@@ -402,6 +350,171 @@ export function ProgressSection(): ReactElement { ); } +function TrainingHeaderActions({ + configItems, + isTrainingRunning, + onOpenStopDialog, + onRequestStop, + stopDialogOpen, + stopRequested, +}: { + configItems: ConfigGroup[]; + isTrainingRunning: boolean; + onOpenStopDialog: (open: boolean) => void; + onRequestStop: (saveCheckpoint: boolean) => Promise; + stopDialogOpen: boolean; + stopRequested: boolean; +}): ReactElement { + return ( +
+ + + + + +
+

Training Config

+ {configItems.map((group) => ( +
+

+ {group.section} +

+ {group.rows.map(([label, value]) => ( +
+ {label} + + {String(value)} + +
+ ))} +
+ ))} +
+ + + + + + + + Stop Training + + Choose how you want to stop the current training run. + + + + Continue Training + onRequestStop(false)} + > + Cancel Training + + onRequestStop(true)}> + Stop and Save + + + + +
+ ); +} + +function MilestoneCallout({ + showCompletedHint, + showHalfwayHint, + onCompareInChat, +}: { + showCompletedHint: boolean; + showHalfwayHint: boolean; + onCompareInChat: () => Promise; +}): ReactElement | null { + if (!(showHalfwayHint || showCompletedHint)) { + return null; + } + + return ( +
+
+
+ {!showCompletedHint && ( +

+ Milestone +

+ )} +

+ {showCompletedHint + ? "Training done. Next step: compare base vs fine-tuned outputs." + : "Halfway done. Training is past 50%."} +

+
+ {!showCompletedHint && ( + + 50%+ + + )} +
+ {showCompletedHint && ( +
+ + +
+ )} +
+ ); +} + +function MetricStat({ + label, + children, + valueClassName, +}: { + label: string; + children: ReactNode; + valueClassName?: string; +}): ReactElement { + return ( +
+

{label}

+

+ {children} +

+
+ ); +} + function lastNonZeroValue(points: { value: number }[]): number | null { for (let i = points.length - 1; i >= 0; i -= 1) { const value = points[i]?.value; @@ -436,7 +549,7 @@ function GpuStat({ pct: number; max?: number; }): ReactElement { - const clamped = Math.min(pct, max ?? 100); + const clamped = Math.max(0, Math.min(pct, max ?? 100)); let barColor = "bg-red-500"; if (clamped < 60) { barColor = "bg-emerald-500"; @@ -445,7 +558,7 @@ function GpuStat({ } return ( -
+
{icon} @@ -453,7 +566,7 @@ function GpuStat({ {value}
-
+
Date: Mon, 9 Mar 2026 15:34:15 +0100 Subject: [PATCH 3/7] feat(studio): centralize chart styling and formatting --- .../sections/charts/eval-loss-chart-card.tsx | 37 +++++++++++-------- .../sections/charts/grad-norm-chart-card.tsx | 20 +++++----- .../charts/learning-rate-chart-card.tsx | 20 +++++++--- .../charts/training-loss-chart-card.tsx | 17 ++++++--- .../features/studio/sections/charts/utils.ts | 29 +++++++++++++++ 5 files changed, 86 insertions(+), 37 deletions(-) diff --git a/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx b/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx index a165f62d3a..2a1e837cf3 100644 --- a/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx +++ b/studio/frontend/src/features/studio/sections/charts/eval-loss-chart-card.tsx @@ -11,7 +11,15 @@ import { ChartAverageIcon } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; import type { ReactElement } from "react"; import { CartesianGrid, Line, LineChart, XAxis, YAxis } from "recharts"; -import { formatMetric, formatStepTick, placeholderEvalData } from "./utils"; +import { + CHART_CONTAINER_CLASS, + DEFAULT_CHART_MARGIN, + DEFAULT_Y_AXIS_WIDTH, + formatAxisMetric, + formatMetric, + formatStepTick, + placeholderEvalData, +} from "./utils"; const evalLossConfig = { loss: { label: "Eval Loss", color: "#ef4444" }, @@ -33,22 +41,17 @@ export function EvalLossChartCard({ return ( - 0 ? "" : " text-muted-foreground"}`} - > + 0 ? "" : " text-muted-foreground"}`}> Eval Loss {data.length > 0 ? ( - + formatMetric(Number(value))} + width={DEFAULT_Y_AXIS_WIDTH} + tickFormatter={(value) => formatAxisMetric(Number(value))} /> - Gradient Norm + Gradient Norm - + { const num = Number(value); if (!Number.isFinite(num)) { return "0"; } const shown = scale === "log" ? fromLog1p(num) : num; - return formatMetric(shown); + return formatAxisMetric(shown); }} /> - Learning Rate + Learning Rate - + { const num = Number(value); if (!Number.isFinite(num)) { diff --git a/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx b/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx index 1edf2bad55..7004108e1d 100644 --- a/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx +++ b/studio/frontend/src/features/studio/sections/charts/training-loss-chart-card.tsx @@ -19,6 +19,10 @@ import { import type { ScaleMode } from "./types"; import { CHART_SYNC_ID, + CHART_CONTAINER_CLASS, + DEFAULT_CHART_MARGIN, + DEFAULT_Y_AXIS_WIDTH, + formatAxisMetric, formatMetric, formatStepTick, fromLog1p, @@ -65,16 +69,16 @@ export function TrainingLossChartCard({ return ( - Training Loss + Training Loss - + { const num = Number(value); if (!Number.isFinite(num)) { return "0"; } const shown = scale === "log" ? fromLog1p(num) : num; - return formatMetric(shown); + return formatAxisMetric(shown); }} /> = 1000) { + decimals = 0; + } else if (abs >= 100) { + decimals = 1; + } else if (abs >= 1) { + decimals = 3; + } else if (abs >= 0.01) { + decimals = 4; + } else { + decimals = 5; + } + + return value + .toFixed(decimals) + .replace(TRAILING_ZEROES_RE, "") + .replace(NEGATIVE_ZERO_RE, "0"); +} + export function formatStepTick(value: number): string { if (value >= 1_000_000) { return `${(value / 1_000_000).toFixed(1)}M`; From c67f4ba29fb767ae86c0fd3bff846caba99f22f7 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 16:04:46 +0100 Subject: [PATCH 4/7] feat(recipe-studio): improve UI responsiveness and fix JSON preview handling --- studio/backend/routes/data_recipe/seed.py | 4 ++-- .../features/recipe-studio/dialogs/config-dialog.tsx | 8 +++++--- .../recipe-studio/dialogs/seed/seed-dialog.tsx | 10 +++++----- 3 files changed, 12 insertions(+), 10 deletions(-) diff --git a/studio/backend/routes/data_recipe/seed.py b/studio/backend/routes/data_recipe/seed.py index ce6073d902..f0396939ca 100644 --- a/studio/backend/routes/data_recipe/seed.py +++ b/studio/backend/routes/data_recipe/seed.py @@ -171,9 +171,9 @@ def _read_preview_rows_from_local_file(path: Path, preview_size: int) -> list[di df = pd.read_json(path, lines=True).head(preview_size) elif ext == ".json": try: - df = pd.read_json(path, lines=True).head(preview_size) - except ValueError: df = pd.read_json(path).head(preview_size) + except ValueError: + df = pd.read_json(path, lines=True).head(preview_size) else: raise HTTPException(status_code=422, detail=f"unsupported file type: {ext}") except HTTPException: diff --git a/studio/frontend/src/features/recipe-studio/dialogs/config-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/config-dialog.tsx index 10c7c6b073..76d0af92b6 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/config-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/config-dialog.tsx @@ -49,7 +49,7 @@ export function ConfigDialog({ position="absolute" overlayPosition="absolute" overlayClassName="bg-transparent" - className="corner-squircle max-h-[650px] overflow-auto sm:max-w-2xl shadow-border" + className="corner-squircle max-h-[650px] overflow-y-auto overflow-x-hidden sm:max-w-2xl shadow-border" > )} {config && ( -
+
{readOnly && (
Recipe locked while execution is active.
)} -
+
{showDropToggle && (
diff --git a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx index 0e2b744b26..67d005dab3 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/seed/seed-dialog.tsx @@ -423,13 +423,13 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl ); return ( - + Config Preview - +
{mode === "hf" && ( <> @@ -774,7 +774,7 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
- +
{previewRows.length === 0 ? (
@@ -795,8 +795,8 @@ export function SeedDialog({ config, onUpdate, open }: SeedDialogProps): ReactEl
Loaded columns: {previewColumns.join(", ") || "None"}
-
- +
+
{previewColumns.map((col) => ( From 1f37b76b19dc714af5cbda8b53d214ed232c859d Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 17:04:15 +0100 Subject: [PATCH 5/7] feat(recipe-studio): remove MCP tools-related dialogs and refactor tool profile management logic --- studio/backend/models/data_recipe.py | 16 + studio/backend/routes/data_recipe/__init__.py | 2 + studio/backend/routes/data_recipe/mcp.py | 77 ++ .../src/features/recipe-studio/api/index.ts | 27 +- .../recipe-studio/blocks/definitions.ts | 20 +- .../recipe-studio/blocks/render-dialog.tsx | 7 + .../recipe-studio/components/block-sheet.tsx | 6 + .../components/inline/inline-llm.tsx | 50 ++ .../components/inline/inline-policy.ts | 3 + .../components/recipe-graph-node.tsx | 47 ++ .../recipe-studio/dialogs/config-dialog.tsx | 3 + .../recipe-studio/dialogs/llm/general-tab.tsx | 54 +- .../recipe-studio/dialogs/llm/llm-dialog.tsx | 29 +- .../dialogs/llm/mcp-tools-tab.tsx | 297 ------- .../llm/mcp-tools/mcp-providers-section.tsx | 272 ------ .../llm/mcp-tools/tool-configs-section.tsx | 182 ---- .../mcp-tools => tool-profile}/helpers.ts | 72 +- .../tool-profile/tool-profile-dialog.tsx | 776 ++++++++++++++++++ .../hooks/use-recipe-editor-graph.ts | 13 + .../hooks/use-recipe-runtime-visuals.ts | 4 + .../recipe-studio/recipe-studio-page.tsx | 14 +- .../recipe-studio/stores/helpers/edge-sync.ts | 35 + .../stores/helpers/model-infra-layout.ts | 51 +- .../stores/helpers/reference-sync.ts | 8 + .../recipe-studio/stores/recipe-studio.ts | 47 +- .../src/features/recipe-studio/types/index.ts | 27 +- .../recipe-studio/utils/config-factories.ts | 24 +- .../utils/graph/recipe-graph-connection.ts | 20 +- .../recipe-studio/utils/graph/relations.ts | 3 + .../utils/graph/runtime-visual-state.ts | 1 + .../recipe-studio/utils/import/edges.ts | 6 + .../recipe-studio/utils/import/importer.ts | 70 +- .../src/features/recipe-studio/utils/index.ts | 1 + .../features/recipe-studio/utils/node-data.ts | 11 + .../utils/payload/build-payload.ts | 55 +- .../utils/payload/builders-llm.ts | 44 +- .../recipe-studio/utils/payload/builders.ts | 7 +- .../recipe-studio/utils/recipe-studio-view.ts | 7 + .../recipe-studio/utils/validation.ts | 38 + .../features/recipe-studio/utils/variables.ts | 6 +- 40 files changed, 1537 insertions(+), 895 deletions(-) create mode 100644 studio/backend/routes/data_recipe/mcp.py delete mode 100644 studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools-tab.tsx delete mode 100644 studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/mcp-providers-section.tsx delete mode 100644 studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/tool-configs-section.tsx rename studio/frontend/src/features/recipe-studio/dialogs/{llm/mcp-tools => tool-profile}/helpers.ts (72%) create mode 100644 studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx diff --git a/studio/backend/models/data_recipe.py b/studio/backend/models/data_recipe.py index 01e7c8cd1f..418ef97afc 100644 --- a/studio/backend/models/data_recipe.py +++ b/studio/backend/models/data_recipe.py @@ -61,3 +61,19 @@ class SeedInspectResponse(BaseModel): preview_rows: list[dict[str, Any]] = Field(default_factory=list) split: str | None = None subset: str | None = None + + +class McpToolsListRequest(BaseModel): + mcp_providers: list[dict[str, Any]] = Field(default_factory=list) + timeout_sec: float | None = Field(default=None, gt=0) + + +class McpToolsProviderResult(BaseModel): + name: str + tools: list[str] = Field(default_factory=list) + error: str | None = None + + +class McpToolsListResponse(BaseModel): + providers: list[McpToolsProviderResult] = Field(default_factory=list) + duplicate_tools: dict[str, list[str]] = Field(default_factory=dict) diff --git a/studio/backend/routes/data_recipe/__init__.py b/studio/backend/routes/data_recipe/__init__.py index 0d3f4febf9..65e7db0279 100644 --- a/studio/backend/routes/data_recipe/__init__.py +++ b/studio/backend/routes/data_recipe/__init__.py @@ -14,6 +14,7 @@ if str(backend_path) not in sys.path: sys.path.insert(0, str(backend_path)) from .jobs import router as jobs_router +from .mcp import router as mcp_router from .seed import router as seed_router from .validate import router as validate_router @@ -21,5 +22,6 @@ router = APIRouter(dependencies=[Depends(get_current_subject)]) router.include_router(seed_router) router.include_router(validate_router) router.include_router(jobs_router) +router.include_router(mcp_router) __all__ = ["router"] diff --git a/studio/backend/routes/data_recipe/mcp.py b/studio/backend/routes/data_recipe/mcp.py new file mode 100644 index 0000000000..5aa5c3247c --- /dev/null +++ b/studio/backend/routes/data_recipe/mcp.py @@ -0,0 +1,77 @@ +"""MCP helper endpoints for data recipe.""" + +from __future__ import annotations + +from collections import defaultdict + +from fastapi import APIRouter + +from core.data_recipe.service import build_mcp_providers +from models.data_recipe import ( + McpToolsListRequest, + McpToolsListResponse, + McpToolsProviderResult, +) + +router = APIRouter() + + +@router.post("/mcp/tools", response_model=McpToolsListResponse) +def list_mcp_tools(payload: McpToolsListRequest) -> McpToolsListResponse: + try: + from data_designer.engine.mcp import io as mcp_io + except ImportError as exc: + return McpToolsListResponse( + providers=[ + McpToolsProviderResult( + name="", + error=f"MCP dependencies unavailable: {exc}", + ) + ] + ) + + providers: list[McpToolsProviderResult] = [] + tool_to_providers: dict[str, list[str]] = defaultdict(list) + + for provider_payload in payload.mcp_providers: + provider_name = str(provider_payload.get("name", "")).strip() + built = build_mcp_providers({"mcp_providers": [provider_payload]}) + if len(built) != 1: + providers.append( + McpToolsProviderResult( + name=provider_name, + error="Unsupported MCP provider config.", + ) + ) + continue + + provider = built[0] + try: + tools = mcp_io.list_tools(provider, timeout_sec=payload.timeout_sec) + tool_names = sorted({tool.name for tool in tools if getattr(tool, "name", "")}) + for tool_name in tool_names: + tool_to_providers[tool_name].append(provider.name) + providers.append( + McpToolsProviderResult( + name=provider.name, + tools=tool_names, + ) + ) + except Exception as exc: + providers.append( + McpToolsProviderResult( + name=provider.name or provider_name, + error=str(exc).strip() or "Failed to load tools.", + ) + ) + + duplicate_tools = { + tool_name: provider_names + for tool_name, provider_names in sorted(tool_to_providers.items()) + if len(provider_names) > 1 + } + + return McpToolsListResponse( + providers=providers, + duplicate_tools=duplicate_tools, + ) diff --git a/studio/frontend/src/features/recipe-studio/api/index.ts b/studio/frontend/src/features/recipe-studio/api/index.ts index c1d4174feb..ef13eb92da 100644 --- a/studio/frontend/src/features/recipe-studio/api/index.ts +++ b/studio/frontend/src/features/recipe-studio/api/index.ts @@ -124,6 +124,25 @@ export type ValidateResponse = { raw_detail?: string | null; }; +export type McpToolsListRequest = { + // biome-ignore lint/style/useNamingConvention: api schema + mcp_providers: Record[]; + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec?: number; +}; + +export type McpToolsProviderResult = { + name: string; + tools: string[]; + error?: string | null; +}; + +export type McpToolsListResponse = { + providers: McpToolsProviderResult[]; + // biome-ignore lint/style/useNamingConvention: api schema + duplicate_tools: Record; +}; + async function parseErrorResponse(response: Response): Promise { const text = (await response.text()).trim(); if (!text) { @@ -261,6 +280,12 @@ export async function inspectSeedUpload( return postJson("/seed/inspect-upload", payload); } +export async function listMcpTools( + payload: McpToolsListRequest, +): Promise { + return postJson("/mcp/tools", payload); +} + export async function streamRecipeJobEvents(options: { jobId: string; signal: AbortSignal; @@ -322,4 +347,4 @@ export async function streamRecipeJobEvents(options: { } } -// NOTE: tools + seed inspect/preview endpoints removed from harness. +// NOTE: preview endpoints removed from harness. diff --git a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts index 2a16d7d825..c2d68b26cf 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/definitions.ts +++ b/studio/frontend/src/features/recipe-studio/blocks/definitions.ts @@ -9,6 +9,7 @@ import { EqualSignIcon, FingerPrintIcon, FunctionIcon, + Plug01Icon, Parabola02Icon, PencilEdit02Icon, Plant01Icon, @@ -29,6 +30,7 @@ import { makeMarkdownNoteConfig, makeModelConfig, makeModelProviderConfig, + makeToolProfileConfig, makeSamplerConfig, makeSeedConfig, makeValidatorConfig, @@ -54,7 +56,8 @@ export type BlockType = | "seed_local" | "seed_unstructured" | "model_provider" - | "model_config"; + | "model_config" + | "tool_config"; export type SeedBlockType = "seed_hf" | "seed_local" | "seed_unstructured"; @@ -83,6 +86,7 @@ export type BlockDialogKey = | "validator" | "model_provider" | "model_config" + | "tool_config" | "expression"; export type BlockDefinition = { @@ -111,7 +115,7 @@ export const BLOCK_GROUPS: BlockGroup[] = [ { kind: "llm", title: "LLM + Models", - description: "Generation, providers, and model aliases.", + description: "Generation, model aliases, and shared tool profiles.", icon: PencilEdit02Icon, }, { @@ -297,6 +301,15 @@ const BLOCK_DEFINITIONS: BlockDefinition[] = [ dialogKey: "model_config", createConfig: (id, existing) => makeModelConfig(id, existing), }, + { + kind: "llm", + type: "tool_config", + title: "Tool Profile", + description: "Reusable MCP servers + allowed tools for one or more LLMs.", + icon: Plug01Icon, + dialogKey: "tool_config", + createConfig: (id, existing) => makeToolProfileConfig(id, existing), + }, { kind: "validator", type: "validator_python", @@ -399,6 +412,9 @@ export function getBlockDefinitionForConfig( if (config.kind === "model_config") { return getBlockDefinition("llm", "model_config"); } + if (config.kind === "tool_config") { + return getBlockDefinition("llm", "tool_config"); + } if (config.kind === "markdown_note") { return getBlockDefinition("note", "markdown_note"); } diff --git a/studio/frontend/src/features/recipe-studio/blocks/render-dialog.tsx b/studio/frontend/src/features/recipe-studio/blocks/render-dialog.tsx index d07aa6b627..d8d8222842 100644 --- a/studio/frontend/src/features/recipe-studio/blocks/render-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/blocks/render-dialog.tsx @@ -16,6 +16,7 @@ import { TimedeltaDialog } from "../dialogs/samplers/timedelta-dialog"; import { UniformDialog } from "../dialogs/samplers/uniform-dialog"; import { UuidDialog } from "../dialogs/samplers/uuid-dialog"; import { MarkdownNoteDialog } from "../dialogs/markdown-note/markdown-note-dialog"; +import { ToolProfileDialog } from "../dialogs/tool-profile/tool-profile-dialog"; import { ValidatorDialog } from "../dialogs/validators/validator-dialog"; export function renderBlockDialog( @@ -24,6 +25,7 @@ export function renderBlockDialog( categoryOptions: SamplerConfig[], modelConfigAliases: string[], modelProviderOptions: string[], + toolProfileAliases: string[], datetimeOptions: string[], onUpdate: (id: string, patch: Partial) => void, ): ReactElement | null { @@ -91,6 +93,7 @@ export function renderBlockDialog( config={config} modelConfigAliases={modelConfigAliases} modelProviderOptions={modelProviderOptions} + toolProfileAliases={toolProfileAliases} onUpdate={update} /> ) : null; @@ -106,6 +109,10 @@ export function renderBlockDialog( onUpdate={update} /> ) : null; + case "tool_config": + return config.kind === "tool_config" ? ( + + ) : null; case "expression": return config.kind === "expression" ? ( diff --git a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx index ffe11076ad..079a0ee004 100644 --- a/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx +++ b/studio/frontend/src/features/recipe-studio/components/block-sheet.tsx @@ -71,6 +71,7 @@ type BlockSheetProps = { onAddLlm: (type: LlmType) => void; onAddModelProvider: () => void; onAddModelConfig: () => void; + onAddToolProfile: () => void; onAddExpression: () => void; onAddValidator: ( type: "validator_python" | "validator_sql" | "validator_oxc", @@ -225,6 +226,7 @@ export function BlockSheet({ onAddLlm, onAddModelProvider, onAddModelConfig, + onAddToolProfile, onAddExpression, onAddValidator, onAddMarkdownNote, @@ -333,6 +335,10 @@ export function BlockSheet({ onAddModelConfig(); return; } + if (type === "tool_config") { + onAddToolProfile(); + return; + } onAddLlm(type as LlmType); return; } diff --git a/studio/frontend/src/features/recipe-studio/components/inline/inline-llm.tsx b/studio/frontend/src/features/recipe-studio/components/inline/inline-llm.tsx index 6a269ca684..1e7bfef6c5 100644 --- a/studio/frontend/src/features/recipe-studio/components/inline/inline-llm.tsx +++ b/studio/frontend/src/features/recipe-studio/components/inline/inline-llm.tsx @@ -52,9 +52,17 @@ export function InlineLlm({ config, onUpdate }: InlineLlmProps): ReactElement { .map((c) => c.name), [configs], ); + const toolProfileAliases = useMemo( + () => + Object.values(configs) + .filter((c) => c.kind === "tool_config") + .map((c) => c.name), + [configs], + ); const aliasInputRef = useRef(config.model_alias); const lastAliasRef = useRef(config.model_alias); const anchorRef = useRef(null); + const toolAnchorRef = useRef(null); if (lastAliasRef.current !== config.model_alias) { lastAliasRef.current = config.model_alias; aliasInputRef.current = config.model_alias; @@ -107,6 +115,48 @@ export function InlineLlm({ config, onUpdate }: InlineLlmProps): ReactElement { + +
+ + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + tool_alias: value ?? "", + }) + } + itemToStringValue={(value) => value} + autoHighlight={true} + > + { + const next = event.target.value; + if (next !== (config.tool_alias ?? "")) { + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + tool_alias: next, + }); + } + }} + /> + + No tool profiles found + + {(alias: string) => ( + + {alias} + + )} + + + +
+ {isCode && ( - onUpdate({ - // biome-ignore lint/style/useNamingConvention: api schema - tool_alias: value, - }) - } - > - - - - - {toolAliasOptions.map((alias) => ( - - {alias} - - ))} - - - ) : ( -

- Add tool config alias first. -

- )} - - { - void loadToolNames(); - }} - onAddToolConfig={addToolConfig} - onUpdateToolConfig={updateToolConfigAt} - onRemoveToolConfig={removeToolConfig} - /> - - - ); -} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/mcp-providers-section.tsx b/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/mcp-providers-section.tsx deleted file mode 100644 index c2b892564b..0000000000 --- a/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/mcp-providers-section.tsx +++ /dev/null @@ -1,272 +0,0 @@ -import { Delete02Icon, PlusSignIcon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; -import { Button } from "@/components/ui/button"; -import { Input } from "@/components/ui/input"; -import { Tabs, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import type { LlmMcpProviderConfig, McpEnvVar } from "../../../types"; -import { FieldLabel } from "../../shared/field-label"; - -type McpProvidersSectionProps = { - providers: LlmMcpProviderConfig[]; - onAddProvider: () => void; - onUpdateProviderAt: ( - index: number, - patch: Partial, - ) => void; - onRemoveProvider: (index: number) => void; - onAddProviderArg: (providerIndex: number) => void; - onUpdateProviderArg: ( - providerIndex: number, - argIndex: number, - value: string, - ) => void; - onRemoveProviderArg: (providerIndex: number, argIndex: number) => void; - onAddProviderEnv: (providerIndex: number) => void; - onUpdateProviderEnv: ( - providerIndex: number, - envIndex: number, - patch: Partial, - ) => void; - onRemoveProviderEnv: (providerIndex: number, envIndex: number) => void; -}; - -export function McpProvidersSection({ - providers, - onAddProvider, - onUpdateProviderAt, - onRemoveProvider, - onAddProviderArg, - onUpdateProviderArg, - onRemoveProviderArg, - onAddProviderEnv, - onUpdateProviderEnv, - onRemoveProviderEnv, -}: McpProvidersSectionProps) { - return ( -
-
- - -
- - {providers.length === 0 && ( -

- Add MCP servers to be referenced by tool config providers. -

- )} - - {providers.map((provider, providerIndex) => { - const args = provider.args && provider.args.length > 0 ? provider.args : [""]; - const envVars = - provider.env && provider.env.length > 0 - ? provider.env - : [{ key: "", value: "" }]; - - return ( -
-
- - - onUpdateProviderAt(providerIndex, { name: event.target.value }) - } - /> -
- - - onUpdateProviderAt(providerIndex, { - // biome-ignore lint/style/useNamingConvention: ui schema - provider_type: value === "stdio" ? "stdio" : "streamable_http", - }) - } - > - - STDIO - Streamable HTTP - - - - {provider.provider_type === "stdio" ? ( -
-
- - - onUpdateProviderAt(providerIndex, { - command: event.target.value, - }) - } - /> -
- -
- - {args.map((arg, argIndex) => ( -
- - onUpdateProviderArg(providerIndex, argIndex, event.target.value) - } - /> - -
- ))} - -
- -
- - {envVars.map((item, envIndex) => ( -
- - onUpdateProviderEnv(providerIndex, envIndex, { - key: event.target.value, - }) - } - /> - - onUpdateProviderEnv(providerIndex, envIndex, { - value: event.target.value, - }) - } - /> - -
- ))} - -
-
- ) : ( -
-
- - - onUpdateProviderAt(providerIndex, { - endpoint: event.target.value, - }) - } - /> -
-
- - - onUpdateProviderAt(providerIndex, { - // biome-ignore lint/style/useNamingConvention: api schema - api_key_env: event.target.value, - }) - } - /> -
-
- - - onUpdateProviderAt(providerIndex, { - // biome-ignore lint/style/useNamingConvention: api schema - api_key: event.target.value, - }) - } - /> -
-
- )} - -
- -
-
- ); - })} -
- ); -} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/tool-configs-section.tsx b/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/tool-configs-section.tsx deleted file mode 100644 index 7dcf601554..0000000000 --- a/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/tool-configs-section.tsx +++ /dev/null @@ -1,182 +0,0 @@ -import { PlusSignIcon } from "@hugeicons/core-free-icons"; -import { HugeiconsIcon } from "@hugeicons/react"; -import { Button } from "@/components/ui/button"; -import { Input } from "@/components/ui/input"; -import { ChipInput } from "../../../components/chip-input"; -import type { LlmToolConfig } from "../../../types"; -import { addUnique, collectToolSuggestions } from "./helpers"; -import { FieldLabel } from "../../shared/field-label"; - -type ToolConfigsSectionProps = { - toolConfigs: LlmToolConfig[]; - providerNameSuggestions: string[]; - toolsByProvider: Record; - loadingTools: boolean; - onFetchTools: () => void; - onAddToolConfig: () => void; - onUpdateToolConfig: (index: number, patch: Partial) => void; - onRemoveToolConfig: (index: number) => void; -}; - -export function ToolConfigsSection({ - toolConfigs, - providerNameSuggestions, - toolsByProvider, - loadingTools, - onFetchTools, - onAddToolConfig, - onUpdateToolConfig, - onRemoveToolConfig, -}: ToolConfigsSectionProps) { - return ( -
-
- -
- - -
-
-

- Define aliases/providers here. Active alias is selected above. -

- {toolConfigs.length === 0 && ( -

- Add at least one tool config to map alias to providers. -

- )} - {toolConfigs.map((toolConfig, index) => ( -
-
- - - onUpdateToolConfig(index, { - // biome-ignore lint/style/useNamingConvention: api schema - tool_alias: event.target.value, - }) - } - /> -
- -
- - - onUpdateToolConfig(index, { - providers: addUnique(toolConfig.providers, value), - }) - } - onRemove={(providerIndex) => - onUpdateToolConfig(index, { - providers: toolConfig.providers.filter( - (_, currentIndex) => currentIndex !== providerIndex, - ), - }) - } - placeholder="Type provider name and press Enter" - /> -
- -
- - - onUpdateToolConfig(index, { - // biome-ignore lint/style/useNamingConvention: api schema - allow_tools: addUnique(toolConfig.allow_tools ?? [], value), - }) - } - onRemove={(toolIndex) => - onUpdateToolConfig(index, { - // biome-ignore lint/style/useNamingConvention: api schema - allow_tools: (toolConfig.allow_tools ?? []).filter( - (_, currentIndex) => currentIndex !== toolIndex, - ), - }) - } - placeholder="Type tool name and press Enter" - /> -
- -
-
- - - onUpdateToolConfig(index, { - // biome-ignore lint/style/useNamingConvention: api schema - max_tool_call_turns: event.target.value, - }) - } - /> -
-
- - - onUpdateToolConfig(index, { - // biome-ignore lint/style/useNamingConvention: api schema - timeout_sec: event.target.value, - }) - } - /> -
-
- -
- -
-
- ))} -
- ); -} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/helpers.ts b/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/helpers.ts similarity index 72% rename from studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/helpers.ts rename to studio/frontend/src/features/recipe-studio/dialogs/tool-profile/helpers.ts index 7c8788beec..332c92fbfa 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/llm/mcp-tools/helpers.ts +++ b/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/helpers.ts @@ -1,13 +1,9 @@ -import type { LlmMcpProviderConfig, LlmToolConfig } from "../../../types"; +import type { LlmMcpProviderConfig } from "../../types"; export function createMcpProviderId(prefix: string, index: number): string { return `${prefix}-mcp-${Date.now()}-${index + 1}`; } -export function createToolConfigId(prefix: string, index: number): string { - return `${prefix}-tool-${Date.now()}-${index + 1}`; -} - export function addUnique(items: string[], value: string): string[] { const trimmed = value.trim(); if (!trimmed || items.includes(trimmed)) { @@ -16,6 +12,32 @@ export function addUnique(items: string[], value: string): string[] { return [...items, trimmed]; } +export function collectToolSuggestions( + providerNames: string[], + toolsByProvider: Record, +): string[] { + return Array.from( + new Set( + providerNames.flatMap( + (providerName) => toolsByProvider[providerName.trim()] ?? [], + ), + ), + ); +} + +export function isProviderReadyForToolFetch( + provider: LlmMcpProviderConfig, +): boolean { + const hasName = provider.name.trim().length > 0; + if (!hasName) { + return false; + } + if (provider.provider_type === "stdio") { + return (provider.command?.trim().length ?? 0) > 0; + } + return (provider.endpoint?.trim().length ?? 0) > 0; +} + export function toApiProvider( provider: LlmMcpProviderConfig, ): Record { @@ -45,43 +67,3 @@ export function toApiProvider( api_key_env: provider.api_key_env?.trim() || undefined, }; } - -export function collectToolSuggestions( - providerNames: string[], - toolsByProvider: Record, -): string[] { - return Array.from( - new Set( - providerNames.flatMap((providerName) => { - return toolsByProvider[providerName.trim()] ?? []; - }), - ), - ); -} - -export function isProviderReadyForToolFetch( - provider: LlmMcpProviderConfig, -): boolean { - const hasName = provider.name.trim().length > 0; - if (!hasName) { - return false; - } - if (provider.provider_type === "stdio") { - return (provider.command?.trim().length ?? 0) > 0; - } - return (provider.endpoint?.trim().length ?? 0) > 0; -} - -export function resolveLlmToolAlias( - toolConfigs: LlmToolConfig[], - previousAlias: string | undefined, -): string { - const toolAliases = toolConfigs - .map((item) => item.tool_alias.trim()) - .filter(Boolean); - const currentAlias = previousAlias?.trim() ?? ""; - if (currentAlias && toolAliases.includes(currentAlias)) { - return currentAlias; - } - return toolAliases[0] ?? ""; -} diff --git a/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx new file mode 100644 index 0000000000..d625c7828d --- /dev/null +++ b/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx @@ -0,0 +1,776 @@ +import { Badge } from "@/components/ui/badge"; +import { Button } from "@/components/ui/button"; +import { + Collapsible, + CollapsibleContent, + CollapsibleTrigger, +} from "@/components/ui/collapsible"; +import { Input } from "@/components/ui/input"; +import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; +import { toastError } from "@/shared/toast"; +import { + ArrowRight01Icon, + Delete02Icon, + PlusSignIcon, +} from "@hugeicons/core-free-icons"; +import { HugeiconsIcon } from "@hugeicons/react"; +import { type ReactElement, useEffect, useMemo, useState } from "react"; +import { listMcpTools } from "../../api"; +import { ChipInput } from "../../components/chip-input"; +import type { LlmMcpProviderConfig, McpEnvVar, ToolProfileConfig } from "../../types"; +import { FieldLabel } from "../shared/field-label"; +import { NameField } from "../shared/name-field"; +import { + addUnique, + collectToolSuggestions, + createMcpProviderId, + isProviderReadyForToolFetch, + toApiProvider, +} from "./helpers"; + +type ToolProfileDialogProps = { + config: ToolProfileConfig; + onUpdate: (patch: Partial) => void; +}; + +function EmptyState({ + title, + description, +}: { + title: string; + description: string; +}): ReactElement { + return ( +
+

{title}

+

{description}

+
+ ); +} + +function isProviderConfigured(provider: LlmMcpProviderConfig): boolean { + const hasName = provider.name.trim().length > 0; + if (!hasName) { + return false; + } + if (provider.provider_type === "stdio") { + return (provider.command?.trim().length ?? 0) > 0; + } + return (provider.endpoint?.trim().length ?? 0) > 0; +} + +function McpServerCard({ + provider, + index, + toolsCount, + error, + open, + onOpenChange, + onUpdateProviderAt, + onRemoveProvider, + onAddProviderArg, + onUpdateProviderArg, + onRemoveProviderArg, + onAddProviderEnv, + onUpdateProviderEnv, + onRemoveProviderEnv, +}: { + provider: LlmMcpProviderConfig; + index: number; + toolsCount?: number; + error?: string; + open: boolean; + onOpenChange: (open: boolean) => void; + onUpdateProviderAt: ( + index: number, + patch: Partial, + ) => void; + onRemoveProvider: (index: number) => void; + onAddProviderArg: (index: number) => void; + onUpdateProviderArg: (index: number, argIndex: number, value: string) => void; + onRemoveProviderArg: (index: number, argIndex: number) => void; + onAddProviderEnv: (index: number) => void; + onUpdateProviderEnv: ( + index: number, + envIndex: number, + patch: Partial, + ) => void; + onRemoveProviderEnv: (index: number, envIndex: number) => void; +}): ReactElement { + const args = provider.args && provider.args.length > 0 ? provider.args : [""]; + const envVars = + provider.env && provider.env.length > 0 + ? provider.env + : [{ key: "", value: "" }]; + const summaryTitle = provider.name.trim() || `MCP server ${index + 1}`; + const transportLabel = + provider.provider_type === "stdio" ? "STDIO" : "Streamable HTTP"; + const toolsLabel = typeof toolsCount === "number" ? `${toolsCount} tools` : null; + const description = + provider.provider_type === "stdio" + ? "Launches a local MCP process over stdio." + : "Calls a remote MCP endpoint from the backend."; + + return ( + +
+
+ + + + +
+ + + {error && ( +
+ {error} +
+ )} + +
+ + + onUpdateProviderAt(index, { name: event.target.value }) + } + /> +
+ + + onUpdateProviderAt(index, { + // biome-ignore lint/style/useNamingConvention: ui schema + provider_type: value === "stdio" ? "stdio" : "streamable_http", + }) + } + > + + STDIO + Streamable HTTP + + + + {provider.provider_type === "stdio" ? ( +
+
+ + + onUpdateProviderAt(index, { command: event.target.value }) + } + /> +
+ +
+
+ + +
+ {args.map((arg, argIndex) => ( +
+ + onUpdateProviderArg(index, argIndex, event.target.value) + } + /> + +
+ ))} +
+ +
+
+ + +
+ {envVars.map((item, envIndex) => ( +
+ + onUpdateProviderEnv(index, envIndex, { + key: event.target.value, + }) + } + /> + + onUpdateProviderEnv(index, envIndex, { + value: event.target.value, + }) + } + /> + +
+ ))} +
+
+ ) : ( +
+
+ + + onUpdateProviderAt(index, { endpoint: event.target.value }) + } + /> +
+
+
+ + + onUpdateProviderAt(index, { + // biome-ignore lint/style/useNamingConvention: api schema + api_key_env: event.target.value, + }) + } + /> +
+
+ + + onUpdateProviderAt(index, { + // biome-ignore lint/style/useNamingConvention: api schema + api_key: event.target.value, + }) + } + /> +
+
+
+ )} + +
+ + ); +} + +export function ToolProfileDialog({ + config, + onUpdate, +}: ToolProfileDialogProps): ReactElement { + const providers = config.mcp_providers; + const [loadingTools, setLoadingTools] = useState(false); + const [toolsByProvider, setToolsByProvider] = useState>( + {}, + ); + const [providerErrors, setProviderErrors] = useState>({}); + const [duplicateTools, setDuplicateTools] = useState>({}); + const [openProviders, setOpenProviders] = useState>({}); + + const providerSignature = useMemo( + () => + JSON.stringify( + providers.map((provider) => ({ + name: provider.name, + // biome-ignore lint/style/useNamingConvention: ui schema + provider_type: provider.provider_type, + command: provider.command, + args: provider.args, + env: provider.env, + endpoint: provider.endpoint, + // biome-ignore lint/style/useNamingConvention: api schema + api_key: provider.api_key, + // biome-ignore lint/style/useNamingConvention: api schema + api_key_env: provider.api_key_env, + })), + ), + [providers], + ); + + useEffect(() => { + setToolsByProvider({}); + setProviderErrors({}); + setDuplicateTools({}); + }, [providerSignature]); + + useEffect(() => { + setOpenProviders((current) => { + const next: Record = {}; + for (const provider of providers) { + next[provider.id] = + current[provider.id] ?? !isProviderConfigured(provider); + } + return next; + }); + }, [providers]); + + function updateProviders(nextProviders: LlmMcpProviderConfig[]): void { + onUpdate({ + // biome-ignore lint/style/useNamingConvention: ui schema + mcp_providers: nextProviders, + }); + } + + function updateProviderAt( + index: number, + patch: Partial, + ): void { + updateProviders( + providers.map((provider, currentIndex) => + currentIndex === index ? { ...provider, ...patch } : provider, + ), + ); + } + + function mutateProviderAt( + index: number, + mapProvider: (provider: LlmMcpProviderConfig) => Partial, + ): void { + const provider = providers[index]; + if (!provider) { + return; + } + updateProviderAt(index, mapProvider(provider)); + } + + function removeProvider(index: number): void { + updateProviders(providers.filter((_, currentIndex) => currentIndex !== index)); + } + + function addProvider(): void { + updateProviders([ + ...providers, + { + id: createMcpProviderId(config.id, providers.length), + name: "", + // biome-ignore lint/style/useNamingConvention: ui schema + provider_type: "stdio", + command: "", + args: [], + env: [], + endpoint: "", + // biome-ignore lint/style/useNamingConvention: api schema + api_key: "", + // biome-ignore lint/style/useNamingConvention: api schema + api_key_env: "", + }, + ]); + } + + function addProviderArg(providerIndex: number): void { + mutateProviderAt(providerIndex, (provider) => ({ + args: [...(provider.args ?? []), ""], + })); + } + + function updateProviderArg( + providerIndex: number, + argIndex: number, + value: string, + ): void { + mutateProviderAt(providerIndex, (provider) => { + const nextArgs = + provider.args && provider.args.length > 0 ? [...provider.args] : [""]; + nextArgs[argIndex] = value; + return { args: nextArgs }; + }); + } + + function removeProviderArg(providerIndex: number, argIndex: number): void { + mutateProviderAt(providerIndex, (provider) => ({ + args: (provider.args ?? []).filter((_, currentIndex) => currentIndex !== argIndex), + })); + } + + function addProviderEnv(providerIndex: number): void { + mutateProviderAt(providerIndex, (provider) => ({ + env: [...(provider.env ?? []), { key: "", value: "" }], + })); + } + + function updateProviderEnv( + providerIndex: number, + envIndex: number, + patch: Partial, + ): void { + mutateProviderAt(providerIndex, (provider) => ({ + env: ( + provider.env && provider.env.length > 0 + ? provider.env + : [{ key: "", value: "" }] + ).map((item, currentIndex) => + currentIndex === envIndex ? { ...item, ...patch } : item, + ), + })); + } + + function removeProviderEnv(providerIndex: number, envIndex: number): void { + mutateProviderAt(providerIndex, (provider) => ({ + env: (provider.env ?? []).filter((_, currentIndex) => currentIndex !== envIndex), + })); + } + + async function loadTools(): Promise { + const readyProviders = providers.filter(isProviderReadyForToolFetch); + if (readyProviders.length === 0) { + toastError( + "No MCP servers ready", + "Add a server name plus command or endpoint first.", + ); + return; + } + + setLoadingTools(true); + try { + const timeoutRaw = config.timeout_sec?.trim(); + const timeoutSec = + timeoutRaw && Number.isFinite(Number(timeoutRaw)) + ? Number(timeoutRaw) + : 15; + const response = await listMcpTools({ + // biome-ignore lint/style/useNamingConvention: api schema + mcp_providers: readyProviders.map(toApiProvider), + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec: timeoutSec, + }); + setToolsByProvider( + Object.fromEntries( + response.providers + .filter((provider) => provider.name.trim()) + .map((provider) => [provider.name.trim(), provider.tools]), + ), + ); + setProviderErrors( + Object.fromEntries( + response.providers + .filter((provider) => provider.name.trim() && provider.error) + .map((provider) => [provider.name.trim(), provider.error ?? "Failed to load tools."]), + ), + ); + setDuplicateTools(response.duplicate_tools ?? {}); + } catch (error) { + toastError( + "Failed to load tools", + error instanceof Error ? error.message : "Could not load MCP tools.", + ); + } finally { + setLoadingTools(false); + } + } + + const providerNames = useMemo( + () => + Array.from( + new Set(providers.map((provider) => provider.name.trim()).filter(Boolean)), + ), + [providers], + ); + const availableTools = useMemo( + () => collectToolSuggestions(providerNames, toolsByProvider), + [providerNames, toolsByProvider], + ); + const hasProviders = providers.length > 0; + + return ( + + + Profile + MCP servers + + + + onUpdate({ name: value })} + /> + + {!hasProviders ? ( + + ) : ( + <> +
+ +
+ {providerNames.map((providerName) => ( + + {providerName} + + ))} +
+
+ +
+
+
+

+ Available tool refs +

+

+ Load tools from backend so users pick tool names instead of guessing. +

+
+ +
+ + {Object.keys(toolsByProvider).length === 0 && + Object.keys(providerErrors).length === 0 && ( +

+ No tools loaded yet. +

+ )} + + {Object.entries(toolsByProvider).map(([providerName, toolNames]) => ( +
+
+

+ {providerName} +

+ + {toolNames.length} + +
+
+ {toolNames.map((toolName) => ( + + {toolName} + + ))} +
+
+ ))} + + {Object.entries(duplicateTools).length > 0 && ( +
+ Duplicate tool names across servers: + {" "} + {Object.entries(duplicateTools) + .map(([toolName, providerList]) => `${toolName} (${providerList.join(", ")})`) + .join("; ")} +
+ )} +
+ +
+ + + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + allow_tools: addUnique(config.allow_tools ?? [], value), + }) + } + onRemove={(toolIndex) => + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + allow_tools: (config.allow_tools ?? []).filter( + (_, currentIndex) => currentIndex !== toolIndex, + ), + }) + } + placeholder="Type tool name and press Enter" + /> +
+ +
+
+ + + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + max_tool_call_turns: event.target.value, + }) + } + /> +
+
+ + + onUpdate({ + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec: event.target.value, + }) + } + /> +
+
+ + )} + + + +
+ + +
+ + {!hasProviders ? ( + + ) : ( +
+ {providers.map((provider, index) => ( + + setOpenProviders((current) => ({ + ...current, + [provider.id]: open, + })) + } + onUpdateProviderAt={updateProviderAt} + onRemoveProvider={removeProvider} + onAddProviderArg={addProviderArg} + onUpdateProviderArg={updateProviderArg} + onRemoveProviderArg={removeProviderArg} + onAddProviderEnv={addProviderEnv} + onUpdateProviderEnv={updateProviderEnv} + onRemoveProviderEnv={removeProviderEnv} + /> + ))} +
+ )} + + + ); +} diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-editor-graph.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-editor-graph.ts index 8e8582b855..11e60c06d7 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-editor-graph.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-editor-graph.ts @@ -67,6 +67,7 @@ type UseRecipeEditorGraphArgs = { addLlmNode: (type: LlmType, position?: XYPosition, openDialog?: boolean) => void; addModelProviderNode: (position?: XYPosition, openDialog?: boolean) => void; addModelConfigNode: (position?: XYPosition, openDialog?: boolean) => void; + addToolProfileNode: (position?: XYPosition, openDialog?: boolean) => void; addExpressionNode: (position?: XYPosition, openDialog?: boolean) => void; addValidatorNode: ( type: "validator_python" | "validator_sql" | "validator_oxc", @@ -90,6 +91,7 @@ type UseRecipeEditorGraphResult = { handleAddLlmFromSheet: (type: LlmType) => void; handleAddModelProviderFromSheet: () => void; handleAddModelConfigFromSheet: () => void; + handleAddToolProfileFromSheet: () => void; handleAddExpressionFromSheet: () => void; handleAddValidatorFromSheet: ( type: "validator_python" | "validator_sql" | "validator_oxc", @@ -113,6 +115,7 @@ export function useRecipeEditorGraph({ addLlmNode, addModelProviderNode, addModelConfigNode, + addToolProfileNode, addExpressionNode, addValidatorNode, addMarkdownNoteNode, @@ -231,6 +234,10 @@ export function useRecipeEditorGraph({ addModelConfigNode(position, false); return; } + if (payload.type === "tool_config") { + addToolProfileNode(position, false); + return; + } addLlmNode(payload.type as LlmType, position, false); }, [ @@ -239,6 +246,7 @@ export function useRecipeEditorGraph({ addMarkdownNoteNode, addModelConfigNode, addModelProviderNode, + addToolProfileNode, addSamplerNode, addSeedNode, addValidatorNode, @@ -290,6 +298,10 @@ export function useRecipeEditorGraph({ addExpressionNode(getViewportCenterPosition()); }, [addExpressionNode, getViewportCenterPosition]); + const handleAddToolProfileFromSheet = useCallback(() => { + addToolProfileNode(getViewportCenterPosition()); + }, [addToolProfileNode, getViewportCenterPosition]); + const handleAddValidatorFromSheet = useCallback( (type: "validator_python" | "validator_sql" | "validator_oxc") => { addValidatorNode(type, getViewportCenterPosition()); @@ -313,6 +325,7 @@ export function useRecipeEditorGraph({ handleAddLlmFromSheet, handleAddModelProviderFromSheet, handleAddModelConfigFromSheet, + handleAddToolProfileFromSheet, handleAddExpressionFromSheet, handleAddValidatorFromSheet, handleAddMarkdownNoteFromSheet, diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-runtime-visuals.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-runtime-visuals.ts index 7d0b259181..5389909b74 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-runtime-visuals.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-runtime-visuals.ts @@ -7,6 +7,7 @@ import { EqualSignIcon, FingerPrintIcon, FunctionIcon, + Plug01Icon, Parabola02Icon, PencilEdit02Icon, Plant01Icon, @@ -78,6 +79,9 @@ function resolveExecutionColumnIcon(config: NodeConfig | null): IconType { if (config.kind === "model_config") { return Plant01Icon; } + if (config.kind === "tool_config") { + return Plug01Icon; + } return PencilEdit02Icon; } diff --git a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx index 3bc2384ac8..d2617cb21f 100644 --- a/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx +++ b/studio/frontend/src/features/recipe-studio/recipe-studio-page.tsx @@ -104,6 +104,7 @@ export function RecipeStudioPage({ addLlmNode, addModelProviderNode, addModelConfigNode, + addToolProfileNode, addExpressionNode, addValidatorNode, addMarkdownNoteNode, @@ -141,6 +142,7 @@ export function RecipeStudioPage({ addLlmNode: state.addLlmNode, addModelProviderNode: state.addModelProviderNode, addModelConfigNode: state.addModelConfigNode, + addToolProfileNode: state.addToolProfileNode, addExpressionNode: state.addExpressionNode, addValidatorNode: state.addValidatorNode, addMarkdownNoteNode: state.addMarkdownNoteNode, @@ -191,6 +193,7 @@ export function RecipeStudioPage({ handleAddLlmFromSheet, handleAddModelProviderFromSheet, handleAddModelConfigFromSheet, + handleAddToolProfileFromSheet, handleAddExpressionFromSheet, handleAddValidatorFromSheet, handleAddMarkdownNoteFromSheet, @@ -210,6 +213,7 @@ export function RecipeStudioPage({ addLlmNode, addModelProviderNode, addModelConfigNode, + addToolProfileNode, addExpressionNode, addValidatorNode, addMarkdownNoteNode, @@ -584,10 +588,11 @@ export function RecipeStudioPage({ onOpenChange={setBlockSheetOpen} onAddSampler={handleAddSamplerFromSheet} onAddSeed={handleAddSeedFromSheet} - onAddLlm={handleAddLlmFromSheet} - onAddModelProvider={handleAddModelProviderFromSheet} - onAddModelConfig={handleAddModelConfigFromSheet} - onAddExpression={handleAddExpressionFromSheet} + onAddLlm={handleAddLlmFromSheet} + onAddModelProvider={handleAddModelProviderFromSheet} + onAddModelConfig={handleAddModelConfigFromSheet} + onAddToolProfile={handleAddToolProfileFromSheet} + onAddExpression={handleAddExpressionFromSheet} onAddValidator={handleAddValidatorFromSheet} onAddMarkdownNote={handleAddMarkdownNoteFromSheet} onOpenProcessors={openProcessorsFromSheet} @@ -651,6 +656,7 @@ export function RecipeStudioPage({ categoryOptions={dialogOptions.categoryOptions} modelConfigAliases={dialogOptions.modelConfigAliases} modelProviderOptions={dialogOptions.modelProviderOptions} + toolProfileAliases={dialogOptions.toolProfileAliases} datetimeOptions={dialogOptions.datetimeOptions} onUpdate={updateConfig} container={sheetContainer} diff --git a/studio/frontend/src/features/recipe-studio/stores/helpers/edge-sync.ts b/studio/frontend/src/features/recipe-studio/stores/helpers/edge-sync.ts index cf4812959f..379a285654 100644 --- a/studio/frontend/src/features/recipe-studio/stores/helpers/edge-sync.ts +++ b/studio/frontend/src/features/recipe-studio/stores/helpers/edge-sync.ts @@ -195,6 +195,41 @@ export function syncEdgesForConfigPatch( } } + const hasToolAliasPatch = Object.prototype.hasOwnProperty.call( + patch, + "tool_alias", + ); + if (current.kind === "llm" && hasToolAliasPatch) { + const nextAlias = + (patch as Partial & { tool_alias?: string }).tool_alias ?? ""; + if (nextAlias.trim() === (current.tool_alias ?? "").trim()) { + return nextEdges; + } + nextEdges = removeTargetEdgesBySource( + nextEdges, + configs, + current.id, + (source) => Boolean(source && source.kind === "tool_config"), + ); + if (nextAlias) { + const toolConfigId = findNodeIdByName(configs, nextAlias); + if (toolConfigId) { + const result = applyRecipeConnection( + { + source: toolConfigId, + sourceHandle: HANDLE_IDS.semanticOut, + target: current.id, + targetHandle: HANDLE_IDS.semanticIn, + }, + configs, + nextEdges, + layoutDirection, + ); + nextEdges = result.edges; + } + } + } + const hasValidatorTargetsPatch = Object.prototype.hasOwnProperty.call( patch, "target_columns", diff --git a/studio/frontend/src/features/recipe-studio/stores/helpers/model-infra-layout.ts b/studio/frontend/src/features/recipe-studio/stores/helpers/model-infra-layout.ts index ce09c46513..cd065103af 100644 --- a/studio/frontend/src/features/recipe-studio/stores/helpers/model-infra-layout.ts +++ b/studio/frontend/src/features/recipe-studio/stores/helpers/model-infra-layout.ts @@ -86,6 +86,15 @@ function isConfigToLlmEdge( return source?.kind === "model_config" && target?.kind === "llm"; } +function isToolConfigToLlmEdge( + edge: Edge, + configs: Record, +): boolean { + const source = configs[edge.source]; + const target = configs[edge.target]; + return source?.kind === "tool_config" && target?.kind === "llm"; +} + function usageKey(nodeId: string, handleId: string): string { return `${nodeId}::${handleId}`; } @@ -263,9 +272,11 @@ export function optimizeModelInfraEdgeHandles( const sourceHandleBefore = normalizeRecipeHandleId(edge.sourceHandle); const targetHandleBefore = normalizeRecipeHandleId(edge.targetHandle); - const isModelSemantic = - isProviderToConfigEdge(edge, configs) || isConfigToLlmEdge(edge, configs); - if (!isModelSemantic) { + const isSemanticInfra = + isProviderToConfigEdge(edge, configs) || + isConfigToLlmEdge(edge, configs) || + isToolConfigToLlmEdge(edge, configs); + if (!isSemanticInfra) { nextEdges.push(edge); continue; } @@ -340,6 +351,7 @@ export function centerModelInfraNodes( ): RecipeNode[] { const nodesById = new Map(nodes.map((node) => [node.id, node] as const)); const configToLlmIds = new Map(); + const toolConfigToLlmIds = new Map(); const providerToConfigIds = new Map(); for (const edge of edges) { @@ -357,6 +369,14 @@ export function centerModelInfraNodes( entries.push(edge.target); } configToLlmIds.set(edge.source, entries); + continue; + } + if (isToolConfigToLlmEdge(edge, configs)) { + const entries = toolConfigToLlmIds.get(edge.source) ?? []; + if (!entries.includes(edge.target)) { + entries.push(edge.target); + } + toolConfigToLlmIds.set(edge.source, entries); } } @@ -370,6 +390,9 @@ export function centerModelInfraNodes( (config) => config.kind === "model_provider" && nodesById.has(config.id), ) .map((config) => config.id); + const toolConfigIds = Object.values(configs) + .filter((config) => config.kind === "tool_config" && nodesById.has(config.id)) + .map((config) => config.id); const occupiedById = new Map( nodes.map((node) => [node.id, toRect(node)] as const), @@ -444,5 +467,27 @@ export function centerModelInfraNodes( placeNode(modelProviderId, preferred); } + for (const toolConfigId of toolConfigIds) { + const llmIds = toolConfigToLlmIds.get(toolConfigId) ?? []; + const targetBounds = collectBounds(llmIds, nodesById); + const toolConfigNode = nodesById.get(toolConfigId); + if (!(targetBounds && toolConfigNode)) { + continue; + } + const width = readNodeWidth(toolConfigNode) ?? DEFAULT_NODE_WIDTH; + const height = readNodeHeight(toolConfigNode) ?? DEFAULT_NODE_HEIGHT; + const preferred = + direction === "LR" + ? { + x: (targetBounds.minX + targetBounds.maxX) / 2 - width / 2, + y: targetBounds.minY - height - clusterGap, + } + : { + x: targetBounds.minX - width - clusterGap, + y: (targetBounds.minY + targetBounds.maxY) / 2 - height / 2, + }; + placeNode(toolConfigId, preferred); + } + return nodes.map((node) => nodesById.get(node.id) ?? node); } diff --git a/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts b/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts index 736745ac4f..c0d7fe892a 100644 --- a/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts +++ b/studio/frontend/src/features/recipe-studio/stores/helpers/reference-sync.ts @@ -83,6 +83,10 @@ export function applyRenameToConfig( const base = next as LlmConfig; next = { ...base, model_alias: to }; } + if (config.kind === "llm" && config.tool_alias === from) { + const base = next as LlmConfig; + next = { ...base, tool_alias: to }; + } if (config.kind === "validator") { const targets = config.target_columns ?? []; if (targets.includes(from)) { @@ -136,6 +140,10 @@ export function applyRemovalToConfig( const base = next as LlmConfig; next = { ...base, model_alias: "" }; } + if (config.kind === "llm" && config.tool_alias === ref) { + const base = next as LlmConfig; + next = { ...base, tool_alias: "" }; + } if (config.kind === "validator") { const targets = (config.target_columns ?? []).filter((target) => target !== ref); if (targets.length !== (config.target_columns ?? []).length) { diff --git a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts index 38f4dd56d4..7c71400c15 100644 --- a/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts +++ b/studio/frontend/src/features/recipe-studio/stores/recipe-studio.ts @@ -95,6 +95,7 @@ type RecipeStudioState = { addLlmNode: (type: LlmType, position?: XYPosition, openDialog?: boolean) => void; addModelProviderNode: (position?: XYPosition, openDialog?: boolean) => void; addModelConfigNode: (position?: XYPosition, openDialog?: boolean) => void; + addToolProfileNode: (position?: XYPosition, openDialog?: boolean) => void; addExpressionNode: (position?: XYPosition, openDialog?: boolean) => void; addValidatorNode: ( type: "validator_python" | "validator_sql" | "validator_oxc", @@ -249,7 +250,8 @@ function isModelSemanticEdge(edge: Edge, configs: Record): b source && target && ((source.kind === "model_provider" && target.kind === "model_config") || - (source.kind === "model_config" && target.kind === "llm")), + (source.kind === "model_config" && target.kind === "llm") || + (source.kind === "tool_config" && target.kind === "llm")), ); } @@ -534,6 +536,49 @@ export const useRecipeStudioStore = create((set, get) => ({ } return { ...added, nodes, edges, configs }; }), + addToolProfileNode: (position, openDialog = true) => + set((state) => { + if (state.executionLocked) { + return state; + } + const added = buildAddedNodeState( + state, + "llm", + "tool_config", + position, + openDialog, + ); + const context = getAddedNodeContext(added); + if (!context) { + return added; + } + let { nodes, configs } = context; + let edges = state.edges; + const unboundLlms = Object.values(configs).filter( + (config) => config.kind === "llm" && !(config.tool_alias?.trim()), + ); + if (!position && unboundLlms.length > 0) { + nodes = placeNodeNear( + nodes, + context.newNodeId, + unboundLlms[0].id, + state.layoutDirection, + "before", + ); + } + if (unboundLlms.length === 1) { + const next = connectSemantic( + edges, + configs, + context.newNodeId, + unboundLlms[0].id, + state.layoutDirection, + ); + edges = next.edges; + configs = next.configs; + } + return { ...added, nodes, edges, configs }; + }), addExpressionNode: (position, openDialog = true) => set((state) => { if (state.executionLocked) { diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index 9c2b8feac7..c62da8fd9f 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -48,7 +48,8 @@ export type RecipeNodeData = { | "seed" | "note" | "model_provider" - | "model_config"; + | "model_config" + | "tool_config"; subtype: string; blockType: | SamplerType @@ -60,7 +61,8 @@ export type RecipeNodeData = { | "seed" | "markdown_note" | "model_provider" - | "model_config"; + | "model_config" + | "tool_config"; layoutDirection?: LayoutDirection; runtimeState?: "idle" | "running" | "done"; executionLocked?: boolean; @@ -173,6 +175,20 @@ export type LlmToolConfig = { timeout_sec?: string; }; +export type ToolProfileConfig = { + id: string; + kind: "tool_config"; + name: string; + // biome-ignore lint/style/useNamingConvention: ui schema + mcp_providers: LlmMcpProviderConfig[]; + // biome-ignore lint/style/useNamingConvention: api schema + allow_tools?: string[]; + // biome-ignore lint/style/useNamingConvention: api schema + max_tool_call_turns?: string; + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec?: string; +}; + export type LlmImageContextConfig = { enabled: boolean; // biome-ignore lint/style/useNamingConvention: api schema @@ -201,10 +217,6 @@ export type LlmConfig = { output_format?: string; // biome-ignore lint/style/useNamingConvention: api schema tool_alias?: string; - // biome-ignore lint/style/useNamingConvention: api schema - tool_configs?: LlmToolConfig[]; - // biome-ignore lint/style/useNamingConvention: ui schema - mcp_providers?: LlmMcpProviderConfig[]; scores?: Score[]; // ui-only, serialized into multi_modal_context for DataDesigner // biome-ignore lint/style/useNamingConvention: ui schema @@ -349,4 +361,5 @@ export type NodeConfig = | MarkdownNoteConfig | SeedConfig | ModelProviderConfig - | ModelConfig; + | ModelConfig + | ToolProfileConfig; diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index 404cf018de..d605e1a775 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -10,6 +10,7 @@ import type { SeedSourceType, SamplerConfig, SamplerType, + ToolProfileConfig, ValidatorCodeLang, ValidatorType, ValidatorConfig, @@ -203,10 +204,6 @@ export function makeLlmConfig( llmType === "structured" ? '{\n "field": "string"\n}' : undefined, // biome-ignore lint/style/useNamingConvention: api schema tool_alias: "", - // biome-ignore lint/style/useNamingConvention: api schema - tool_configs: [], - // biome-ignore lint/style/useNamingConvention: ui schema - mcp_providers: [], // biome-ignore lint/style/useNamingConvention: ui schema image_context: { enabled: false, @@ -268,6 +265,25 @@ export function makeModelConfig( }; } +export function makeToolProfileConfig( + id: string, + existing: NodeConfig[], +): ToolProfileConfig { + return { + id, + kind: "tool_config", + name: nextName(existing, "tools"), + // biome-ignore lint/style/useNamingConvention: ui schema + mcp_providers: [], + // biome-ignore lint/style/useNamingConvention: api schema + allow_tools: [], + // biome-ignore lint/style/useNamingConvention: api schema + max_tool_call_turns: "5", + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec: "", + }; +} + export function makeExpressionConfig( id: string, existing: NodeConfig[], diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts b/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts index 6699ab1007..bcafb09ebc 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/recipe-graph-connection.ts @@ -58,7 +58,11 @@ function syncSubcategoryMapping( } function isModelInfraNode(config: NodeConfig): boolean { - return config.kind === "model_provider" || config.kind === "model_config"; + return ( + config.kind === "model_provider" || + config.kind === "model_config" || + config.kind === "tool_config" + ); } function isSemanticLane(connection: Connection): boolean { @@ -80,6 +84,7 @@ function isDataLane(connection: Connection): boolean { type SingleRefRelation = | "provider" | "model_alias" + | "tool_alias" | "reference_column_name" | "subcategory_parent" | "validator_target_columns"; @@ -94,6 +99,9 @@ function getSingleRefRelation( if (source.kind === "model_config" && target.kind === "llm") { return "model_alias"; } + if (source.kind === "tool_config" && target.kind === "llm") { + return "tool_alias"; + } if ( source.kind === "sampler" && source.sampler_type === "datetime" && @@ -134,6 +142,9 @@ function isCompetingIncomingEdge( if (relation === "model_alias") { return source.kind === "model_config"; } + if (relation === "tool_alias") { + return source.kind === "tool_config"; + } if (relation === "subcategory_parent") { return isCategoryConfig(source); } @@ -146,7 +157,8 @@ function isCompetingIncomingEdge( function isModelSemanticRelation(source: NodeConfig, target: NodeConfig): boolean { return ( (source.kind === "model_provider" && target.kind === "model_config") || - (source.kind === "model_config" && target.kind === "llm") + (source.kind === "model_config" && target.kind === "llm") || + (source.kind === "tool_config" && target.kind === "llm") ); } @@ -378,6 +390,10 @@ export function applyRecipeConnection( const next = { ...target, model_alias: source.name }; return { edges: nextEdges, configs: { ...configs, [target.id]: next } }; } + if (source.kind === "tool_config" && target.kind === "llm") { + const next = { ...target, tool_alias: source.name }; + return { edges: nextEdges, configs: { ...configs, [target.id]: next } }; + } if ( source.kind === "sampler" && source.sampler_type === "datetime" && diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/relations.ts b/studio/frontend/src/features/recipe-studio/utils/graph/relations.ts index 8513002d9b..a032cc65d8 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/relations.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/relations.ts @@ -10,6 +10,9 @@ export function isSemanticRelation( if (source.kind === "model_config" && target.kind === "llm") { return true; } + if (source.kind === "tool_config" && target.kind === "llm") { + return true; + } if ( source.kind === "llm" && source.llm_type === "code" && diff --git a/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts b/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts index d36014d4c9..4806892d27 100644 --- a/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts +++ b/studio/frontend/src/features/recipe-studio/utils/graph/runtime-visual-state.ts @@ -22,6 +22,7 @@ const DONE_UPSTREAM_KINDS: ReadonlySet = new Set([ "llm", "model_config", "model_provider", + "tool_config", ]); export type GraphRuntimeVisualState = { diff --git a/studio/frontend/src/features/recipe-studio/utils/import/edges.ts b/studio/frontend/src/features/recipe-studio/utils/import/edges.ts index 4e02877b6c..127b2015bb 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/edges.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/edges.ts @@ -20,6 +20,9 @@ function isSemanticConnection(source: NodeConfig, target: NodeConfig): boolean { if (source.kind === "model_config" && target.kind === "llm") { return true; } + if (source.kind === "tool_config" && target.kind === "llm") { + return true; + } if ( source.kind === "llm" && source.llm_type === "code" && @@ -163,6 +166,9 @@ export function buildEdges( if (config.kind === "llm" && config.model_alias) { addEdgeByName(config.model_alias, config.name); } + if (config.kind === "llm" && config.tool_alias) { + addEdgeByName(config.tool_alias, config.name); + } if (config.kind === "validator") { for (const targetColumn of config.target_columns ?? []) { if (targetColumn.trim()) { diff --git a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts index def08974bc..b8823882f7 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts @@ -8,6 +8,7 @@ import type { SeedConfig, SamplerConfig, SeedSourceType, + ToolProfileConfig, ValidatorConfig, } from "../../types"; import { buildEdges } from "./edges"; @@ -216,15 +217,6 @@ function parseToolConfigs(input: unknown): Map { return toolConfigs; } -function cloneToolConfig(config: LlmToolConfig): LlmToolConfig { - return { - ...config, - providers: [...config.providers], - // biome-ignore lint/style/useNamingConvention: api schema - allow_tools: [...(config.allow_tools ?? [])], - }; -} - function cloneMcpProvider(config: LlmMcpProviderConfig): LlmMcpProviderConfig { return { ...config, @@ -296,28 +288,28 @@ function applyAdvancedOpen( config.advancedOpen = advancedOpenByNode[config.name] === true; } -function attachLlmTooling( - config: LlmConfig, +function buildToolProfileConfig( + toolConfig: LlmToolConfig, toolConfigsByAlias: Map, mcpProvidersByName: Map, -): void { - const toolAlias = config.tool_alias?.trim(); - if (!toolAlias) { - config.tool_alias = ""; - config.tool_configs = []; - config.mcp_providers = []; - return; - } - const toolConfig = toolConfigsByAlias.get(toolAlias); - if (!toolConfig) { - config.tool_configs = []; - config.mcp_providers = []; - return; - } - config.tool_configs = [cloneToolConfig(toolConfig)]; - config.mcp_providers = toolConfig.providers - .map((providerName) => mcpProvidersByName.get(providerName)) - .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])); + id: string, +): ToolProfileConfig { + const canonical = toolConfigsByAlias.get(toolConfig.tool_alias) ?? toolConfig; + return { + id, + kind: "tool_config", + name: canonical.tool_alias, + // biome-ignore lint/style/useNamingConvention: ui schema + mcp_providers: canonical.providers + .map((providerName) => mcpProvidersByName.get(providerName)) + .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])), + // biome-ignore lint/style/useNamingConvention: api schema + allow_tools: [...(canonical.allow_tools ?? [])], + // biome-ignore lint/style/useNamingConvention: api schema + max_tool_call_turns: canonical.max_tool_call_turns ?? "5", + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec: canonical.timeout_sec ?? "", + }; } export function importRecipePayload(input: string): ImportResult { @@ -471,6 +463,23 @@ export function importRecipePayload(input: string): ImportResult { }); } + for (const toolConfig of toolConfigsByAlias.values()) { + const id = `n${nextId}`; + nextId += 1; + const config = buildToolProfileConfig( + toolConfig, + toolConfigsByAlias, + mcpProvidersByName, + id, + ); + if (nameToId.has(config.name)) { + errors.push(`Duplicate column name: ${config.name}.`); + continue; + } + nameToId.set(config.name, config.id); + configs.push(config); + } + recipe.columns.forEach((column, index) => { if (!isRecord(column)) { errors.push(`Column ${index + 1}: invalid object.`); @@ -482,9 +491,6 @@ export function importRecipePayload(input: string): ImportResult { if (!config) { return; } - if (config.kind === "llm") { - attachLlmTooling(config, toolConfigsByAlias, mcpProvidersByName); - } applyAdvancedOpen(config, uiAdvancedOpenByNode); if (nameToId.has(config.name)) { errors.push(`Duplicate column name: ${config.name}.`); diff --git a/studio/frontend/src/features/recipe-studio/utils/index.ts b/studio/frontend/src/features/recipe-studio/utils/index.ts index fa03a50e40..a30db80f98 100644 --- a/studio/frontend/src/features/recipe-studio/utils/index.ts +++ b/studio/frontend/src/features/recipe-studio/utils/index.ts @@ -6,6 +6,7 @@ export { makeModelProviderConfig, makeSamplerConfig, makeSeedConfig, + makeToolProfileConfig, makeValidatorConfig, } from "./config-factories"; export { diff --git a/studio/frontend/src/features/recipe-studio/utils/node-data.ts b/studio/frontend/src/features/recipe-studio/utils/node-data.ts index c887a979c3..195a5f03bb 100644 --- a/studio/frontend/src/features/recipe-studio/utils/node-data.ts +++ b/studio/frontend/src/features/recipe-studio/utils/node-data.ts @@ -97,6 +97,17 @@ export function nodeDataFromConfig( layoutDirection, }; } + if (config.kind === "tool_config") { + const providerCount = config.mcp_providers.length; + return { + title: "Tool Profile", + kind: "tool_config", + subtype: providerCount === 1 ? "1 MCP server" : `${providerCount} MCP servers`, + blockType: "tool_config", + name: config.name, + layoutDirection, + }; + } return { title: "LLM", kind: "llm", diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts index f037586af7..615d48b739 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts @@ -24,14 +24,13 @@ import { readNodeWidth } from "../rf-node-dimensions"; import { buildExpressionColumn, buildLlmColumn, - buildLlmMcpProvider, - buildLlmToolConfig, buildModelConfig, buildModelProvider, buildProcessors, buildSamplerColumn, buildSeedConfig, buildSeedDropProcessor, + buildToolProfilePayload, buildValidatorColumn, pickFirstSeedConfig, } from "./builders"; @@ -169,34 +168,6 @@ export function buildRecipePayload( } } columns.push(buildLlmColumn(config, errors)); - for (const provider of config.mcp_providers ?? []) { - const builtProvider = buildLlmMcpProvider(provider, errors); - if (!builtProvider) { - continue; - } - pushUniqueJson( - "MCP provider", - String(builtProvider.name), - builtProvider, - mcpProviderJsonByName, - mcpProviders, - errors, - ); - } - for (const toolConfig of config.tool_configs ?? []) { - const builtToolConfig = buildLlmToolConfig(toolConfig, errors); - if (!builtToolConfig) { - continue; - } - pushUniqueJson( - "Tool config", - String(builtToolConfig.tool_alias), - builtToolConfig, - toolConfigJsonByAlias, - toolConfigs, - errors, - ); - } if (config.model_alias) { modelAliases.add(config.model_alias); } @@ -230,6 +201,30 @@ export function buildRecipePayload( modelProviderConfigs.push(config); continue; } + if (config.kind === "tool_config") { + const built = buildToolProfilePayload(config, errors); + for (const provider of built.mcp_providers) { + pushUniqueJson( + "MCP provider", + String(provider.name), + provider, + mcpProviderJsonByName, + mcpProviders, + errors, + ); + } + if (built.tool_config) { + pushUniqueJson( + "Tool config", + String(built.tool_config.tool_alias), + built.tool_config, + toolConfigJsonByAlias, + toolConfigs, + errors, + ); + } + continue; + } modelConfigs.push(buildModelConfig(config, errors)); modelConfigConfigs.push(config); } diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts index f142a63701..f805a046f7 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders-llm.ts @@ -1,4 +1,9 @@ -import type { LlmConfig, LlmMcpProviderConfig, LlmToolConfig } from "../../types"; +import type { + LlmConfig, + LlmMcpProviderConfig, + LlmToolConfig, + ToolProfileConfig, +} from "../../types"; function buildImageContext( config: LlmConfig, @@ -200,3 +205,40 @@ export function buildLlmToolConfig( timeout_sec: timeoutSec, }; } + +export function buildToolProfilePayload( + config: ToolProfileConfig, + errors: string[], +): { + // biome-ignore lint/style/useNamingConvention: api schema + mcp_providers: Record[]; + // biome-ignore lint/style/useNamingConvention: api schema + tool_config: Record | null; +} { + const mcpProviders = config.mcp_providers + .map((provider) => buildLlmMcpProvider(provider, errors)) + .flatMap((provider) => (provider ? [provider] : [])); + const toolConfig = buildLlmToolConfig( + { + id: config.id, + // biome-ignore lint/style/useNamingConvention: api schema + tool_alias: config.name, + providers: mcpProviders + .map((provider) => String(provider.name ?? "").trim()) + .filter(Boolean), + // biome-ignore lint/style/useNamingConvention: api schema + allow_tools: config.allow_tools, + // biome-ignore lint/style/useNamingConvention: api schema + max_tool_call_turns: config.max_tool_call_turns, + // biome-ignore lint/style/useNamingConvention: api schema + timeout_sec: config.timeout_sec, + }, + errors, + ); + return { + // biome-ignore lint/style/useNamingConvention: api schema + mcp_providers: mcpProviders, + // biome-ignore lint/style/useNamingConvention: api schema + tool_config: toolConfig, + }; +} diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/builders.ts b/studio/frontend/src/features/recipe-studio/utils/payload/builders.ts index fb756cd159..b963d6bc0b 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/builders.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/builders.ts @@ -1,4 +1,9 @@ -export { buildLlmColumn, buildLlmMcpProvider, buildLlmToolConfig } from "./builders-llm"; +export { + buildLlmColumn, + buildLlmMcpProvider, + buildLlmToolConfig, + buildToolProfilePayload, +} from "./builders-llm"; export { buildModelConfig, buildModelProvider } from "./builders-model"; export { buildExpressionColumn, buildProcessors } from "./builders-processors"; export { buildSamplerColumn } from "./builders-sampler"; diff --git a/studio/frontend/src/features/recipe-studio/utils/recipe-studio-view.ts b/studio/frontend/src/features/recipe-studio/utils/recipe-studio-view.ts index b75dc454ba..d16175933e 100644 --- a/studio/frontend/src/features/recipe-studio/utils/recipe-studio-view.ts +++ b/studio/frontend/src/features/recipe-studio/utils/recipe-studio-view.ts @@ -4,6 +4,7 @@ export type DialogOptions = { categoryOptions: SamplerConfig[]; modelConfigAliases: string[]; modelProviderOptions: string[]; + toolProfileAliases: string[]; datetimeOptions: string[]; }; @@ -11,6 +12,7 @@ export function buildDialogOptions(configList: NodeConfig[]): DialogOptions { const categoryOptions: SamplerConfig[] = []; const modelConfigAliases: string[] = []; const modelProviderOptions: string[] = []; + const toolProfileAliases: string[] = []; const datetimeOptions: string[] = []; for (const config of configList) { @@ -29,6 +31,10 @@ export function buildDialogOptions(configList: NodeConfig[]): DialogOptions { } if (config.kind === "model_provider") { modelProviderOptions.push(config.name); + continue; + } + if (config.kind === "tool_config") { + toolProfileAliases.push(config.name); } } @@ -36,6 +42,7 @@ export function buildDialogOptions(configList: NodeConfig[]): DialogOptions { categoryOptions, modelConfigAliases, modelProviderOptions, + toolProfileAliases, datetimeOptions, }; } diff --git a/studio/frontend/src/features/recipe-studio/utils/validation.ts b/studio/frontend/src/features/recipe-studio/utils/validation.ts index b7805ce0c3..7aa755971a 100644 --- a/studio/frontend/src/features/recipe-studio/utils/validation.ts +++ b/studio/frontend/src/features/recipe-studio/utils/validation.ts @@ -195,6 +195,44 @@ export function getConfigErrors(config: NodeConfig | null): string[] { errors.push("Expression is required."); } } + if (config.kind === "tool_config") { + if (config.mcp_providers.length === 0) { + errors.push("Add at least one MCP server."); + } + const serverNames = new Set(); + for (const provider of config.mcp_providers) { + const name = provider.name.trim(); + if (!name) { + errors.push("Each MCP server needs a name."); + continue; + } + if (serverNames.has(name)) { + errors.push(`Duplicate MCP server name: ${name}.`); + } + serverNames.add(name); + if (provider.provider_type === "stdio") { + if (!provider.command?.trim()) { + errors.push(`MCP server ${name}: command is required.`); + } + } else if (!provider.endpoint?.trim()) { + errors.push(`MCP server ${name}: endpoint is required.`); + } + } + const maxTurnsRaw = config.max_tool_call_turns?.trim(); + if ( + maxTurnsRaw && + (!Number.isFinite(Number(maxTurnsRaw)) || Number(maxTurnsRaw) < 1) + ) { + errors.push("Max tool call turns must be >= 1."); + } + const timeoutRaw = config.timeout_sec?.trim(); + if ( + timeoutRaw && + (!Number.isFinite(Number(timeoutRaw)) || Number(timeoutRaw) <= 0) + ) { + errors.push("Timeout must be > 0."); + } + } if (config.kind === "validator") { const targets = (config.target_columns ?? []) .map((value) => value.trim()) diff --git a/studio/frontend/src/features/recipe-studio/utils/variables.ts b/studio/frontend/src/features/recipe-studio/utils/variables.ts index f98f0662bc..7a34ad7006 100644 --- a/studio/frontend/src/features/recipe-studio/utils/variables.ts +++ b/studio/frontend/src/features/recipe-studio/utils/variables.ts @@ -29,7 +29,11 @@ export function getAvailableVariableEntries( if (config.id === currentId) { continue; } - if (config.kind === "model_provider" || config.kind === "model_config") { + if ( + config.kind === "model_provider" || + config.kind === "model_config" || + config.kind === "tool_config" + ) { continue; } From 542d9126cc5ead80b3c04535c43d99779d54d25c Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 17:27:02 +0100 Subject: [PATCH 6/7] chore(data-recipe): bump data-designer to 0.5.2 and pin duckdb<1.5 --- studio/backend/core/data_recipe/service.py | 1 + .../backend/requirements/single-env/data-designer-deps.txt | 3 ++- studio/backend/requirements/single-env/data-designer.txt | 6 +++--- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/studio/backend/core/data_recipe/service.py b/studio/backend/core/data_recipe/service.py index f2c25e62f8..01855f5b1d 100644 --- a/studio/backend/core/data_recipe/service.py +++ b/studio/backend/core/data_recipe/service.py @@ -184,6 +184,7 @@ def build_mcp_providers( MCPProvider( name=str(provider.get("name", "")), endpoint=str(provider.get("endpoint", "")), + provider_type=str(provider_type), api_key=str(api_key) if api_key else None, ) ) diff --git a/studio/backend/requirements/single-env/data-designer-deps.txt b/studio/backend/requirements/single-env/data-designer-deps.txt index cbf8856073..dfc5b9bf21 100644 --- a/studio/backend/requirements/single-env/data-designer-deps.txt +++ b/studio/backend/requirements/single-env/data-designer-deps.txt @@ -1,6 +1,7 @@ # Data Designer runtime deps installed explicitly (single-env mode). +# DuckDB 1.5 removed Relation.record_batch(); keep <1.5 until upstream ships the fix. anyascii<1,>=0.3.3 -duckdb<2,>=1.1.3 +duckdb<1.5,>=1.1.3 faker<21,>=20.1.0 httpx<1,>=0.27.2 httpx-retries<1,>=0.4.2 diff --git a/studio/backend/requirements/single-env/data-designer.txt b/studio/backend/requirements/single-env/data-designer.txt index c5ddcddc36..8daa1eca43 100644 --- a/studio/backend/requirements/single-env/data-designer.txt +++ b/studio/backend/requirements/single-env/data-designer.txt @@ -1,5 +1,5 @@ # Install Data Designer in same env as Unsloth. -data-designer==0.5.1 -data-designer-config==0.5.1 -data-designer-engine==0.5.1 +data-designer==0.5.2 +data-designer-config==0.5.2 +data-designer-engine==0.5.2 prompt-toolkit>=3,<4 From 1fe8995f1c76a4e3b3b4ce9f94716e39b3f65125 Mon Sep 17 00:00:00 2001 From: Shine1i Date: Mon, 9 Mar 2026 19:19:14 +0100 Subject: [PATCH 7/7] feat(recipe-studio): add support for managing tools by provider in tool profiles --- .../tool-profile/tool-profile-dialog.tsx | 42 +++++++++++++---- .../hooks/use-recipe-persistence.ts | 10 +++++ .../src/features/recipe-studio/types/index.ts | 2 + .../recipe-studio/utils/config-factories.ts | 2 + .../recipe-studio/utils/import/importer.ts | 45 +++++++++++++++++++ .../utils/payload/build-payload.ts | 25 +++++++++++ .../recipe-studio/utils/payload/types.ts | 3 +- 7 files changed, 119 insertions(+), 10 deletions(-) diff --git a/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx b/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx index d625c7828d..76a8e5b449 100644 --- a/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx +++ b/studio/frontend/src/features/recipe-studio/dialogs/tool-profile/tool-profile-dialog.tsx @@ -14,7 +14,7 @@ import { PlusSignIcon, } from "@hugeicons/core-free-icons"; import { HugeiconsIcon } from "@hugeicons/react"; -import { type ReactElement, useEffect, useMemo, useState } from "react"; +import { type ReactElement, useEffect, useMemo, useRef, useState } from "react"; import { listMcpTools } from "../../api"; import { ChipInput } from "../../components/chip-input"; import type { LlmMcpProviderConfig, McpEnvVar, ToolProfileConfig } from "../../types"; @@ -351,11 +351,12 @@ export function ToolProfileDialog({ const providers = config.mcp_providers; const [loadingTools, setLoadingTools] = useState(false); const [toolsByProvider, setToolsByProvider] = useState>( - {}, + config.fetched_tools_by_provider ?? {}, ); const [providerErrors, setProviderErrors] = useState>({}); const [duplicateTools, setDuplicateTools] = useState>({}); const [openProviders, setOpenProviders] = useState>({}); + const previousProviderSignatureRef = useRef(null); const providerSignature = useMemo( () => @@ -378,10 +379,30 @@ export function ToolProfileDialog({ ); useEffect(() => { + const previousSignature = previousProviderSignatureRef.current; + previousProviderSignatureRef.current = providerSignature; + if (previousSignature === null) { + setToolsByProvider(config.fetched_tools_by_provider ?? {}); + return; + } + if (previousSignature === providerSignature) { + return; + } setToolsByProvider({}); setProviderErrors({}); setDuplicateTools({}); - }, [providerSignature]); + if (Object.keys(config.fetched_tools_by_provider ?? {}).length > 0) { + onUpdate({ + // biome-ignore lint/style/useNamingConvention: ui schema + fetched_tools_by_provider: {}, + }); + } + }, [config.fetched_tools_by_provider, onUpdate, providerSignature]); + + useEffect(() => { + const tools = config.fetched_tools_by_provider ?? {}; + setToolsByProvider(tools); + }, [config.fetched_tools_by_provider]); useEffect(() => { setOpenProviders((current) => { @@ -523,13 +544,16 @@ export function ToolProfileDialog({ // biome-ignore lint/style/useNamingConvention: api schema timeout_sec: timeoutSec, }); - setToolsByProvider( - Object.fromEntries( - response.providers - .filter((provider) => provider.name.trim()) - .map((provider) => [provider.name.trim(), provider.tools]), - ), + const nextToolsByProvider = Object.fromEntries( + response.providers + .filter((provider) => provider.name.trim()) + .map((provider) => [provider.name.trim(), provider.tools]), ); + setToolsByProvider(nextToolsByProvider); + onUpdate({ + // biome-ignore lint/style/useNamingConvention: ui schema + fetched_tools_by_provider: nextToolsByProvider, + }); setProviderErrors( Object.fromEntries( response.providers diff --git a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts index 099dfb2f6b..7ea6f8726e 100644 --- a/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts +++ b/studio/frontend/src/features/recipe-studio/hooks/use-recipe-persistence.ts @@ -61,6 +61,16 @@ function stripApiKeys(value: unknown): unknown { } output[key] = stripApiKeys(entry); } + if ( + output.provider_type === "stdio" && + output.env && + typeof output.env === "object" && + !Array.isArray(output.env) + ) { + output.env = Object.fromEntries( + Object.keys(output.env as Record).map((envKey) => [envKey, ""]), + ); + } return output; } diff --git a/studio/frontend/src/features/recipe-studio/types/index.ts b/studio/frontend/src/features/recipe-studio/types/index.ts index c62da8fd9f..4e0470db71 100644 --- a/studio/frontend/src/features/recipe-studio/types/index.ts +++ b/studio/frontend/src/features/recipe-studio/types/index.ts @@ -181,6 +181,8 @@ export type ToolProfileConfig = { name: string; // biome-ignore lint/style/useNamingConvention: ui schema mcp_providers: LlmMcpProviderConfig[]; + // biome-ignore lint/style/useNamingConvention: ui schema + fetched_tools_by_provider?: Record; // biome-ignore lint/style/useNamingConvention: api schema allow_tools?: string[]; // biome-ignore lint/style/useNamingConvention: api schema diff --git a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts index d605e1a775..697a819ecf 100644 --- a/studio/frontend/src/features/recipe-studio/utils/config-factories.ts +++ b/studio/frontend/src/features/recipe-studio/utils/config-factories.ts @@ -275,6 +275,8 @@ export function makeToolProfileConfig( name: nextName(existing, "tools"), // biome-ignore lint/style/useNamingConvention: ui schema mcp_providers: [], + // biome-ignore lint/style/useNamingConvention: ui schema + fetched_tools_by_provider: {}, // biome-ignore lint/style/useNamingConvention: api schema allow_tools: [], // biome-ignore lint/style/useNamingConvention: api schema diff --git a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts index b8823882f7..41460fb7bf 100644 --- a/studio/frontend/src/features/recipe-studio/utils/import/importer.ts +++ b/studio/frontend/src/features/recipe-studio/utils/import/importer.ts @@ -252,6 +252,46 @@ function parseUiMarkdownNoteNodes(input: unknown): UiMarkdownNoteNode[] { return noteNodes; } +function parseUiToolProfileNodes(input: unknown): Map> { + const toolProfiles = new Map>(); + if (!Array.isArray(input)) { + return toolProfiles; + } + for (const node of input) { + if (!isRecord(node)) { + continue; + } + const nodeType = readString(node.node_type) ?? readString(node.type); + if (nodeType !== "tool_config") { + continue; + } + const name = readString(node.name) ?? readString(node.id); + if (!name?.trim()) { + continue; + } + const rawToolsByProvider = isRecord(node.tools_by_provider) + ? node.tools_by_provider + : null; + if (!rawToolsByProvider) { + continue; + } + const toolsByProvider = Object.fromEntries( + Object.entries(rawToolsByProvider).flatMap(([providerName, tools]) => { + const trimmedName = providerName.trim(); + if (!trimmedName || !Array.isArray(tools)) { + return []; + } + const values = Array.from( + new Set(tools.map((value) => String(value).trim()).filter(Boolean)), + ); + return values.length > 0 ? [[trimmedName, values]] : []; + }), + ); + toolProfiles.set(name.trim(), toolsByProvider); + } + return toolProfiles; +} + function parseAdvancedOpenByNode(input: unknown): Record { if (!isRecord(input)) { return {}; @@ -292,6 +332,7 @@ function buildToolProfileConfig( toolConfig: LlmToolConfig, toolConfigsByAlias: Map, mcpProvidersByName: Map, + fetchedToolsByProfileName: Map>, id: string, ): ToolProfileConfig { const canonical = toolConfigsByAlias.get(toolConfig.tool_alias) ?? toolConfig; @@ -303,6 +344,8 @@ function buildToolProfileConfig( mcp_providers: canonical.providers .map((providerName) => mcpProvidersByName.get(providerName)) .flatMap((provider) => (provider ? [cloneMcpProvider(provider)] : [])), + // biome-ignore lint/style/useNamingConvention: ui schema + fetched_tools_by_provider: fetchedToolsByProfileName.get(canonical.tool_alias) ?? {}, // biome-ignore lint/style/useNamingConvention: api schema allow_tools: [...(canonical.allow_tools ?? [])], // biome-ignore lint/style/useNamingConvention: api schema @@ -370,6 +413,7 @@ export function importRecipePayload(input: string): ImportResult { ); const uiAdvancedOpenByNode = parseAdvancedOpenByNode(ui?.advanced_open_by_node); const uiMarkdownNotes = parseUiMarkdownNoteNodes(ui?.nodes); + const uiToolProfilesByName = parseUiToolProfileNodes(ui?.nodes); for (const note of uiMarkdownNotes) { const id = `n${nextId}`; @@ -470,6 +514,7 @@ export function importRecipePayload(input: string): ImportResult { toolConfig, toolConfigsByAlias, mcpProvidersByName, + uiToolProfilesByName, id, ); if (nameToId.has(config.name)) { diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts index 615d48b739..758fcd4986 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/build-payload.ts @@ -267,6 +267,31 @@ export function buildRecipePayload( }, ]; } + if (config.kind === "tool_config") { + const toolsByProvider = Object.fromEntries( + Object.entries(config.fetched_tools_by_provider ?? {}).flatMap( + ([providerName, tools]) => { + const name = providerName.trim(); + const values = Array.from( + new Set(tools.map((tool) => tool.trim()).filter(Boolean)), + ); + return name && values.length > 0 ? [[name, values]] : []; + }, + ), + ); + return [ + { + id: config.name, + x: node.position.x, + y: node.position.y, + ...(width !== null ? { width } : {}), + node_type: "tool_config" as const, + ...(Object.keys(toolsByProvider).length > 0 && { + tools_by_provider: toolsByProvider, + }), + }, + ]; + } return [ { id: config.name, diff --git a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts index e8e207a3e7..69db100d96 100644 --- a/studio/frontend/src/features/recipe-studio/utils/payload/types.ts +++ b/studio/frontend/src/features/recipe-studio/utils/payload/types.ts @@ -37,11 +37,12 @@ export type RecipePayload = { x: number; y: number; width?: number; - node_type?: "markdown_note"; + node_type?: "markdown_note" | "tool_config"; name?: string; markdown?: string; note_color?: string; note_opacity?: string; + tools_by_provider?: Record; }>; edges: { from: string;