-
Notifications
You must be signed in to change notification settings - Fork 347
Description
For various reasons (testing, separation of concerns, etc) I need to get the plain interface of the BamlAsyncClient. Here's what I mean by plain interface:
// Current "type" of the BamlAsyncClient class ('typeof b')
interface IBamlAsyncClient {
stream(): xyz;
request(): xyz;
// etc
// Actual methods
ExtractResume(resume: string, __baml_options__?: BamlCallOptions): Promise<Resume>;
}
// What I want generated
interface IBamlAsyncClient {
ExtractResume(resume: string): Promise<Resume>;
}This interface allows me to do things like:
- Type test stubs correctly
- Have agent implementations that have the prompt methods injected rather than the underlying client injected
- Separate my "agent" code from the underlying physical call to the BAML Client
- This is useful because our agents usually run in Temporal.io workflows - I can take the above interface, and expose the underlying call to the BAML client via Temporal 'activities'
My current workaround is to export a barrel file as follows (this is generic and I generate it automatically for all my Prompt modules):
// index-types.ts
// Generated by Gemini, hence all the comments :P
import type { b } from './baml_client';
type BamlCallOptions = {
tb?: any
clientRegistry?: any
collector?: any
}
type BamlAsyncClient = typeof b;
// Utility type to modify the function signature to remove the last argument
// if it is BamlCallOptions or optional BamlCallOptions.
// It uses conditional type inference with rest parameters to achieve this.
type ModifyMethodSignature<F> =
// Case 1: Check if F matches a function signature where the last argument is exactly BamlCallOptions.
F extends (...args: [...infer P, BamlCallOptions]) => infer R
// If it matches, return a new function signature with parameters P (all args except the last one).
? (...args: P) => R
// Case 2: Check if F matches a function signature where the last argument is optional BamlCallOptions.
// The `(BamlCallOptions | undefined)?` pattern handles the optional case.
: F extends (...args: [...infer P, (BamlCallOptions | undefined)?]) => infer R
// If it matches, return a new function signature with parameters P.
? (...args: P) => R
// Default: If neither case matches (e.g., the function doesn't have BamlCallOptions as the last arg,
// or F is not a function type), return the original type F.
: F;
// Prompts type will contain only the async methods from BamlAsyncClient,
// with the BamlCallOptions argument removed if present as the last argument.
type Prompts = {
// Iterate over keys K of BamlAsyncClient
[K in keyof BamlAsyncClient as
// Filter keys: only include K if the property BamlAsyncClient[K] is an async function
// (i.e., a function returning a Promise).
BamlAsyncClient[K] extends (...args: any[]) => Promise<any> ? K : never
]:
// For the included keys K, get the type of the property BamlAsyncClient[K].
BamlAsyncClient[K] extends (...args: any[]) => Promise<any>
// If it's confirmed to be an async function, apply ModifyMethodSignature to its type.
? ModifyMethodSignature<BamlAsyncClient[K]>
// This 'never' branch should not be reachable due to the 'as K' filter above,
// but it's needed for type completeness.
: never;
};
export default Prompts;This exports the Prompts as:
type Prompts = {
ExtractResume: (resume: string) => Promise<Resume>;
}One key thing for me also is that the file this is generated in imports no concrete runtime code - ideally this is a types-only file. This means that when the Temporal workflow sandbox bundler runs (which bans any imports of modules that can cause side-effects, calls to fetch, etc), I can import this file directly and it all just works.
I imagine this is a fairly simple lift in the codegen so if this is amenable to the maintainers then I'm happy to open a PR to do it.