Code editor artifact

apps/dbagent/package.json

@@ -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",
@@ -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",
@@ -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",
@@ -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",
@@ -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",

apps/dbagent/src/components/chat/artifacts/artifact.tsx

@@ -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 {

apps/dbagent/src/components/chat/artifacts/code/client.tsx

@@ -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'
+        });
+      }
+    }
+  ]
+});

apps/dbagent/src/components/chat/artifacts/code/server.ts

@@ -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;
+  }
+});

apps/dbagent/src/components/chat/artifacts/server.ts

@@ -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';
 
@@ -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;

apps/dbagent/src/components/chat/query-editor.module.css

@@ -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);
+}