Merge pull request #12 from docker/cm/jsonrpc

Add JSON RPC stream

Author: ColinMcNeil

src/extension/ui/src/App.tsx

@@ -2,12 +2,18 @@ import React, { useEffect } from 'react';
 import Button from '@mui/material/Button';
 import DelIcon from '@mui/icons-material/Delete';
 import { createDockerDesktopClient } from '@docker/extension-api-client';
-import { IconButton, Link, List, ListItem, ListItemButton, ListItemText, Paper, Stack, TextField, Typography } from '@mui/material';
+import { Chip, IconButton, Link, List, ListItem, ListItemButton, ListItemText, Paper, Stack, TextField, Typography } from '@mui/material';
 import { getRunArgs } from './args';
 import Convert from 'ansi-to-html';
 
 const convert = new Convert({ newline: true });
 
+type RPCMessage = {
+  jsonrpc?: string;
+  method: string;
+  params: any;
+}
+
 // Note: This line relies on Docker Desktop's presence as a host application.
 // If you're running this React app in a browser, it won't work properly.
 const client = createDockerDesktopClient();
@@ -33,10 +39,11 @@ export function App() {
 
   const [promptInput, setPromptInput] = React.useState<string>('');
 
-  const [runOut, setRunOut] = React.useState<string>('');
+  const [runOut, setRunOut] = React.useState<RPCMessage[]>([]);
 
   const scrollRef = React.useRef<HTMLDivElement>(null);
 
+  const [showDebug, setShowDebug] = React.useState(false);
 
   useEffect(() => {
     localStorage.setItem('projects', JSON.stringify(projects));
@@ -90,12 +97,38 @@ export function App() {
   const delim = client.host.platform === 'win32' ? '\\' : '/';
 
   const startPrompt = async () => {
-    let output = ""
-    const updateOutput = (data: string) => {
-      output += data;
+    let output: RPCMessage[] = []
+    const updateOutput = (line: RPCMessage) => {
+      if (line.method === 'functions') {
+        const functions = line.params;
+        for (const func of functions) {
+          const functionId = func.id;
+          const existingFunction = output.find(o =>
+            o.method === 'functions'
+            &&
+            o.params.find((p: { id: string }) => p.id === functionId)
+          );
+          if (existingFunction) {
+            const existingFunctionParamsIndex = existingFunction.params.findIndex((p: { id: string }) => p.id === functionId);
+            existingFunction.params[existingFunctionParamsIndex] = { ...existingFunction.params[existingFunctionParamsIndex], ...func };
+            output = output.map(
+              o => o.method === 'functions'
+                ?
+                { ...o, params: o.params.map((p: { id: string }) => p.id === functionId ? { ...p, ...func } : p) }
+                :
+                o
+            );
+          } else {
+            output = [...output, line];
+          }
+        }
+      }
+      else {
+        output = [...output, line];
+      }
       setRunOut(output);
     }
-    updateOutput("Pulling images\n")
+    updateOutput({ method: 'message', params: { debug: 'Pulling images' } })
     try {
       const pullWriteFiles = await client.docker.cli.exec("pull", ["vonwig/function_write_files"]);
       const pullPrompts = await client.docker.cli.exec("pull", ["vonwig/prompts"]);
@@ -106,12 +139,12 @@ export function App() {
         "vonwig/function_write_files",
         `'` + JSON.stringify({ files: [{ path: ".openai-api-key", content: openAIKey, executable: false }] }) + `'`
       ]);
-      updateOutput(JSON.stringify({ pullWriteFiles, pullPrompts, writeKey }));
+      updateOutput({ method: 'message', params: { debug: JSON.stringify({ pullWriteFiles, pullPrompts, writeKey }) } });
     }
     catch (e) {
-      updateOutput(JSON.stringify(e));
+      updateOutput({ method: 'message', params: { debug: JSON.stringify(e) } });
     }
-    updateOutput("Running prompts\n")
+    updateOutput({ method: 'message', params: { debug: 'Running prompts...' } })
     const args = getRunArgs(selectedPrompt!, selectedProject!, "", client.host.platform)
 
     client.docker.cli.exec("run", args, {
@@ -120,24 +153,33 @@ export function App() {
         onOutput: ({ stdout, stderr }) => {
           if (stdout && stdout.startsWith('{')) {
             let rpcMessage = stdout.split('}Content-Length:')[0]
-            if (!rpcMessage.endsWith('}')) {
+            if (!rpcMessage.endsWith('}}')) {
               rpcMessage += '}'
             }
             const json = JSON.parse(rpcMessage)
-            if (json.params.content) {
-              output += json.params.content
-            }
+            updateOutput(json)
+            //   {
+            //     "jsonrpc": "2.0",
+            //     "method": "functions",
+            //     "params": [
+            //         {
+            //             "function": {
+            //                 "name": "run-eslint",
+            //                 "arguments": "{\n  \""
+            //             },
+            //             "id": "call_53E2o4fq1QEmIHixWcKZmOqo"
+            //         }
+            //     ]
+            // }
           }
           if (stderr) {
-            output += stderr
+            updateOutput({ method: 'message', params: { debug: stderr } });
           }
-          setRunOut(output);
         },
         onError: (err) => {
           console.error(err);
-          output += err;
-          setRunOut(output);
-        }
+          updateOutput({ method: 'message', params: { debug: err } });
+        },
       }
     });
   }
@@ -196,18 +238,31 @@ export function App() {
         {/* Prompts column */}
         <Paper sx={{ padding: 1 }}>
           <Typography variant="h3">Prompts</Typography>
-          <TextField
-            sx={{ width: '100%', mt: 1 }}
-            placeholder='Enter GitHub ref or URL'
-            value={promptInput}
-            onChange={(e) => setPromptInput(e.target.value)}
-          />
-          {promptInput.length > 0 && (
+          <Stack direction='row' spacing={1} alignItems={'center'} justifyContent={'space-between'}>
+            <TextField
+              fullWidth
+              placeholder='Enter GitHub ref or URL'
+              value={promptInput}
+              onChange={(e) => setPromptInput(e.target.value)}
+            />
+            {promptInput.length > 0 && (
+              <Button onClick={() => {
+                setPrompts([...prompts, promptInput]);
+                setPromptInput('');
+              }}>Add prompt</Button>
+            )}
             <Button onClick={() => {
-              setPrompts([...prompts, promptInput]);
-              setPromptInput('');
-            }}>Add prompt</Button>
-          )}
+              client.desktopUI.dialog.showOpenDialog({
+                properties: ['openDirectory', 'multiSelections']
+              }).then((result) => {
+                if (result.canceled) {
+                  return;
+                }
+                setPrompts([...prompts, ...result.filePaths.map(p => `local://${p}`)]);
+              });
+            }}>Add local prompt</Button>
+          </Stack>
+
           <List>
             {prompts.map((prompt) => (
               <ListItem
@@ -232,8 +287,12 @@ export function App() {
                 }>
                 <ListItemButton sx={{ padding: 0, pl: 1.5 }} onClick={() => {
                   setSelectedPrompt(prompt);
-                }}>
-                  <ListItemText primary={prompt.split(delim).pop()} secondary={prompt} />
+                }}>{
+                    prompt.startsWith('local://') ?
+                      <><ListItemText primary={<>{prompt.split(delim).pop()}<Chip sx={{ ml: 1 }} label='local' /></>} secondary={prompt.replace('local://', '')} /></>
+                      :
+                      <ListItemText primary={prompt.split('/').pop()} secondary={prompt} />
+                  }
                 </ListItemButton>
               </ListItem>
             ))}
@@ -259,10 +318,30 @@ export function App() {
         )}
         {/* Show run output */}
         {
-          runOut && (
+          runOut.length > 0 && (
             <Paper sx={{ p: 1 }}>
-              <Typography variant='h3'>Run output</Typography>
-              <div style={{ whiteSpace: 'pre-wrap' }} dangerouslySetInnerHTML={{ __html: convert.toHtml(runOut) }} />
+              <Stack direction='row' spacing={1} alignItems={'center'} justifyContent={'space-between'}>
+                <Typography variant='h3'>Run output</Typography>
+                <Button onClick={() => setShowDebug(!showDebug)}>{showDebug ? 'Hide' : 'Show'} debug</Button>
+              </Stack>
+
+              <div style={{ overflow: 'auto', maxHeight: '100vh' }}>
+                {runOut.map((line, i) => {
+                  if (line.method === 'message') {
+                    if (line.params.debug) {
+                      return showDebug ? <Typography key={i} variant='body1' sx={theme => ({ color: theme.palette.docker.grey[400] })}>{line.params.debug}</Typography> : null;
+                    }
+                    if (line.params.role === 'assistant') {
+                      return <Typography key={i} variant='body1' sx={theme => ({ color: theme.palette.docker.blue[400] })}>{line.params.content}</Typography>
+                    }
+                    return <pre key={i} style={{ whiteSpace: 'pre-wrap', display: 'inline' }} dangerouslySetInnerHTML={{ __html: convert.toHtml(line.params.content) }} />
+                  }
+                  if (line.method === 'functions') {
+                    return <Typography key={i} variant='body1' sx={{ whiteSpace: 'pre-wrap' }}>{JSON.stringify(line.params, null, 2)}</Typography>
+                  }
+                  return <Typography key={i} variant='body1'>{JSON.stringify(line)}</Typography>
+                })}
+              </div>
             </Paper>
           )
         }

src/extension/ui/src/args.ts

@@ -1,18 +1,32 @@
-export const getRunArgs = (prompt_ref: string, project_dir: string, username: string, platform: string) => {
-    return [
+export const getRunArgs = (promptRef: string, projectDir: string, username: string, platform: string) => {
+    const isLocal = promptRef.startsWith('local://');
+    let promptArgs: string[] = ["--prompts", promptRef];
+    let mountArgs: string[] = [];
+
+    if (isLocal) {
+        const localPromptPath = promptRef.replace('local://', '');
+        const pathSeparator = platform === 'win32' ? '\\' : '/';
+        promptRef = localPromptPath.split(pathSeparator).pop() || 'unknown-local-prompt';
+        promptArgs = ["--prompts-dir", `/app/${promptRef}`];
+        mountArgs = ["--mount", `type=bind,source=${localPromptPath},target=/app/${promptRef}`];
+    }
+
+    const baseArgs: string[] = [
         '--rm',
-        '-v',
-        '/var/run/docker.sock:/var/run/docker.sock',
-        '-v',
-        'openai_key:/root',
-        '--mount',
-        'type=volume,source=docker-prompts,target=/prompts',
+        '-v', '/var/run/docker.sock:/var/run/docker.sock',
+        '-v', 'openai_key:/root',
+        '--mount', 'type=volume,source=docker-prompts,target=/prompts'
+    ];
+
+    const runArgs: string[] = [
         'vonwig/prompts:latest',
         'run',
-        "--host-dir", project_dir,
+        "--host-dir", projectDir,
         "--user", username,
         "--platform", platform,
-        "--prompts", prompt_ref,
+        ...promptArgs,
         '--jsonrpc'
     ];
+
+    return [...baseArgs, ...mountArgs, ...runArgs];
 }