Skip to content

Commit 192f61d

Browse files
mscolnickmscolnick
authored
improvement: better UX for AI context items (#6649)
AI Context Improvements - Improved table context format: added sample values for columns and inline formatting - Styling of Resource's attached when using the agent (support text/plain, markdown, and images) - Reorganized completion sections with explicit ranking system (boosts dont apply across sections, so need to rank the sections) - Hide empty datasources from context menu Various UI fixes - Added hover states and improved visual hierarchy for resource and resource-link blocks - Implemented image preview popover for image resource links - Added overflow handling to prevent horizontal scrolling in user messages and tool notifications - Fixed truncation for long accordion titles and adjusted z-index for session tabs
1 parent 570a5b0 commit 192f61d

File tree

18 files changed

+208
-176
lines changed

18 files changed

+208
-176
lines changed

‎frontend/src/components/chat/acp/__tests__/context-utils.test.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -133,7 +133,7 @@ describe("parseContextFromPrompt", () => {
133133
type: "resource",
134134
resource: {
135135
uri: "context.md",
136-
mimeType: "text/markdown",
136+
mimeType: "text/plain",
137137
text: "formatted context",
138138
},
139139
});

‎frontend/src/components/chat/acp/agent-panel.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -867,7 +867,7 @@ const AgentPanel: React.FC = () => {
867867
type: "resource",
868868
resource: {
869869
uri: "marimo_rules.md",
870-
mimeType: "text/markdown",
870+
mimeType: "text/plain",
871871
text: getAgentPrompt(filename),
872872
},
873873
},

‎frontend/src/components/chat/acp/blocks.tsx‎

Lines changed: 46 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,6 @@ import {
3030
} from "@/components/ui/popover";
3131
import { logNever } from "@/utils/assertNever";
3232
import { cn } from "@/utils/cn";
33-
import { type Base64String, base64ToDataURL } from "@/utils/json/base64";
3433
import { Strings } from "@/utils/strings";
3534
import { MarkdownRenderer } from "../markdown-renderer";
3635
import { SimpleAccordion } from "./common";
@@ -289,7 +288,7 @@ export const PlansBlock = (props: { data: PlanNotificationEvent[] }) => {
289288

290289
export const UserMessagesBlock = (props: { data: UserNotificationEvent[] }) => {
291290
return (
292-
<div className="flex flex-col gap-2 text-muted-foreground border p-2 bg-background rounded break-words">
291+
<div className="flex flex-col gap-2 text-muted-foreground border p-2 bg-background rounded break-words overflow-x-hidden">
293292
<ContentBlocks data={props.data.map((item) => item.content)} />
294293
</div>
295294
);
@@ -364,58 +363,28 @@ export const ResourceBlock = (props: { data: ContentBlockOf<"resource"> }) => {
364363
return (
365364
<Popover>
366365
<PopoverTrigger>
367-
<span className="flex items-center gap-1">
366+
<span className="flex items-center gap-1 hover:bg-muted rounded-md px-1">
368367
{props.data.resource.mimeType && (
369368
<MimeIcon mimeType={props.data.resource.mimeType} />
370369
)}
371370
{props.data.resource.uri}
372371
</span>
373372
</PopoverTrigger>
374-
<PopoverContent className="max-h-96 overflow-y-auto scrollbar-thin">
375-
<MarkdownRenderer content={props.data.resource.text} />
373+
<PopoverContent className="max-h-96 overflow-y-auto scrollbar-thin whitespace-pre-wrap w-full max-w-[500px]">
374+
<span className="text-muted-foreground text-xs mb-1 italic">
375+
Formatted for agents, not humans.
376+
</span>
377+
{props.data.resource.mimeType === "text/plain" ? (
378+
<pre className="text-xs whitespace-pre-wrap p-2 bg-muted rounded-md break-words">
379+
{props.data.resource.text}
380+
</pre>
381+
) : (
382+
<MarkdownRenderer content={props.data.resource.text} />
383+
)}
376384
</PopoverContent>
377385
</Popover>
378386
);
379387
}
380-
381-
if ("blob" in props.data.resource) {
382-
if (props.data.resource.mimeType?.startsWith("image/")) {
383-
return (
384-
<ImageBlock
385-
data={{
386-
type: "image",
387-
mimeType: props.data.resource.mimeType,
388-
data: props.data.resource.blob,
389-
}}
390-
/>
391-
);
392-
}
393-
if (props.data.resource.mimeType?.startsWith("audio/")) {
394-
return (
395-
<AudioBlock
396-
data={{
397-
type: "audio",
398-
mimeType: props.data.resource.mimeType,
399-
data: props.data.resource.blob,
400-
}}
401-
/>
402-
);
403-
}
404-
const dataURL = base64ToDataURL(
405-
props.data.resource.blob as Base64String,
406-
props.data.resource.mimeType ?? "",
407-
);
408-
return (
409-
<a href={dataURL} className="flex items-center gap-1" download={true}>
410-
{props.data.resource.mimeType && (
411-
<MimeIcon mimeType={props.data.resource.mimeType} />
412-
)}
413-
{props.data.resource.uri}
414-
</a>
415-
);
416-
}
417-
logNever(props.data.resource);
418-
return null;
419388
};
420389

421390
export const ResourceLinkBlock = (props: {
@@ -427,38 +396,62 @@ export const ResourceLinkBlock = (props: {
427396
href={props.data.uri}
428397
target="_blank"
429398
rel="noopener noreferrer"
430-
className="text-link hover:underline"
399+
className="text-link hover:underline px-1"
431400
>
432401
{props.data.name}
433402
</a>
434403
);
435404
}
436405

406+
// Show image in popover for image mime types
407+
if (props.data.mimeType?.startsWith("image/")) {
408+
return (
409+
<div>
410+
<Popover>
411+
<PopoverTrigger>
412+
<span className="flex items-center gap-1 hover:bg-muted rounded-md px-1 cursor-pointer">
413+
<MimeIcon mimeType={props.data.mimeType} />
414+
{props.data.name || props.data.title || props.data.uri}
415+
</span>
416+
</PopoverTrigger>
417+
<PopoverContent className="w-auto max-w-[500px] p-2">
418+
<img
419+
src={props.data.uri}
420+
alt={props.data.name || props.data.title || "Image"}
421+
className="max-w-full max-h-96 object-contain"
422+
/>
423+
</PopoverContent>
424+
</Popover>
425+
</div>
426+
);
427+
}
428+
437429
return (
438-
<span className="flex items-center gap-1">
430+
<span className="flex items-center gap-1 px-1">
439431
{props.data.mimeType && <MimeIcon mimeType={props.data.mimeType} />}
440432
{props.data.name || props.data.title || props.data.uri}
441433
</span>
442434
);
443435
};
444436

445437
export const MimeIcon = (props: { mimeType: string }) => {
438+
const classNames = "h-2 w-2 flex-shrink-0";
446439
if (props.mimeType.startsWith("image/")) {
447-
return <FileImageIcon className="h-2 w-2" />;
440+
return <FileImageIcon className={classNames} />;
448441
}
449442
if (props.mimeType.startsWith("audio/")) {
450-
return <FileAudio2Icon className="h-2 w-2" />;
443+
return <FileAudio2Icon className={classNames} />;
451444
}
452445
if (props.mimeType.startsWith("video/")) {
453-
return <FileVideoCameraIcon className="h-2 w-2" />;
446+
return <FileVideoCameraIcon className={classNames} />;
454447
}
455448
if (props.mimeType.startsWith("text/")) {
456-
return <FileTextIcon className="h-2 w-2" />;
449+
return <FileTextIcon className={classNames} />;
457450
}
458451
if (props.mimeType.startsWith("application/")) {
459-
return <FileJsonIcon className="h-2 w-2" />;
452+
return <FileJsonIcon className={classNames} />;
460453
}
461-
return <FileIcon className="h-2 w-2" />;
454+
return <FileIcon className={classNames} />;
462455
};
463456

464457
export const SessionNotificationsBlock = <
@@ -533,7 +526,7 @@ export const ToolNotificationsBlock = (props: {
533526
const toolCalls = mergeToolCalls(props.data);
534527

535528
return (
536-
<div className="flex flex-col text-muted-foreground">
529+
<div className="flex flex-col text-muted-foreground overflow-x-hidden">
537530
{toolCalls.map((item) => (
538531
<SimpleAccordion
539532
key={item.toolCallId}

‎frontend/src/components/chat/acp/common.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -68,7 +68,7 @@ export const SimpleAccordion: React.FC<SimpleAccordionProps> = ({
6868
>
6969
<span className="flex items-center gap-1">
7070
{getStatusIcon()}
71-
<code className="font-mono text-xs">{title}</code>
71+
<code className="font-mono text-xs truncate">{title}</code>
7272
</span>
7373
</AccordionTrigger>
7474
<AccordionContent className="p-2">

‎frontend/src/components/chat/acp/context-utils.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,7 @@ export async function parseContextFromPrompt(
9393
type: "resource",
9494
resource: {
9595
uri: "context.md",
96-
mimeType: "text/markdown",
96+
mimeType: "text/plain",
9797
text: contextString,
9898
},
9999
});

‎frontend/src/components/chat/acp/session-tabs.tsx‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@ const SessionTab: React.FC<SessionTabProps> = memo(
3030
<div
3131
className={cn(
3232
"flex items-center gap-1 px-2 py-1 text-xs border-r border-border bg-muted/30 hover:bg-muted/50 cursor-pointer min-w-0",
33-
isActive && "bg-background border-b-0 relative z-10",
33+
isActive && "bg-background border-b-0 relative z-1",
3434
)}
3535
onClick={() => onSelect(session.tabId)}
3636
>

‎frontend/src/components/editor/ai/__tests__/completion-utils.test.ts‎

Lines changed: 23 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -55,15 +55,13 @@ describe("getAICompletionBody", () => {
5555
expect(result).toMatchInlineSnapshot(`
5656
{
5757
"context": {
58-
"plainText": "<data name="dataset1" source="unknown">
59-
Columns:
60-
- col1: number
61-
- col2: string</data>
62-
63-
<data name="dataset2" source="unknown">
64-
Columns:
65-
- col3: boolean
66-
- col4: date</data>",
58+
"plainText": "<data name="dataset1" source="unknown">Columns:
59+
col1 (number)
60+
col2 (string)</data>
61+
62+
<data name="dataset2" source="unknown">Columns:
63+
col3 (boolean)
64+
col4 (date)</data>",
6765
"schema": [],
6866
"variables": [],
6967
},
@@ -108,10 +106,9 @@ describe("getAICompletionBody", () => {
108106
expect(result).toMatchInlineSnapshot(`
109107
{
110108
"context": {
111-
"plainText": "<data name="existingDataset" source="unknown">
112-
Columns:
113-
- col1: number
114-
- col2: string</data>",
109+
"plainText": "<data name="existingDataset" source="unknown">Columns:
110+
col1 (number)
111+
col2 (string)</data>",
115112
"schema": [],
116113
"variables": [],
117114
},
@@ -144,14 +141,12 @@ describe("getAICompletionBody", () => {
144141
expect(result).toMatchInlineSnapshot(`
145142
{
146143
"context": {
147-
"plainText": "<data name="dataset.with.dots" source="unknown">
148-
Columns:
149-
- col1: number
150-
- col2: string</data>
151-
152-
<data name="regular_dataset" source="unknown">
153-
Columns:
154-
- col3: boolean</data>",
144+
"plainText": "<data name="dataset.with.dots" source="unknown">Columns:
145+
col1 (number)
146+
col2 (string)</data>
147+
148+
<data name="regular_dataset" source="unknown">Columns:
149+
col3 (boolean)</data>",
155150
"schema": [],
156151
"variables": [],
157152
},
@@ -198,9 +193,8 @@ describe("getAICompletionBody", () => {
198193
expect(result).toMatchInlineSnapshot(`
199194
{
200195
"context": {
201-
"plainText": "<data name="table1" source="unknown">
202-
Columns:
203-
- col1: number</data>",
196+
"plainText": "<data name="table1" source="unknown">Columns:
197+
col1 (number)</data>",
204198
"schema": [],
205199
"variables": [],
206200
},
@@ -276,10 +270,9 @@ describe("getAICompletionBody", () => {
276270
expect(result).toMatchInlineSnapshot(`
277271
{
278272
"context": {
279-
"plainText": "<data name="dataset1" source="unknown">
280-
Columns:
281-
- col1: number
282-
- col2: string</data>
273+
"plainText": "<data name="dataset1" source="unknown">Columns:
274+
col1 (number)
275+
col2 (string)</data>
283276
284277
<variable name="var1" dataType="string">"string value"</variable>",
285278
"schema": [],
@@ -346,9 +339,8 @@ describe("getAICompletionBody", () => {
346339
expect(result).toMatchInlineSnapshot(`
347340
{
348341
"context": {
349-
"plainText": "<data name="conflict" source="unknown">
350-
Columns:
351-
- col1: number</data>",
342+
"plainText": "<data name="conflict" source="unknown">Columns:
343+
col1 (number)</data>",
352344
"schema": [],
353345
"variables": [],
354346
},

‎frontend/src/core/ai/context/providers/__tests__/__snapshots__/tables.test.ts.snap‎

Lines changed: 13 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -1,34 +1,28 @@
11
// Vitest Snapshot v1, https://vitest.dev/guide/snapshot.html
22

33
exports[`TableContextProvider > formatContext > should format context for basic table > basic-table-context 1`] = `
4-
"<data name="products" source="memory">
5-
Shape: 100 rows, 3 columns
4+
"<data name="products" source="memory">Shape: 100 rows, 3 columns
65
Columns:
7-
- id: integer
8-
- name: string
9-
- active: boolean</data>"
6+
id (integer) - samples: [sample_id_1, sample_id_2]
7+
name (string) - samples: [sample_name_1, sample_name_2]
8+
active (boolean) - samples: [sample_active_1, sample_active_2]</data>"
109
`;
1110

1211
exports[`TableContextProvider > formatContext > should format context for remote database table > remote-table-context 1`] = `
13-
"<data name="remote_table" source="postgresql://localhost:5432/mydb">
14-
Shape: 100 rows, 3 columns
12+
"<data name="remote_table" source="postgresql://localhost:5432/mydb">Shape: 100 rows, 3 columns
1513
Columns:
16-
- uuid: string
17-
- created_at: string
18-
- metadata: string</data>"
14+
uuid (string) - samples: [sample_uuid_1, sample_uuid_2]
15+
created_at (string) - samples: [sample_created_at_1, sample_created_at_2]
16+
metadata (string) - samples: [sample_metadata_1, sample_metadata_2]</data>"
1917
`;
2018

21-
exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `
22-
"<data name="no_columns" source="memory">
23-
Shape: 100 rows, 3 columns</data>"
24-
`;
19+
exports[`TableContextProvider > formatContext > should format context for table without columns > no-columns-table-context 1`] = `"<data name="no_columns" source="memory">Shape: 100 rows, 3 columns</data>"`;
2520

2621
exports[`TableContextProvider > formatContext > should format context for table without shape info > no-shape-table-context 1`] = `
27-
"<data name="no_shape" source="memory">
28-
Columns:
29-
- id: integer
30-
- name: string
31-
- active: boolean</data>"
22+
"<data name="no_shape" source="memory">Columns:
23+
id (integer) - samples: [sample_id_1, sample_id_2]
24+
name (string) - samples: [sample_name_1, sample_name_2]
25+
active (boolean) - samples: [sample_active_1, sample_active_2]</data>"
3226
`;
3327

3428
exports[`TableContextProvider > getItems > should handle dataframe tables with variable names > dataframe-table 1`] = `

‎frontend/src/core/ai/context/providers/__tests__/cell-output.test.ts‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -187,7 +187,6 @@ describe("CellOutputContextProvider", () => {
187187
expect(completion.displayLabel).toBe(item.data.cellName);
188188
expect(completion.detail).toContain("output");
189189
expect(completion.type).toBe("cell-output");
190-
expect(completion.section).toBe("Cell Output");
191190
expect(typeof completion.info).toBe("function");
192191
});
193192
});

0 commit comments

Comments
 (0)