Skip to content

Commit

Permalink
feat(js/client): added defineRemoteAction remote action client
Browse files Browse the repository at this point in the history
  • Loading branch information
pavelgj committed Feb 5, 2025
1 parent 58d0bb8 commit 0401f1d
Show file tree
Hide file tree
Showing 4 changed files with 121 additions and 1 deletion.
9 changes: 9 additions & 0 deletions js/core/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -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": {
Expand All @@ -116,6 +122,9 @@
],
"schema": [
"lib/schema"
],
"zod": [
"lib/zod"
]
}
}
Expand Down
17 changes: 17 additions & 0 deletions js/core/src/zod.ts
Original file line number Diff line number Diff line change
@@ -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';
60 changes: 60 additions & 0 deletions js/genkit/src/client/client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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';

Expand Down
36 changes: 35 additions & 1 deletion js/genkit/src/client/index.ts
Original file line number Diff line number Diff line change
@@ -1,4 +1,6 @@
/**
* @license
*
* Copyright 2024 Google LLC
*
* Licensed under the Apache License, Version 2.0 (the "License");
Expand All @@ -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';

0 comments on commit 0401f1d

Please sign in to comment.