Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 13 additions & 9 deletions apps/dbagent/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -20,13 +20,13 @@
"dependencies": {
"@ai-sdk/anthropic": "^1.2.10",
"@ai-sdk/deepseek": "^0.2.13",
"@ai-sdk/google": "^1.2.13",
"@ai-sdk/google": "^1.2.14",
"@ai-sdk/openai": "^1.3.20",
"@ai-sdk/provider": "^1.1.3",
"@ai-sdk/react": "^1.2.9",
"@ai-sdk/ui-utils": "^1.2.8",
"@aws-sdk/client-cloudwatch": "^3.797.0",
"@aws-sdk/client-rds": "^3.797.0",
"@ai-sdk/react": "^1.2.10",
"@ai-sdk/ui-utils": "^1.2.9",
"@aws-sdk/client-cloudwatch": "^3.799.0",
"@aws-sdk/client-rds": "^3.799.0",
"@fluentui/react-icons": "^2.0.298",
"@google-cloud/logging": "^11.2.0",
"@google-cloud/monitoring": "^5.0.1",
Expand All @@ -36,9 +36,12 @@
"@mastra/core": "^0.9.0",
"@mastra/evals": "^0.1.19",
"@tailwindcss/postcss": "^4.1.4",
"@tanstack/react-query": "^5.74.7",
"@tanstack/react-query": "^5.74.11",
"@vercel/functions": "^2.0.0",
"ai": "^4.3.10",
"@xata.io/code-editor": "^0.0.6",
"@xata.io/components": "^0.0.7",
"@xata.io/theme": "^1.0.1",
"ai": "^4.3.11",
"bytes": "^3.1.2",
"canvas-confetti": "^1.9.3",
"class-variance-authority": "^0.7.1",
Expand Down Expand Up @@ -74,6 +77,7 @@
"react-data-grid": "7.0.0-beta.52",
"react-dom": "19.1.0",
"react-markdown": "^10.1.0",
"react-split-pane": "^0.1.92",
"react-syntax-highlighter": "^15.6.1",
"remark-gfm": "^4.0.1",
"server-only": "^0.0.1",
Expand All @@ -94,7 +98,7 @@
"@types/papaparse": "^5.3.15",
"@types/pg": "^8.11.14",
"@types/react": "^19.1.2",
"@types/react-dom": "^19.1.2",
"@types/react-dom": "^19.1.3",
"@types/react-syntax-highlighter": "^15.5.13",
"autoprefixer": "^10.4.21",
"dockerode": "^4.0.6",
Expand All @@ -105,7 +109,7 @@
"pg": "^8.15.6",
"postcss": "^8.5.3",
"tailwindcss": "^4.1.4",
"tsx": "^4.19.3"
"tsx": "^4.19.4"
},
"engines": {
"node": "22.x",
Expand Down
3 changes: 2 additions & 1 deletion apps/dbagent/src/components/chat/artifacts/artifact.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -13,13 +13,14 @@ import { fetcher } from '../utils';
import { ArtifactActions } from './artifact-actions';
import { ArtifactCloseButton } from './artifact-close-button';
import { ArtifactMessages } from './artifact-messages';
import { codeArtifact } from './code/client';
import { sheetArtifact } from './sheet/client';
import { textArtifact } from './text/client';
import { Toolbar } from './toolbar';
import { useArtifact } from './use-artifact';
import { VersionFooter } from './version-footer';

export const artifactDefinitions = [textArtifact, sheetArtifact];
export const artifactDefinitions = [textArtifact, sheetArtifact, codeArtifact];
export type ArtifactKind = (typeof artifactDefinitions)[number]['kind'];

export interface UIArtifact {
Expand Down
162 changes: 162 additions & 0 deletions apps/dbagent/src/components/chat/artifacts/code/client.tsx
Original file line number Diff line number Diff line change
@@ -0,0 +1,162 @@
import { toast } from '@internal/components';
import { CopyIcon, LogsIcon, PlayIcon, RedoIcon, UndoIcon } from 'lucide-react';
import { QueryEditor } from '../../query-editor';
import { generateUUID } from '../../utils';
import { Artifact } from '../create-artifact';

type ConsoleOutput = {
id: string;
contents: Array<ConsoleOutputContent>;
status: 'in_progress' | 'loading_packages' | 'completed' | 'failed';
};

type ConsoleOutputContent = {
type: 'text';
value: string;
};

interface Metadata {
outputs: Array<ConsoleOutput>;
}

export const codeArtifact = new Artifact<'code', Metadata>({
kind: 'code',
description: 'Useful for code generation; Code execution is only available for python code.',
initialize: async ({ setMetadata }) => {
setMetadata({
outputs: []
});
},
onStreamPart: ({ streamPart, setArtifact }) => {
if (streamPart.type === 'code-delta') {
setArtifact((draftArtifact) => ({
...draftArtifact,
content: streamPart.content as string,
isVisible:
draftArtifact.status === 'streaming' &&
draftArtifact.content.length > 300 &&
draftArtifact.content.length < 310
? true
: draftArtifact.isVisible,
status: 'streaming'
}));
}
},
content: ({ metadata, setMetadata, ...props }) => {
return (
<>
<div className="px-1">
<QueryEditor />
</div>

{metadata?.outputs && (
<div className="mt-2 flex flex-col gap-2">
{metadata.outputs.map((output) => (
<div key={output.id} className="rounded border p-2">
{output.status === 'in_progress' && <p>Running...</p>}
{output.status === 'loading_packages' && <p>Loading packages...</p>}
{output.status === 'completed' && (
<div>
{output.contents.map((content, index) => (
<div key={index}>
{content.type === 'text' ? <p>{content.value}</p> : <img src={content.value} alt="Output" />}
</div>
))}
</div>
)}
{output.status === 'failed' && <p>Error: {output.contents[0]?.value}</p>}
</div>
))}
</div>
)}
</>
);
},
actions: [
{
icon: <PlayIcon size={18} />,
label: 'Run',
description: 'Execute code',
onClick: async ({ content, setMetadata }) => {
const runId = generateUUID();
const outputContent: Array<ConsoleOutputContent> = [];

setMetadata((metadata) => ({
...metadata,
outputs: [
...metadata.outputs,
{
id: runId,
contents: [],
status: 'in_progress'
}
]
}));

try {
// Run the code in a sandboxed environment
} catch (error: any) {
setMetadata((metadata) => ({
...metadata,
outputs: [
...metadata.outputs.filter((output) => output.id !== runId),
{
id: runId,
contents: [{ type: 'text', value: error.message }],
status: 'failed'
}
]
}));
}
}
},
{
icon: <UndoIcon size={18} />,
description: 'View Previous version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('prev');
},
isDisabled: ({ currentVersionIndex }) => {
if (currentVersionIndex === 0) {
return true;
}

return false;
}
},
{
icon: <RedoIcon size={18} />,
description: 'View Next version',
onClick: ({ handleVersionChange }) => {
handleVersionChange('next');
},
isDisabled: ({ isCurrentVersion }) => {
if (isCurrentVersion) {
return true;
}

return false;
}
},
{
icon: <CopyIcon size={18} />,
description: 'Copy code to clipboard',
onClick: ({ content }) => {
navigator.clipboard.writeText(content);
toast.success('Copied to clipboard!');
}
}
],
toolbar: [
{
icon: <LogsIcon />,
description: 'Add logs',
onClick: ({ appendMessage }) => {
appendMessage({
role: 'user',
content: 'Add logs to the code snippet for debugging'
});
}
}
]
});
73 changes: 73 additions & 0 deletions apps/dbagent/src/components/chat/artifacts/code/server.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
import { streamObject } from 'ai';
import { z } from 'zod';
import { getModelInstance } from '~/lib/ai/agent';
import { updateDocumentPrompt } from '~/lib/ai/prompts';
import { createDocumentHandler } from '../server';

export const codeDocumentHandler = createDocumentHandler<'code'>({
kind: 'code',
onCreateDocument: async ({ title, dataStream }) => {
let draftContent = '';

const { fullStream } = streamObject({
model: await getModelInstance('chat'),
system: 'Generate code based on the given title. Output only the code.',
prompt: title,
schema: z.object({
code: z.string()
})
});

for await (const delta of fullStream) {
const { type } = delta;

if (type === 'object') {
const { object } = delta;
const { code } = object;

if (code) {
dataStream.writeData({
type: 'code-delta',
content: code ?? ''
});

draftContent = code;
}
}
}

return draftContent;
},
onUpdateDocument: async ({ document, description, dataStream }) => {
let draftContent = '';

const { fullStream } = streamObject({
model: await getModelInstance('chat'),
system: updateDocumentPrompt(document.content, 'code'),
prompt: description,
schema: z.object({
code: z.string()
})
});

for await (const delta of fullStream) {
const { type } = delta;

if (type === 'object') {
const { object } = delta;
const { code } = object;

if (code) {
dataStream.writeData({
type: 'code-delta',
content: code ?? ''
});

draftContent = code;
}
}
}

return draftContent;
}
});
9 changes: 7 additions & 2 deletions apps/dbagent/src/components/chat/artifacts/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ import { saveDocument } from '~/lib/db/chats';
import { DBAccess } from '~/lib/db/db';
import { ArtifactDocument } from '~/lib/db/schema';
import { ArtifactKind } from './artifact';
import { codeDocumentHandler } from './code/server';
import { sheetDocumentHandler } from './sheet/server';
import { textDocumentHandler } from './text/server';

Expand Down Expand Up @@ -93,6 +94,10 @@ export function createDocumentHandler<T extends ArtifactKind>(config: {
/*
* Use this array to define the document handlers for each artifact kind.
*/
export const documentHandlersByArtifactKind: Array<DocumentHandler> = [textDocumentHandler, sheetDocumentHandler];
export const documentHandlersByArtifactKind: Array<DocumentHandler> = [
textDocumentHandler,
sheetDocumentHandler,
codeDocumentHandler
];

export const artifactKinds = ['text', 'sheet'] as const;
export const artifactKinds = ['text', 'sheet', 'code'] as const;
30 changes: 30 additions & 0 deletions apps/dbagent/src/components/chat/query-editor.module.css
Original file line number Diff line number Diff line change
@@ -0,0 +1,30 @@
.Resizer {
height: 4px;
cursor: grab;
transition: 0.3s transform ease;
flex-shrink: 0;
width: 4px;
background-color: var(--muted);
}

.Resizer:hover {
background-color: var(--subtle);
}

@media (min-width: 768px) {
.Resizer {
background-image: linear-gradient(to bottom, var(--subtle));
width: 2px;
height: 100%;
}

.Resizer:hover,
.Resizer:active {
transform: scaleY(2);
}
}

.Resizer:active {
cursor: grabbing;
background-color: var(--muted);
}
Loading
Loading