Fixes for CURL (#735)

* Remove vscode requirement and use VSCodeAPIWrapper to support prompt fiddle * render_curl uses async method * show errors when prompt fails to render while in curl mode
BoundaryML · Jun 29, 2024 · e16b581 · e16b581
1 parent 423faa1
commit e16b581
Show file tree

Hide file tree

Showing 5 changed files with 59 additions and 55 deletions.
diff --git a/engine/baml-runtime/src/runtime_interface.rs b/engine/baml-runtime/src/runtime_interface.rs
@@ -136,6 +136,7 @@ pub trait InternalRuntimeInterface {
         node_index: Option<usize>,
     ) -> Result<(RenderedPrompt, OrchestrationScope)>;
 
+    #[warn(async_fn_in_trait)]
     async fn render_raw_curl(
         &self,
         function_name: &str,

diff --git a/engine/baml-schema-wasm/src/runtime_wasm/mod.rs b/engine/baml-schema-wasm/src/runtime_wasm/mod.rs
@@ -994,14 +994,14 @@ impl WasmFunction {
                 RenderedPrompt::Chat(chat_messages) => chat_messages,
                 RenderedPrompt::Completion(_) => vec![], // or handle this case differently
             },
-            Err(e) => return Err(wasm_bindgen::JsError::new(format!("{:?}", e).as_str())),
+            Err(e) => return Err(wasm_bindgen::JsError::new(format!("{:#?}", e).as_str())),
         };
 
         rt.runtime
             .internal()
             .render_raw_curl(&self.name, &ctx, &final_prompt, stream, None)
             .await
-            .map_err(|e| wasm_bindgen::JsError::new(format!("{e:?}").as_str()))
+            .map_err(|e| wasm_bindgen::JsError::new(format!("{e:#?}").as_str()))
     }
 
     #[wasm_bindgen]

diff --git a/typescript/playground-common/src/baml_wasm_web/EventListener.tsx b/typescript/playground-common/src/baml_wasm_web/EventListener.tsx
@@ -10,7 +10,7 @@ import CustomErrorBoundary from '../utils/ErrorFallback'
 import { sessionStore, vscodeLocalStorageStore } from './JotaiProvider'
 import { availableProjectsAtom, projectFamilyAtom, projectFilesAtom, runtimeFamilyAtom } from './baseAtoms'
 import type { WasmDiagnosticError, WasmParam, WasmRuntime } from '@gloo-ai/baml-schema-wasm-web/baml_schema_build'
-const vscode = acquireVsCodeApi()
+import { vscode } from '../utils/vscode'
 
 // const wasm = await import("@gloo-ai/baml-schema-wasm-web/baml_schema_build");
 // const { WasmProject, WasmRuntime, WasmRuntimeContext, version: RuntimeVersion } = wasm;
@@ -347,30 +347,29 @@ export const availableFunctionsAtom = atom((get) => {
   return runtime.list_functions()
 })
 
-export const rawCurlAtom = atom((get) => {
+const asyncCurlAtom = atom(async (get) => {
   const runtime = get(selectedRuntimeAtom)
   const func = get(selectedFunctionAtom)
   const test_case = get(selectedTestCaseAtom)
 
   if (!runtime || !func || !test_case) {
-    return null
+    return 'Not yet ready'
   }
   const params = Object.fromEntries(
     test_case.inputs
       .filter((i): i is WasmParam & { value: string } => i.value !== undefined)
       .map((input) => [input.name, JSON.parse(input.value)]),
   )
   try {
-    return func.render_raw_curl(runtime, params, false)
+    return await func.render_raw_curl(runtime, params, false)
   } catch (e) {
-    if (e instanceof Error) {
-      return e.message
-    } else {
-      return `${e}`
-    }
+    console.error(e)
+    return 'Error rendering curl command'
   }
 })
 
+export const curlAtom = unwrap(asyncCurlAtom)
+
 export const renderPromptAtom = atom((get) => {
   const runtime = get(selectedRuntimeAtom)
   const func = get(selectedFunctionAtom)
@@ -621,7 +620,7 @@ export const EventListener: React.FC<{ children: React.ReactNode }> = ({ childre
 
   return (
     <>
-      <div className='absolute flex flex-row gap-2 text-xs bg-transparent right-2 bottom-2'>
+      <div className='absolute flex flex-row gap-2 text-xs bg-transparent right-2 bottom-2 z-50'>
         <ErrorCount /> <span>Runtime Version: {version}</span>
       </div>
       {selectedProject === null ? (

diff --git a/typescript/playground-common/src/baml_wasm_web/JotaiProvider.tsx b/typescript/playground-common/src/baml_wasm_web/JotaiProvider.tsx
@@ -5,20 +5,15 @@ import { createJSONStorage } from 'jotai/utils'
 import type { SyncStorage } from 'jotai/vanilla/utils/atomWithStorage'
 import { DevTools } from 'jotai-devtools'
 import 'jotai-devtools/styles.css'
+import { vscode } from '../utils/vscode'
 
 export const atomStore = createStore()
 
-export const vscodeAPI = () => {
-  if (typeof acquireVsCodeApi === 'function') {
-    return acquireVsCodeApi()
-  }
-  return undefined
-}
 function setVSCodeState(state: any) {
-  vscodeAPI()?.setState(state)
+  vscode.setState(state)
 }
 function getLocalStorage() {
-  const state = vscodeAPI()?.getState() || { localStorage: {} }
+  const state = vscode.getState() || { localStorage: {} }
   return (state as any).localStorage
 }
 

diff --git a/typescript/playground-common/src/shared/FunctionPanel.tsx b/typescript/playground-common/src/shared/FunctionPanel.tsx
@@ -1,7 +1,7 @@
 /// Content once a function has been selected.
 import { useAppState } from './AppStateContext'
 import { useAtomValue, useSetAtom } from 'jotai'
-import { renderPromptAtom, selectedFunctionAtom, rawCurlAtom } from '../baml_wasm_web/EventListener'
+import { renderPromptAtom, selectedFunctionAtom, curlAtom } from '../baml_wasm_web/EventListener'
 import TestResults from '../baml_wasm_web/test_uis/test_result'
 import { ResizableHandle, ResizablePanel, ResizablePanelGroup } from '../components/ui/resizable'
 import { TooltipProvider } from '../components/ui/tooltip'
@@ -10,13 +10,44 @@ import FunctionTestSnippet from './TestSnippet'
 import { Copy } from 'lucide-react'
 import { Button } from '../components/ui/button'
 import { CheckboxHeader } from './CheckboxHeader'
+import CustomErrorBoundary from '../utils/ErrorFallback'
 const handleCopy = (text: string) => () => {
   navigator.clipboard.writeText(text)
 }
+
+const CurlSnippet: React.FC = () => {
+  const rawCurl = useAtomValue(curlAtom) ?? 'Loading...'
+
+  return (
+    <div>
+      <div className='flex justify-end'>
+        <Button
+          onClick={handleCopy(rawCurl)}
+          className='copy-button bg-transparent text-white m-0 py-0 hover:bg-indigo-500 text-xs'
+        >
+          <Copy size={16} />
+        </Button>
+      </div>
+      <PromptChunk
+        text={rawCurl}
+        client={{
+          identifier: {
+            end: 0,
+            source_file: '',
+            start: 0,
+            value: 'Curl Request',
+          },
+          provider: '',
+          model: '',
+        }}
+        showCopy={true}
+      />
+    </div>
+  )
+}
+
 const PromptPreview: React.FC = () => {
   const promptPreview = useAtomValue(renderPromptAtom)
-  const rawCurl = useAtomValue(rawCurlAtom) ?? 'Not yet available'
-
   const { showCurlRequest } = useAppState()
 
   if (!promptPreview) {
@@ -28,36 +59,7 @@ const PromptPreview: React.FC = () => {
     )
   }
 
-  if (showCurlRequest) {
-    return (
-      <div>
-        <div className='flex justify-end'>
-          <Button
-            onClick={handleCopy(rawCurl)}
-            className='copy-button bg-transparent text-white m-0 py-0 hover:bg-indigo-500 text-xs'
-          >
-            <Copy size={16} />
-          </Button>
-        </div>
-        <PromptChunk
-          text={rawCurl}
-          client={{
-            identifier: {
-              end: 0,
-              source_file: '',
-              start: 0,
-              value: 'Curl Request',
-            },
-            provider: '',
-            model: '',
-          }}
-          showCopy={true}
-        />
-      </div>
-    )
-  }
-
-  if (typeof promptPreview === 'string')
+  if (typeof promptPreview === 'string') {
     return (
       <PromptChunk
         text={promptPreview}
@@ -74,6 +76,11 @@ const PromptPreview: React.FC = () => {
         }}
       />
     )
+  }
+
+  if (showCurlRequest) {
+    return <CurlSnippet />
+  }
 
   return (
     <div className='flex flex-col w-full h-full gap-4 px-2'>
@@ -164,9 +171,11 @@ enum Topic {
           <ResizablePanel id='top-panel' className='flex w-full px-1' defaultSize={50}>
             <div className='w-full'>
               <ResizablePanelGroup direction='horizontal' className='h-full'>
-                <div className='relative w-full h-full overflow-y-auto'>
+                <div className='w-full h-full'>
                   <CheckboxHeader />
-                  <PromptPreview />
+                  <div className='relative w-full overflow-y-auto' style={{ height: 'calc(100% - 32px)' }}>
+                    <PromptPreview />
+                  </div>
                 </div>
               </ResizablePanelGroup>