From 66d9ec9cadfd41c3bfd3e8cae736ec290a2bdcb2 Mon Sep 17 00:00:00 2001
From: John Smith <john@johnjcsmith.com>
Date: Thu, 19 Dec 2024 17:41:34 +1030
Subject: [PATCH] feat: Allow passing run context in UI

---
 app/components/chat/prompt-textarea.tsx |  28 ++++--
 app/components/chat/run-config.tsx      | 108 ++++++++++++++++++------
 app/components/chat/sdk-commands.tsx    |  10 +--
 3 files changed, 107 insertions(+), 39 deletions(-)

diff --git a/app/components/chat/prompt-textarea.tsx b/app/components/chat/prompt-textarea.tsx
index fea1f632..0e0308c5 100644
--- a/app/components/chat/prompt-textarea.tsx
+++ b/app/components/chat/prompt-textarea.tsx
@@ -16,13 +16,14 @@ import { TemplateSearch } from "./template-search";
 
 export type RunConfiguration = {
   attachedFunctions: string[];
-  structuredOutput: string | null;
+  resultSchema: string | null;
   reasoningTraces: boolean;
   prompt: string;
   template?: {
     id: string;
     input: Record<string, string>;
   };
+  runContext: string | null;
 };
 
 export function PromptTextarea({ clusterId }: { clusterId: string }) {
@@ -48,14 +49,16 @@ export function PromptTextarea({ clusterId }: { clusterId: string }) {
 
   const [runConfig, setRunConfig] = useState({
     attachedFunctions: [] as string[],
-    structuredOutput: null as string | null,
+    resultSchema: null as string | null,
     reasoningTraces: true,
+    runContext: null as string | null,
   });
 
   const handleConfigChange = (newConfig: {
     attachedFunctions: string[];
-    structuredOutput: string | null;
+    resultSchema: string | null;
     reasoningTraces: boolean;
+    runContext: string | null;
   }) => {
     setRunConfig(newConfig);
   };
@@ -80,8 +83,12 @@ export function PromptTextarea({ clusterId }: { clusterId: string }) {
         body.initialPrompt = config.prompt.replace(/{{.*?}}/g, "");
       }
 
-      if (config.structuredOutput) {
-        body.resultSchema = JSON.parse(config.structuredOutput);
+      if(config.runContext) {
+        body.context = JSON.parse(config.runContext);
+      }
+
+      if (config.resultSchema) {
+        body.resultSchema = JSON.parse(config.resultSchema);
       }
 
       if (config.attachedFunctions && config.attachedFunctions.length > 0) {
@@ -119,8 +126,9 @@ export function PromptTextarea({ clusterId }: { clusterId: string }) {
         setStoredInputs({});
         setRunConfig({
           attachedFunctions: [],
-          structuredOutput: null,
+          resultSchema: null,
           reasoningTraces: true,
+          runContext: null
         });
       }
 
@@ -149,9 +157,10 @@ export function PromptTextarea({ clusterId }: { clusterId: string }) {
 
     onSubmit({
       attachedFunctions: runConfig.attachedFunctions,
-      structuredOutput: runConfig.structuredOutput,
+      resultSchema: runConfig.resultSchema,
       reasoningTraces: runConfig.reasoningTraces,
       prompt: updatedPrompt,
+      runContext: runConfig.runContext,
       template: selectedTemplate
         ? {
             id: selectedTemplate.id,
@@ -177,9 +186,10 @@ export function PromptTextarea({ clusterId }: { clusterId: string }) {
   }) => {
     setRunConfig({
       attachedFunctions: template.attachedFunctions,
-      structuredOutput: template.structuredOutput
+      resultSchema: template.structuredOutput
         ? JSON.stringify(template.structuredOutput)
         : null,
+      runContext: null,
       reasoningTraces: true,
     });
     template.initialPrompt && setPrompt(template.initialPrompt);
@@ -249,7 +259,7 @@ export function PromptTextarea({ clusterId }: { clusterId: string }) {
             clusterId={clusterId}
             config={{
               attachedFunctions: runConfig.attachedFunctions,
-              structuredOutput: runConfig.structuredOutput,
+              resultSchema: runConfig.resultSchema,
               prompt,
             }}
           />
diff --git a/app/components/chat/run-config.tsx b/app/components/chat/run-config.tsx
index c4e96e0c..9e925c4b 100644
--- a/app/components/chat/run-config.tsx
+++ b/app/components/chat/run-config.tsx
@@ -11,13 +11,15 @@ import { Checkbox } from "../ui/checkbox";
 type RunConfigProps = {
   config: {
     attachedFunctions: string[];
-    structuredOutput: string | null;
+    resultSchema: string | null;
     reasoningTraces: boolean;
+    runContext: string | null;
   };
   onConfigChange: (newConfig: {
     attachedFunctions: string[];
-    structuredOutput: string | null;
+    resultSchema: string | null;
     reasoningTraces: boolean;
+    runContext: string | null;
   }) => void;
 };
 
@@ -27,36 +29,53 @@ const RunConfig: React.FC<RunConfigProps> = ({ config, onConfigChange }) => {
     config.attachedFunctions.filter(Boolean),
   );
 
-  const [structuredOutput, setStructuredOutput] = useState<string | null>(
-    config.structuredOutput,
+
+  const [runContext, setRunContext] = useState<string | null>(
+    config.runContext ? JSON.stringify(config.runContext) : null,
+  );
+
+  const [resultSchema, setResultSchema] = useState<string | null>(
+    config.resultSchema,
   );
 
   const [reasoningTraces, setReasoningTraces] = useState(
     config.reasoningTraces,
   );
 
-  const hasConfig = attachedFunctions.length > 0 || !!structuredOutput;
+  const hasConfig = attachedFunctions.length > 0 || !!resultSchema;
 
   const handleApplyConfig = () => {
-    if (structuredOutput) {
+    if (resultSchema) {
       try {
-        JSON.parse(structuredOutput);
-        setJsonInvalid(false);
+        JSON.parse(resultSchema);
+        setResultSchemaJsonInvalid(false);
       } catch (e) {
-        toast.error("Invalid JSON in structured output");
+        toast.error("Invalid JSON in result schema");
+        return;
+      }
+    }
+
+    if (runContext) {
+      try {
+        JSON.parse(runContext);
+        setRunContextJsonInvalid(false);
+      } catch (e) {
+        toast.error("Invalid JSON in run context");
         return;
       }
     }
 
     onConfigChange({
       attachedFunctions,
-      structuredOutput,
+      resultSchema,
       reasoningTraces,
+      runContext,
     });
     setIsOpen(false);
   };
 
-  const [jsonInvalid, setJsonInvalid] = useState(false);
+  const [resultSchemaJsonInvalid, setResultSchemaJsonInvalid] = useState(false);
+  const [runContextJsonInvalid, setRunContextJsonInvalid] = useState(false);
 
   return (
     <Popover open={isOpen} onOpenChange={setIsOpen}>
@@ -72,6 +91,11 @@ const RunConfig: React.FC<RunConfigProps> = ({ config, onConfigChange }) => {
       <PopoverContent className="w-[400px]">
         <div className="space-y-4">
           <h3 className="font-semibold">Run Configuration</h3>
+          <p className="text-xs text-gray-500">
+            <a href="https://docs.inferable.ai/pages/runs#options" className="underline">
+              Options
+            </a> to invoke the run with
+          </p>
           <div className="space-x-2 flex flex-row items-center">
             <Checkbox
               checked={reasoningTraces}
@@ -91,34 +115,68 @@ const RunConfig: React.FC<RunConfigProps> = ({ config, onConfigChange }) => {
             />
           </div>
           <div className="space-y-2">
-            <Label htmlFor="input-structured-output">
-              Structured Output (JSON Schema)
+            <Label htmlFor="input-result-schema">
+              Result Schema (JSON Schema)
+            </Label>
+            <p className="text-xs text-gray-500">
+              A{" "}
+              <a href="https://docs.inferable.ai/pages/runs#resultschema" className="underline">
+                 Result Schema
+              </a>{" "}
+              that the run must return
+            </p>
+            <Textarea
+              id="input-result-schema"
+              value={resultSchema ?? ""}
+              onChange={(e) => setResultSchema(e.target.value || null)}
+              onBlur={() => {
+                if (resultSchema) {
+                  try {
+                    JSON.parse(resultSchema);
+                    setResultSchemaJsonInvalid(false);
+                  } catch (e) {
+                    setResultSchemaJsonInvalid(true);
+                  }
+                } else {
+                  setResultSchemaJsonInvalid(false);
+                }
+              }}
+            />
+            {resultSchemaJsonInvalid && (
+              <p className="text-xs text-red-500">Invalid JSON</p>
+            )}
+
+          </div>
+
+          <div className="space-y-2">
+            <Label htmlFor="input-run-context">
+              Run Context
             </Label>
             <p className="text-xs text-gray-500">
               A{" "}
-              <a href="https://json-schema.org/" className="underline">
-                JSON Schema
+              <a href="https://docs.inferable.ai/pages/runs#context" className="underline">
+                Run Context
               </a>{" "}
-              that describes the structured output the run must return.
+              object which is passed to all calls
             </p>
             <Textarea
-              id="input-structured-output"
-              value={structuredOutput ?? ""}
-              onChange={(e) => setStructuredOutput(e.target.value || null)}
+              id="input-run-context"
+              value={runContext ?? ""}
+              onChange={(e) => setRunContext(e.target.value || null)}
               onBlur={() => {
-                if (structuredOutput) {
+                if (runContext) {
                   try {
-                    JSON.parse(structuredOutput);
-                    setJsonInvalid(false);
+                    JSON.parse(runContext);
+                    setRunContextJsonInvalid(false);
                   } catch (e) {
-                    setJsonInvalid(true);
+                    setRunContextJsonInvalid(true);
                   }
                 } else {
-                  setJsonInvalid(false);
+                  setRunContextJsonInvalid(false);
                 }
               }}
             />
-            {jsonInvalid && (
+            {runContextJsonInvalid && (
               <p className="text-xs text-red-500">Invalid JSON</p>
             )}
           </div>
diff --git a/app/components/chat/sdk-commands.tsx b/app/components/chat/sdk-commands.tsx
index c1712bb3..ab5d102c 100644
--- a/app/components/chat/sdk-commands.tsx
+++ b/app/components/chat/sdk-commands.tsx
@@ -10,7 +10,7 @@ type SDKCommandsProps = {
   clusterId: string;
   config: {
     attachedFunctions: string[];
-    structuredOutput: string | null;
+    resultSchema: string | null;
     prompt: string;
   };
 };
@@ -30,9 +30,9 @@ const SDKCommands: React.FC<SDKCommandsProps> = ({ clusterId, config }) => {
       message: config.prompt,
     };
 
-    if (config.structuredOutput) {
+    if (config.resultSchema) {
       body.result = {
-        schema: config.structuredOutput,
+        schema: config.resultSchema,
       };
     }
 
@@ -54,8 +54,8 @@ const SDKCommands: React.FC<SDKCommandsProps> = ({ clusterId, config }) => {
       command += `\\\n  --attachedFunctions '${config.attachedFunctions.join(",")}' `;
     }
 
-    if (config.structuredOutput) {
-      command += `\\\n  --resultSchema '${JSON.stringify(JSON.parse(config.structuredOutput), null, 2)}' `;
+    if (config.resultSchema) {
+      command += `\\\n  --resultSchema '${JSON.stringify(JSON.parse(config.resultSchema), null, 2)}' `;
     }
 
     command += `"${config.prompt}" `;