From 0401f1dc078043a195dcdba5a4de5c647b7ee68d Mon Sep 17 00:00:00 2001
From: Pavel Jbanov <pavelj@google.com>
Date: Tue, 4 Feb 2025 23:27:08 -0500
Subject: [PATCH] feat(js/client): added defineRemoteAction remote action
 client

---
 js/core/package.json           |  9 +++++
 js/core/src/zod.ts             | 17 ++++++++++
 js/genkit/src/client/client.ts | 60 ++++++++++++++++++++++++++++++++++
 js/genkit/src/client/index.ts  | 36 +++++++++++++++++++-
 4 files changed, 121 insertions(+), 1 deletion(-)
 create mode 100644 js/core/src/zod.ts

diff --git a/js/core/package.json b/js/core/package.json
index 6dd8671fd..605ac7c51 100644
--- a/js/core/package.json
+++ b/js/core/package.json
@@ -92,6 +92,12 @@
       "require": "./lib/schema.js",
       "import": "./lib/schema.mjs",
       "default": "./lib/schema.js"
+    },
+    "./zod": {
+      "types": "./lib/zod.d.ts",
+      "require": "./lib/zod.js",
+      "import": "./lib/zod.mjs",
+      "default": "./lib/zod.js"
     }
   },
   "typesVersions": {
@@ -116,6 +122,9 @@
       ],
       "schema": [
         "lib/schema"
+      ],
+      "zod": [
+        "lib/zod"
       ]
     }
   }
diff --git a/js/core/src/zod.ts b/js/core/src/zod.ts
new file mode 100644
index 000000000..31d911f70
--- /dev/null
+++ b/js/core/src/zod.ts
@@ -0,0 +1,17 @@
+/**
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+export { z } from 'zod';
diff --git a/js/genkit/src/client/client.ts b/js/genkit/src/client/client.ts
index dd42cabb8..368da2749 100644
--- a/js/genkit/src/client/client.ts
+++ b/js/genkit/src/client/client.ts
@@ -15,6 +15,66 @@
  */
 
 import { Channel } from '@genkit-ai/core/async';
+import { z } from '@genkit-ai/core/zod';
+
+interface RunOptions {
+  headers?: Record<string, string>;
+}
+
+export type RemoteAction<I = any, O = any, S = any> = ((
+  input?: I,
+  options?: RunOptions
+) => Promise<O>) & {
+  stream(input?: I, opts?: RunOptions): StreamingResponse<O, S>;
+};
+
+/**
+ * Streaming response from an action.
+ */
+export interface StreamingResponse<O = any, S = any> {
+  /** Iterator over the streaming chunks. */
+  stream: AsyncIterable<S>;
+  /** Final output of the action. */
+  output: Promise<O>;
+}
+
+/**
+ * Defines a remote action which can be invoked as a function or streamed.
+ */
+export function defineRemoteAction<
+  I extends z.ZodTypeAny = z.ZodTypeAny,
+  O extends z.ZodTypeAny = z.ZodTypeAny,
+  S extends z.ZodTypeAny = z.ZodTypeAny,
+>(actionOptions: {
+  url: string;
+  headers?: Record<string, string>;
+  inputSchema?: I;
+  outputSchema?: O;
+  streamSchema?: S;
+}): RemoteAction<z.infer<I>, z.infer<O>, z.infer<S>> {
+  const action = (
+    input?: z.infer<I>,
+    runOptions?: RunOptions
+  ): Promise<z.infer<O>> => {
+    return runFlow({
+      url: actionOptions.url,
+      headers: runOptions?.headers ?? actionOptions.headers,
+      input,
+    });
+  };
+  (action as RemoteAction<z.infer<I>, z.infer<O>, z.infer<S>>).stream = (
+    input?: I,
+    runOptions?: RunOptions
+  ): StreamingResponse<z.infer<O>, z.infer<S>> => {
+    return streamFlow({
+      url: actionOptions.url,
+      headers: runOptions?.headers ?? actionOptions.headers,
+      input,
+    });
+  };
+
+  return action as RemoteAction<z.infer<I>, z.infer<O>, z.infer<S>>;
+}
 
 const __flowStreamDelimiter = '\n\n';
 
diff --git a/js/genkit/src/client/index.ts b/js/genkit/src/client/index.ts
index 154de1604..e6b923be6 100644
--- a/js/genkit/src/client/index.ts
+++ b/js/genkit/src/client/index.ts
@@ -1,4 +1,6 @@
 /**
+ * @license
+ *
  * Copyright 2024 Google LLC
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
@@ -14,4 +16,36 @@
  * limitations under the License.
  */
 
-export { runFlow, streamFlow } from './client.js';
+/**
+ * A client library for remotely invoking deployed flows and other actions (e.g. models).
+ *
+ * ```ts
+ * import { defineRemoteAction } from 'genkit/beta/client';
+ *
+ * const myFlow = defineRemoteAction({
+ *   url: 'http://.../myFlow',
+ *   inputSchema: z.object({ foo: z.string() }),
+ *   outputSchema: z.string(),
+ *   streamSchema: z.string(),
+ * });
+ *
+ * const { stream } = myFlow.stream(
+ *   { foo: 'bar' },
+ *   { headers: { authentication: getAuthToken() } }
+ * );
+ *
+ * for await (const chunk of stream) {
+ *   console.log(chunk);
+ * }
+ * ```
+ *
+ * @module
+ */
+
+export {
+  RemoteAction,
+  StreamingResponse,
+  defineRemoteAction,
+  runFlow,
+  streamFlow,
+} from './client.js';