Skip to content

Commit 9eb285d

Browse files
taylordowns2000elias-ba
taylordowns2000
and
elias-ba
authored
Optionally include scrubbed I/O data in AI chat (#4146)
* initial commit * clean up * guards and disabling checkboxes * dialyzer * alpha sort * Extract build_message_options to reduce duplication Consolidate the duplicated message options map construction that existed in both extract_session_options and extract_message_options into a single build_message_options helper function. --------- Co-authored-by: Elias W. BA <eliaswalyba@gmail.com>
1 parent 049e947 commit 9eb285d

File tree

19 files changed

+1141
-75
lines changed

19 files changed

+1141
-75
lines changed

‎assets/js/collaborative-editor/components/AIAssistantPanel.tsx‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@ import {
1010
useAIHasCompletedSessionLoad,
1111
useAISessionListCommands,
1212
} from '../hooks/useAIAssistant';
13+
import { useSelectedStepId, useSelectedRunId } from '../hooks/useHistory';
1314

1415
import { ChatInput } from './ChatInput';
1516
import { DisclaimerScreen } from './DisclaimerScreen';
@@ -62,6 +63,8 @@ interface AIAssistantPanelProps {
6263
interface MessageOptions {
6364
attach_code?: boolean;
6465
attach_logs?: boolean;
66+
attach_io_data?: boolean;
67+
step_id?: string;
6568
}
6669

6770
/**
@@ -112,6 +115,8 @@ export function AIAssistantPanel({
112115
const hasSessionContext = useAIHasSessionContext();
113116
const hasCompletedSessionLoad = useAIHasCompletedSessionLoad();
114117
const { loadSessionList } = useAISessionListCommands();
118+
const selectedStepId = useSelectedStepId();
119+
const selectedRunId = useSelectedRunId();
115120

116121
useEffect(() => {
117122
if (prevViewRef.current !== view) {
@@ -412,6 +417,9 @@ export function AIAssistantPanel({
412417
focusTrigger={(focusTrigger ?? 0) + internalFocusTrigger}
413418
placeholder={placeholderText}
414419
disabledMessage={disabledMessage}
420+
selectedStepId={selectedStepId}
421+
selectedRunId={selectedRunId}
422+
selectedJobId={jobCodeContext?.job_id ?? null}
415423
/>
416424

417425
{/* About AI Assistant Modal */}

‎assets/js/collaborative-editor/components/AIAssistantPanelWrapper.tsx‎

Lines changed: 16 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -359,7 +359,12 @@ export function AIAssistantPanelWrapper() {
359359
const sendMessage = useCallback(
360360
(
361361
content: string,
362-
messageOptions?: { attach_code?: boolean; attach_logs?: boolean }
362+
messageOptions?: {
363+
attach_code?: boolean;
364+
attach_logs?: boolean;
365+
attach_io_data?: boolean;
366+
step_id?: string;
367+
}
363368
) => {
364369
const currentState = aiStore.getSnapshot();
365370

@@ -374,6 +379,8 @@ export function AIAssistantPanelWrapper() {
374379
// Include attach_code/attach_logs so backend knows to include them in first message
375380
...(messageOptions?.attach_code && { attach_code: true }),
376381
...(messageOptions?.attach_logs && { attach_logs: true }),
382+
...(messageOptions?.attach_io_data && { attach_io_data: true }),
383+
...(messageOptions?.step_id && { step_id: messageOptions.step_id }),
377384
};
378385

379386
// Add workflow YAML if in workflow mode
@@ -411,9 +418,15 @@ export function AIAssistantPanelWrapper() {
411418

412419
// For existing sessions, prepare options and send
413420
let options:
414-
| { attach_code?: boolean; attach_logs?: boolean; code?: string }
421+
| {
422+
attach_code?: boolean;
423+
attach_logs?: boolean;
424+
attach_io_data?: boolean;
425+
step_id?: string;
426+
code?: string;
427+
}
415428
| undefined = {
416-
...messageOptions, // Include attach_code and attach_logs
429+
...messageOptions, // Include attach_code, attach_logs, attach_io_data, step_id
417430
};
418431

419432
if (currentState.sessionType === 'workflow_template') {

‎assets/js/collaborative-editor/components/ChatInput.tsx‎

Lines changed: 190 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -21,11 +21,19 @@ interface ChatInputProps {
2121
placeholder?: string | undefined;
2222
/** Message to show in tooltip when input is disabled */
2323
disabledMessage?: string | undefined;
24+
/** Selected step ID for attaching I/O data */
25+
selectedStepId?: string | null;
26+
/** Selected run ID for attaching logs */
27+
selectedRunId?: string | null;
28+
/** Selected job ID for attaching code */
29+
selectedJobId?: string | null;
2430
}
2531

2632
interface MessageOptions {
2733
attach_code?: boolean;
2834
attach_logs?: boolean;
35+
attach_io_data?: boolean;
36+
step_id?: string;
2937
}
3038

3139
const MIN_TEXTAREA_HEIGHT = 52;
@@ -40,6 +48,9 @@ export function ChatInput({
4048
focusTrigger,
4149
placeholder = 'Ask me anything...',
4250
disabledMessage,
51+
selectedStepId,
52+
selectedRunId,
53+
selectedJobId,
4354
}: ChatInputProps) {
4455
const [input, setInput] = useState('');
4556

@@ -69,6 +80,19 @@ export function ChatInput({
6980
}
7081
});
7182

83+
const [attachIoData, setAttachIoData] = useState(() => {
84+
if (!storageKey) {
85+
return false;
86+
}
87+
try {
88+
const key = `${storageKey}:attach-io-data`;
89+
const saved = localStorage.getItem(key);
90+
return saved === 'true';
91+
} catch {
92+
return false;
93+
}
94+
});
95+
7296
const textareaRef = useRef<HTMLTextAreaElement>(null);
7397
const isLoadingFromStorageRef = useRef(false);
7498

@@ -107,6 +131,15 @@ export function ChatInput({
107131
// Ignore localStorage errors
108132
}
109133

134+
try {
135+
const ioDataKey = `${storageKey}:attach-io-data`;
136+
const savedIoData = localStorage.getItem(ioDataKey);
137+
const ioDataValue = savedIoData === 'true';
138+
setAttachIoData(ioDataValue);
139+
} catch {
140+
// Ignore localStorage errors
141+
}
142+
110143
setTimeout(() => {
111144
isLoadingFromStorageRef.current = false;
112145
}, 0);
@@ -132,6 +165,19 @@ export function ChatInput({
132165
}
133166
}, [attachLogs, storageKey]);
134167

168+
useEffect(() => {
169+
if (!storageKey) return;
170+
if (isLoadingFromStorageRef.current) return;
171+
try {
172+
localStorage.setItem(
173+
`${storageKey}:attach-io-data`,
174+
String(attachIoData)
175+
);
176+
} catch {
177+
// Ignore localStorage errors
178+
}
179+
}, [attachIoData, storageKey]);
180+
135181
useEffect(() => {
136182
if (enableAutoFocus && textareaRef.current) {
137183
const timeoutId = setTimeout(() => {
@@ -155,8 +201,16 @@ export function ChatInput({
155201

156202
const options: MessageOptions = {};
157203
if (showJobControls) {
158-
options.attach_code = attachCode;
159-
options.attach_logs = attachLogs;
204+
if (selectedRunId) {
205+
options.attach_logs = attachLogs;
206+
}
207+
if (selectedJobId) {
208+
options.attach_code = attachCode;
209+
}
210+
if (selectedStepId) {
211+
options.attach_io_data = attachIoData;
212+
options.step_id = selectedStepId;
213+
}
160214
}
161215

162216
onSendMessage?.(input.trim(), options);
@@ -219,31 +273,140 @@ export function ChatInput({
219273
<div className="flex items-center gap-3">
220274
{showJobControls ? (
221275
<>
222-
<label className="flex items-center gap-1.5 cursor-pointer group">
223-
<input
224-
type="checkbox"
225-
checked={attachCode}
226-
onChange={e => setAttachCode(e.target.checked)}
227-
className="w-3.5 h-3.5 rounded border-gray-300 text-primary-600
228-
focus:ring-primary-500 focus:ring-offset-0 cursor-pointer"
229-
/>
230-
<span className="text-[11px] font-medium text-gray-600 group-hover:text-gray-900">
231-
Include job code
232-
</span>
233-
</label>
234-
235-
<label className="flex items-center gap-1.5 cursor-pointer group">
236-
<input
237-
type="checkbox"
238-
checked={attachLogs}
239-
onChange={e => setAttachLogs(e.target.checked)}
240-
className="w-3.5 h-3.5 rounded border-gray-300 text-primary-600
241-
focus:ring-primary-500 focus:ring-offset-0 cursor-pointer"
242-
/>
243-
<span className="text-[11px] font-medium text-gray-600 group-hover:text-gray-900">
244-
Include run logs
245-
</span>
246-
</label>
276+
<Tooltip
277+
content={
278+
selectedJobId
279+
? undefined
280+
: 'Select a job to include code'
281+
}
282+
side="top"
283+
>
284+
<label
285+
className={cn(
286+
'flex items-center gap-1.5 group',
287+
selectedJobId
288+
? 'cursor-pointer'
289+
: 'cursor-not-allowed opacity-50'
290+
)}
291+
>
292+
<input
293+
type="checkbox"
294+
// NOTE: Regardless of preferences, we show it
295+
// unchecked if no job is selected because code
296+
// can't be sent without a job
297+
checked={attachCode && !!selectedJobId}
298+
onChange={e => setAttachCode(e.target.checked)}
299+
disabled={!selectedJobId}
300+
className={cn(
301+
'w-3.5 h-3.5 rounded border-gray-300 text-primary-600',
302+
'focus:ring-primary-500 focus:ring-offset-0',
303+
selectedJobId
304+
? 'cursor-pointer'
305+
: 'cursor-not-allowed'
306+
)}
307+
/>
308+
<span
309+
className={cn(
310+
'text-[11px] font-medium',
311+
selectedJobId
312+
? 'text-gray-600 group-hover:text-gray-900'
313+
: 'text-gray-400'
314+
)}
315+
>
316+
Send code
317+
</span>
318+
</label>
319+
</Tooltip>
320+
321+
<Tooltip
322+
content={
323+
selectedRunId
324+
? undefined
325+
: 'Select a run to include logs'
326+
}
327+
side="top"
328+
>
329+
<label
330+
className={cn(
331+
'flex items-center gap-1.5 group',
332+
selectedRunId
333+
? 'cursor-pointer'
334+
: 'cursor-not-allowed opacity-50'
335+
)}
336+
>
337+
<input
338+
type="checkbox"
339+
// NOTE: Regardless of preferences, we show it
340+
// unchecked if no run is selected because logs
341+
// can't be sent without a run
342+
checked={attachLogs && !!selectedRunId}
343+
onChange={e => setAttachLogs(e.target.checked)}
344+
disabled={!selectedRunId}
345+
className={cn(
346+
'w-3.5 h-3.5 rounded border-gray-300 text-primary-600',
347+
'focus:ring-primary-500 focus:ring-offset-0',
348+
selectedRunId
349+
? 'cursor-pointer'
350+
: 'cursor-not-allowed'
351+
)}
352+
/>
353+
<span
354+
className={cn(
355+
'text-[11px] font-medium',
356+
selectedRunId
357+
? 'text-gray-600 group-hover:text-gray-900'
358+
: 'text-gray-400'
359+
)}
360+
>
361+
Send logs
362+
</span>
363+
</label>
364+
</Tooltip>
365+
366+
<Tooltip
367+
content={
368+
selectedStepId
369+
? 'Include scrubbed I/O data structure (values removed)'
370+
: 'Select a step to include I/O data'
371+
}
372+
side="top"
373+
>
374+
<label
375+
className={cn(
376+
'flex items-center gap-1.5 group',
377+
selectedStepId
378+
? 'cursor-pointer'
379+
: 'cursor-not-allowed opacity-50'
380+
)}
381+
>
382+
<input
383+
type="checkbox"
384+
// NOTE: Regardless of preferences, we show it
385+
// unchecked if no step is selected because I/O
386+
// can't be sent without a step
387+
checked={attachIoData && !!selectedStepId}
388+
onChange={e => setAttachIoData(e.target.checked)}
389+
disabled={!selectedStepId}
390+
className={cn(
391+
'w-3.5 h-3.5 rounded border-gray-300 text-primary-600',
392+
'focus:ring-primary-500 focus:ring-offset-0',
393+
selectedStepId
394+
? 'cursor-pointer'
395+
: 'cursor-not-allowed'
396+
)}
397+
/>
398+
<span
399+
className={cn(
400+
'text-[11px] font-medium',
401+
selectedStepId
402+
? 'text-gray-600 group-hover:text-gray-900'
403+
: 'text-gray-400'
404+
)}
405+
>
406+
Send scrubbed I/O
407+
</span>
408+
</label>
409+
</Tooltip>
247410
</>
248411
) : (
249412
<div className="flex items-center gap-1.5">

‎assets/js/collaborative-editor/hooks/useHistory.ts‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -191,6 +191,17 @@ export const useSelectedStepId = (): string | null => {
191191
return useSyncExternalStore(historyStore.subscribe, selectStepId);
192192
};
193193

194+
/**
195+
* Hook to get the ID of the currently active run
196+
*/
197+
export const useSelectedRunId = (): string | null => {
198+
const historyStore = useHistoryStore();
199+
const selectRunId = historyStore.withSelector(
200+
state => state.activeRun?.id ?? null
201+
);
202+
return useSyncExternalStore(historyStore.subscribe, selectRunId);
203+
};
204+
194205
/**
195206
* Hook to get currently selected step (with lookup)
196207
*/

‎assets/js/collaborative-editor/lib/AIChannelRegistry.ts‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -700,6 +700,12 @@ export class AIChannelRegistry {
700700
if (context.attach_logs) {
701701
params['attach_logs'] = true;
702702
}
703+
if (context.attach_io_data) {
704+
params['attach_io_data'] = true;
705+
}
706+
if (context.step_id) {
707+
params['step_id'] = context.step_id;
708+
}
703709
} else {
704710
// WorkflowTemplateContext
705711
params['project_id'] = context.project_id;

‎assets/js/collaborative-editor/types/ai-assistant.ts‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,8 @@ export interface JobCodeContext {
5252
job_id: string;
5353
attach_code?: boolean;
5454
attach_logs?: boolean;
55+
attach_io_data?: boolean;
56+
step_id?: string;
5557
follow_run_id?: string;
5658
content?: string;
5759

@@ -169,6 +171,8 @@ export interface AIAssistantStore {
169171
export interface MessageOptions {
170172
attach_code?: boolean;
171173
attach_logs?: boolean;
174+
attach_io_data?: boolean;
175+
step_id?: string;
172176

173177
code?: string;
174178
errors?: string;

0 commit comments

Comments
 (0)