Simplify the GenSX api.

Author: jmoseley

playground/index.tsx

@@ -1,8 +1,4 @@
-import {
-  executeJsxWorkflow,
-  withWorkflowComponent,
-  withWorkflowFunction,
-} from "@/index";
+import { Component, gensx } from "@/index";
 
 // Pure workflow steps
 const pureResearchBrainstorm = async ({ prompt }: { prompt: string }) => {
@@ -28,66 +24,78 @@ const pureEditor = async ({ draft }: { draft: string }) => {
 };
 
 // Wrapped workflow steps
-const LLMResearchBrainstorm = withWorkflowFunction(pureResearchBrainstorm);
-const LLMResearch = withWorkflowFunction(
-  async ({ topic }: { topic: string }) => {
-    console.log("📚 Researching topic:", topic);
-    return await Promise.resolve(`research results for ${topic}`);
-  },
-);
-const LLMWriter = withWorkflowFunction(pureWriter);
-const LLMEditor = withWorkflowFunction(pureEditor);
-const WebResearcher = withWorkflowFunction(
-  async ({ prompt }: { prompt: string }) => {
-    console.log("🌐 Researching web for:", prompt);
-    const results = await Promise.resolve([
-      "web result 1",
-      "web result 2",
-      "web result 3",
-    ]);
-    return results;
-  },
-);
-
-// Research collection component
-const ResearchCollection = withWorkflowComponent<
-  [string[], string[]],
-  { prompt: string }
->(({ prompt }: { prompt: string }) => (
-  <>
-    <LLMResearchBrainstorm prompt={prompt}>
-      {topics => (
-        <>
-          {topics.map(topic => (
-            <LLMResearch topic={topic} />
-          ))}
-        </>
-      )}
-    </LLMResearchBrainstorm>
-    <WebResearcher prompt={prompt} />
-  </>
-));
+const LLMResearchBrainstorm = Component(pureResearchBrainstorm);
+const LLMResearch = Component(async ({ topic }: { topic: string }) => {
+  console.log("📚 Researching topic:", topic);
+  return await Promise.resolve(`research results for ${topic}`);
+});
+const LLMWriter = Component(pureWriter);
+const LLMEditor = Component(pureEditor);
+const WebResearcher = Component(async ({ prompt }: { prompt: string }) => {
+  console.log("🌐 Researching web for:", prompt);
+  const results = await Promise.resolve([
+    "web result 1",
+    "web result 2",
+    "web result 3",
+  ]);
+  return results;
+});
 
-const BlogWritingWorkflow = withWorkflowComponent<string, { prompt: string }>(
+// When building a workflow out of components, there are two options:
+// 1. Use the Component function to wrap it and specify the input and output types. This gets you a function with type safe inputs and outputs (if you just call it as a function).
+// 2. Don't wrap it in the Component function, and do not specify the output type. You get a function that matches the JSX.Element signature, but it has an unknown output type. If you try to specify the output type, you get a type error (unknown is not assignable to X).
+const ParallelResearch = Component<{ prompt: string }, [string[], string[]]>(
   ({ prompt }: { prompt: string }) => (
-    <ResearchCollection prompt={prompt}>
-      {([catalogResearch, webResearch]) => {
-        console.log("🧠 Research:", { catalogResearch, webResearch });
-        return (
-          <LLMWriter research={catalogResearch.join("\n")} prompt={prompt}>
-            {draft => <LLMEditor draft={draft} />}
-          </LLMWriter>
-        );
-      }}
-    </ResearchCollection>
+    <>
+      <LLMResearchBrainstorm prompt={prompt}>
+        {topics => (
+          <>
+            {topics.map(topic => (
+              <LLMResearch topic={topic} />
+            ))}
+          </>
+        )}
+      </LLMResearchBrainstorm>
+      <WebResearcher prompt={prompt} />
+    </>
   ),
 );
+
+const BlogWritingWorkflow = async ({ prompt }: { prompt: string }) => (
+  <ParallelResearch prompt={prompt}>
+    {([catalogResearch, webResearch]) => {
+      console.log("🧠 Research:", { catalogResearch, webResearch });
+      return (
+        <LLMWriter research={catalogResearch.join("\n")} prompt={prompt}>
+          {draft => <LLMEditor draft={draft} />}
+        </LLMWriter>
+      );
+    }}
+  </ParallelResearch>
+);
 async function main() {
   console.log("🚀 Starting blog writing workflow");
-  const result = await executeJsxWorkflow<string>(
+
+  // Use the gensx function to execute the workflow and annotate with the output type.
+  const result = await gensx<string>(
     <BlogWritingWorkflow prompt="Write a blog post about the future of AI" />,
   );
-  console.log("✅ Final result:", result);
+
+  // Or just call the workflow as a function, and cast to the output type.
+  const result2 = (await (
+    <BlogWritingWorkflow prompt="Write a blog post about using GenSX to 10x your productivity." />
+  )) as string;
+
+  // Still need to cast here, because we didn't use the Component helper to wrap the workflow.
+  const result3 = (await BlogWritingWorkflow({
+    prompt: "Write a blog post about the future of AI",
+  })) as string;
+
+  // Don't need to cast here, because we used the Component helper to wrap the workflow.
+  const researchResult = await ParallelResearch({
+    prompt: "Write a blog post about the future of AI",
+  });
+  console.log("✅ Final result:", { result, result2, result3, researchResult });
 }
 
 await main();

src/component.ts

@@ -0,0 +1,14 @@
+import { MaybePromise } from "./jsx-runtime";
+
+export function Component<TInput extends Record<string, unknown>, TOutput>(
+  fn: (input: TInput) => MaybePromise<TOutput> | JSX.Element,
+) {
+  function WorkflowFunction(
+    props: TInput & {
+      children?: (output: TOutput) => MaybePromise<TOutput> | JSX.Element;
+    },
+  ): Promise<TOutput> {
+    return Promise.resolve(fn(props)) as Promise<TOutput>;
+  }
+  return WorkflowFunction;
+}

src/components/Collect.tsx

@@ -0,0 +1,3 @@
+export async function gensx<TOutput>(node: JSX.Element): Promise<TOutput> {
+  return (await node) as TOutput;
+}

src/gensx.ts

@@ -12,10 +12,7 @@
 7. Support parallel execution of steps (either dynamic via something liek a collector, or static via a few explicitly defined siblings)
  */
 
-import { executeJsxWorkflow } from "./workflow/execute";
-import {
-  withWorkflowComponent,
-  withWorkflowFunction,
-} from "./workflow/withWorkflowStep";
+import { Component } from "./component";
+import { gensx } from "./gensx";
 
-export { executeJsxWorkflow, withWorkflowComponent, withWorkflowFunction };
+export { Component, gensx };