Add output schemas to the playgrounds (#522)

Add output schema information to the playgrounds ![image](https://github.com/langchain-ai/langserve/assets/3205522/ce0e2ef4-3fa8-4ecd-8c21-08e1319879d5) --------- Co-authored-by: jacoblee93 <[email protected]>
langchain-ai · Mar 11, 2024 · e83dc5a · e83dc5a
1 parent 3f1fc38
commit e83dc5a
Show file tree

Hide file tree

Showing 11 changed files with 139 additions and 52 deletions.
diff --git a/langserve/api_handler.py b/langserve/api_handler.py
@@ -1361,6 +1361,7 @@ async def playground(
         return await serve_playground(
             self._runnable.with_config(config),
             self._runnable.with_config(config).input_schema,
+            self._runnable.with_config(config).output_schema,
             self._config_keys,
             playground_url,
             file_path,

diff --git a/..._playground/dist/assets/index-9018861a.js → ..._playground/dist/assets/index-c5f8c3dc.js b/..._playground/dist/assets/index-9018861a.js → ..._playground/dist/assets/index-c5f8c3dc.js
diff --git a/langserve/chat_playground/dist/index.html b/langserve/chat_playground/dist/index.html
@@ -5,18 +5,35 @@
     <link rel="icon" href="/____LANGSERVE_BASE_URL/favicon.ico" />
     <meta name="viewport" content="width=device-width, initial-scale=1.0" />
     <title>Chat Playground</title>
-    <script type="module" crossorigin src="/____LANGSERVE_BASE_URL/assets/index-9018861a.js"></script>
+    <script type="module" crossorigin src="/____LANGSERVE_BASE_URL/assets/index-c5f8c3dc.js"></script>
     <link rel="stylesheet" href="/____LANGSERVE_BASE_URL/assets/index-b47ed17e.css">
   </head>
   <body>
     <div id="root"></div>
     <script>
       try {
         window.CONFIG_SCHEMA = ____LANGSERVE_CONFIG_SCHEMA;
+      } catch (e) {
+        // pass
+      }
+      try {
         window.INPUT_SCHEMA = ____LANGSERVE_INPUT_SCHEMA;
+      } catch (e) {
+        // pass
+      }
+      try {
+        window.OUTPUT_SCHEMA = ____LANGSERVE_OUTPUT_SCHEMA;
+      } catch (e) {
+        // pass
+      }
+      try {
         window.FEEDBACK_ENABLED = ____LANGSERVE_FEEDBACK_ENABLED;
+      } catch (e) {
+        // pass
+      }
+      try {
         window.PUBLIC_TRACE_LINK_ENABLED = ____LANGSERVE_PUBLIC_TRACE_LINK_ENABLED;
-      } catch (error) {
+      } catch (e) {
         // pass
       }
     </script>

diff --git a/langserve/chat_playground/index.html b/langserve/chat_playground/index.html
@@ -11,10 +11,27 @@
     <script>
       try {
         window.CONFIG_SCHEMA = ____LANGSERVE_CONFIG_SCHEMA;
+      } catch (e) {
+        // pass
+      }
+      try {
         window.INPUT_SCHEMA = ____LANGSERVE_INPUT_SCHEMA;
+      } catch (e) {
+        // pass
+      }
+      try {
+        window.OUTPUT_SCHEMA = ____LANGSERVE_OUTPUT_SCHEMA;
+      } catch (e) {
+        // pass
+      }
+      try {
         window.FEEDBACK_ENABLED = ____LANGSERVE_FEEDBACK_ENABLED;
+      } catch (e) {
+        // pass
+      }
+      try {
         window.PUBLIC_TRACE_LINK_ENABLED = ____LANGSERVE_PUBLIC_TRACE_LINK_ENABLED;
-      } catch (error) {
+      } catch (e) {
         // pass
       }
     </script>

diff --git a/langserve/chat_playground/src/App.tsx b/langserve/chat_playground/src/App.tsx
@@ -2,32 +2,43 @@ import "./App.css";
 
 import { ChatWindow } from "./components/ChatWindow";
 import { AppCallbackContext, useAppStreamCallbacks } from "./useStreamCallback";
-import { useInputSchema } from "./useSchemas";
+import { useInputSchema, useOutputSchema } from "./useSchemas";
 import { useStreamLog } from "./useStreamLog";
-import { resolveApiUrl } from "./utils/url";
 
 export function App() {
   const { context, callbacks } = useAppStreamCallbacks();
   const { startStream, stopStream } = useStreamLog(callbacks);
   const inputSchema = useInputSchema({});
+  const outputSchema = useOutputSchema({});
   const inputProps = inputSchema?.data?.schema?.properties;
-  const isLoading = inputProps === undefined;
+  const outputDataSchema = outputSchema?.data?.schema;
+  const isLoading = inputProps === undefined || outputDataSchema === undefined;
   const inputKeys = Object.keys(inputProps ?? {});
-  const isSupported = isLoading || (inputKeys.length === 1 && inputProps[inputKeys[0]].type === "array");
+  const isSupported = isLoading || (
+    inputKeys.length === 1 &&
+    inputProps[inputKeys[0]].type === "array" &&
+    (
+      outputDataSchema.anyOf?.find((option) => option.properties?.type?.enum?.includes("ai")) ||
+      outputDataSchema.type === "string"
+    )
+  );
   return (
     <div className="flex items-center flex-col text-ls-black bg-background">
       <AppCallbackContext.Provider value={context}>
-        {isSupported 
-          ? <ChatWindow 
+        {isSupported
+          ? <ChatWindow
               startStream={startStream}
               stopStream={stopStream}
               inputKey={inputKeys[0]}
-            ></ChatWindow> 
+            ></ChatWindow>
           : <div className="h-[100vh] w-[100vw] flex justify-center items-center text-xl">
               <span className="text-center">
-                The chat playground is only supported for chains that take a single array of messages as input.
+                The chat playground is only supported for chains that take a single array of messages as input
                 <br/>
-                You can test this chain in the standard <a href={resolveApiUrl("/playground").toString()}>LangServe playground</a>.
+                and return either an AIMessage or a string.
+                <br />
+                <br />
+                You can test this chain in the default LangServe playground instead. Please set <code>playground_type="default"</code>.
               </span>
             </div>}
       </AppCallbackContext.Provider>

diff --git a/langserve/chat_playground/src/useSchemas.tsx b/langserve/chat_playground/src/useSchemas.tsx
@@ -1,7 +1,7 @@
-import { resolveApiUrl } from "./utils/url";
-import { simplifySchema } from "./utils/simplifySchema";
 import { JsonSchema } from "@jsonforms/core";
 import { compressToEncodedURIComponent } from "lz-string";
+import { resolveApiUrl } from "./utils/url";
+import { simplifySchema } from "./utils/simplifySchema";
 
 import useSWR from "swr";
 import defaults from "./utils/defaults";
@@ -13,6 +13,8 @@ declare global {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     INPUT_SCHEMA?: any;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    OUTPUT_SCHEMA?: any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     FEEDBACK_ENABLED?: any;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     PUBLIC_TRACE_LINK_ENABLED?: any;
@@ -96,3 +98,34 @@ export function useInputSchema(configData?: unknown) {
     { keepPreviousData: true }
   );
 }
+
+export function useOutputSchema(configData?: unknown) {
+  return useSWR(
+    ["/output_schema", configData],
+    async ([, configData]) => {
+      // TODO: this won't work if we're already seeing a prefixed URL
+      const prefix = configData
+        ? `/c/${compressToEncodedURIComponent(JSON.stringify(configData))}`
+        : "";
+
+      let schema: JsonSchema | null = null;
+
+      if (!prefix && !import.meta.env.DEV && window.OUTPUT_SCHEMA) {
+        schema = await simplifySchema(window.OUTPUT_SCHEMA);
+      } else {
+        const response = await fetch(resolveApiUrl(`${prefix}/output_schema`));
+        if (!response.ok) throw new Error(await response.text());
+
+        const json = await response.json();
+        schema = await simplifySchema(json);
+      }
+
+      if (schema == null) return null;
+      return {
+        schema,
+        defaults: defaults(schema),
+      };
+    },
+    { keepPreviousData: true }
+  );
+}
diff --git a/langserve/chat_playground/vite.config.ts b/langserve/chat_playground/vite.config.ts
@@ -8,7 +8,7 @@ export default defineConfig({
   plugins: [svgr(), react()],
   server: {
     proxy: {
-      "^/____LANGSERVE_BASE_URL.*/(config_schema|input_schema|stream_log|feedback|public_trace_link)(/[a-zA-Z0-9-]*)?$": {
+      "^/____LANGSERVE_BASE_URL.*/(config_schema|input_schema|output_schema|stream_log|feedback|public_trace_link)(/[a-zA-Z0-9-]*)?$": {
         target: "http://127.0.0.1:8000",
         changeOrigin: true,
         rewrite: (path) => path.replace("/____LANGSERVE_BASE_URL", ""),

diff --git a/langserve/playground.py b/langserve/playground.py
@@ -50,6 +50,7 @@ def _get_mimetype(path: str) -> str:
 async def serve_playground(
     runnable: Runnable,
     input_schema: Type[BaseModel],
+    output_schema: Type[BaseModel],
     config_keys: Sequence[str],
     base_url: str,
     file_path: str,
@@ -92,6 +93,7 @@ async def serve_playground(
                         runnable.config_schema(include=config_keys).schema()
                     ),
                     LANGSERVE_INPUT_SCHEMA=json.dumps(input_schema.schema()),
+                    LANGSERVE_OUTPUT_SCHEMA=json.dumps(output_schema.schema()),
                     LANGSERVE_FEEDBACK_ENABLED=json.dumps(
                         "true" if feedback_enabled else "false"
                     ),

diff --git a/langserve/playground/dist/index.html b/langserve/playground/dist/index.html
@@ -14,6 +14,7 @@
       try {
         window.CONFIG_SCHEMA = ____LANGSERVE_CONFIG_SCHEMA;
         window.INPUT_SCHEMA = ____LANGSERVE_INPUT_SCHEMA;
+        window.OUTPUT_SCHEMA = ____LANGSERVE_OUTPUT_SCHEMA;
         window.FEEDBACK_ENABLED = ____LANGSERVE_FEEDBACK_ENABLED;
         window.PUBLIC_TRACE_LINK_ENABLED = ____LANGSERVE_PUBLIC_TRACE_LINK_ENABLED;
       } catch (error) {

diff --git a/langserve/playground/index.html b/langserve/playground/index.html
@@ -12,6 +12,7 @@
       try {
         window.CONFIG_SCHEMA = ____LANGSERVE_CONFIG_SCHEMA;
         window.INPUT_SCHEMA = ____LANGSERVE_INPUT_SCHEMA;
+        window.OUTPUT_SCHEMA = ____LANGSERVE_OUTPUT_SCHEMA;
         window.FEEDBACK_ENABLED = ____LANGSERVE_FEEDBACK_ENABLED;
         window.PUBLIC_TRACE_LINK_ENABLED = ____LANGSERVE_PUBLIC_TRACE_LINK_ENABLED;
       } catch (error) {

diff --git a/langserve/playground/src/useSchemas.tsx b/langserve/playground/src/useSchemas.tsx
@@ -13,7 +13,11 @@ declare global {
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
     INPUT_SCHEMA?: any;
     // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    OUTPUT_SCHEMA?: any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
     FEEDBACK_ENABLED?: any;
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    PUBLIC_TRACE_LINK_ENABLED?: any;
   }
 }