Only use forevervm

Author: Jake Lazaroff

src/components/altair/Altair.tsx

@@ -25,6 +25,16 @@ const fvm = new ForeverVM({ token: await generateReplToken() });
 
 const repl = fvm.repl();
 
+const env = repl.exec(
+  `import os
+
+os.environ['AWS_ACCESS_KEY_ID'] = '${process.env.REACT_APP_AWS_ACCESS_KEY_ID}'
+os.environ['AWS_SECRET_ACCESS_KEY'] = '${process.env.REACT_APP_AWS_SECRET_ACCESS_KEY}'
+os.environ['AWS_DEFAULT_REGION'] = 'us-east-1'
+`.trim(),
+);
+await env.result;
+
 const renderAltair: FunctionDeclaration = {
   name: "render_altair",
   description: "Displays an altair graph in json format.",
@@ -41,10 +51,10 @@ const renderAltair: FunctionDeclaration = {
   },
 };
 
-const runPython: FunctionDeclaration = {
-  name: "run_python",
+const forevervm: FunctionDeclaration = {
+  name: "forevervm",
   description:
-    "Run Python code in a stateful read-eval-print loop. Variables, imports, and functions are persisted between calls. Imports of common packages, including requests, matplotlib, and pandas, are permitted, but you may not install libraries.",
+    "Run Python code in a stateful read-eval-print loop. Variables, imports, and functions are persisted between calls. Imports of common packages, including requests, matplotlib, and pandas and boto3 are permitted, but you may not install libraries.",
   parameters: {
     type: SchemaType.OBJECT,
     properties: {
@@ -58,7 +68,8 @@ const runPython: FunctionDeclaration = {
 };
 
 function AltairComponent() {
-  const [jsonString, setJSONString] = useState<string>("");
+  const [vegaSource, setVegaSource] = useState<string>("");
+  const [imgSource, setImgSource] = useState<string>("");
   const { client, setConfig } = useLiveAPIContext();
 
   useEffect(() => {
@@ -73,14 +84,14 @@ function AltairComponent() {
       systemInstruction: {
         parts: [
           {
-            text: 'You are my helpful assistant. Any time I ask you for a graph call the "render_altair" function I have provided you. Dont ask for additional information just make your best judgement.',
+            text: 'You are my helpful assistant. Any time I ask you to execute code, use the "forevervm" function I have provided you. You can display charts and graphs using matplotlib with plt.show(). Dont ask for additional information just make your best judgement.',
           },
         ],
       },
       tools: [
         // there is a free-tier quota for search
         { googleSearch: {} },
-        { functionDeclarations: [renderAltair, runPython] },
+        { functionDeclarations: [forevervm] },
       ],
     });
   }, [setConfig]);
@@ -95,20 +106,20 @@ function AltairComponent() {
         switch (fc?.name) {
           case renderAltair.name: {
             const str = (fc.args as any).json_graph;
-            setJSONString(str);
+            setVegaSource(str);
             functionResponses.push({ id: fc.id, response: { output: { success: true } } });
             break;
           }
 
-          case runPython.name: {
+          case forevervm.name: {
             const code = (fc.args as any).code;
 
             const instruction = repl.exec(code);
 
             let output: StandardOutput[] = [];
             for await (const line of instruction.output) output.push(line);
-
             const result = await repl.exec(code).result;
+            if (result.data?.png) setImgSource(result.data.png as string);
 
             functionResponses.push({ id: fc.id, response: { output: { output, result } } });
             break;
@@ -130,10 +141,18 @@ function AltairComponent() {
   const embedRef = useRef<HTMLDivElement>(null);
 
   useEffect(() => {
-    if (embedRef.current && jsonString) {
-      vegaEmbed(embedRef.current, JSON.parse(jsonString));
-    }
-  }, [embedRef, jsonString]);
+    if (!embedRef.current || !vegaSource) return;
+    vegaEmbed(embedRef.current, JSON.parse(vegaSource));
+  }, [embedRef, vegaSource]);
+
+  useEffect(() => {
+    if (!embedRef.current || !imgSource) return;
+
+    const img = document.createElement("img");
+    img.src = "data:image/png;base64," + imgSource;
+    embedRef.current.replaceChildren(img);
+  }, [embedRef, imgSource]);
+
   return <div className="vega-embed" ref={embedRef} />;
 }