Skip to content

Commit

Permalink
feat: move @boundaryml/baml errors into it's own file
Browse files Browse the repository at this point in the history
feat: update types /client/ server to work better with streaming
  • Loading branch information
seawatts committed Jan 27, 2025
1 parent ee273b1 commit f367583
Show file tree
Hide file tree
Showing 26 changed files with 213 additions and 619 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -8,54 +8,53 @@ import type {
NonStreamingProps,
StreamingHookResult,
NonStreamingHookResult,
BamlHookProps,
BamlHookResult
HookProps,
HookResult
} from './types';
import type { RecursivePartialNull } from '../types';
import type { RecursivePartialNull, Check, Checked } from '../types';
import type { Image, Audio } from "@boundaryml/baml"
import * as ServerActions from './server';
import type { Check, Checked, RecursivePartialNull } from "../types"

import type {
{%- for t in types %}{{ t }}{% if !loop.last %}, {% endif %}{% endfor -%}
} from "../types"

// Type guard functions
function isPartialResponse<T>(obj: any): obj is PartialResponse<T> {
return obj && 'partial' in obj && !('final' in obj);
function isPartialResponse<T>(obj: unknown): obj is PartialResponse<T> {
return obj && typeof obj === 'object' && 'partial' in obj && !('final' in obj);
}

function isFinalResponse<T>(obj: any): obj is FinalResponse<T> {
return obj && 'final' in obj && !('partial' in obj);
function isFinalResponse<T>(obj: unknown): obj is FinalResponse<T> {
return obj && typeof obj === 'object' && 'final' in obj && !('partial' in obj);
}

/**
* Type guard to check if props are for streaming mode
*/
function isStreamingProps<Action>(
props: BamlHookProps<Action>
props: HookProps<Action>
): props is StreamingProps<Action> {
return props.stream === true;
}

interface BamlHookState<TPartial, TFinal> {
interface HookState<TPartial, TFinal> {
isSuccess: boolean;
error: Error | null;
data: TFinal | null;
partialData: TPartial | null;
}

type BamlHookStateAction<TPartial, TFinal> =
type HookStateAction<TPartial, TFinal> =
| { type: 'START_REQUEST' }
| { type: 'SET_ERROR'; payload: Error }
| { type: 'SET_PARTIAL'; payload: TPartial }
| { type: 'SET_FINAL'; payload: TFinal }
| { type: 'RESET' };

function bamlHookReducer<TPartial, TFinal>(
state: BamlHookState<TPartial, TFinal>,
action: BamlHookStateAction<TPartial, TFinal>
): BamlHookState<TPartial, TFinal> {
function hookReducer<TPartial, TFinal>(
state: HookState<TPartial, TFinal>,
action: HookStateAction<TPartial, TFinal>
): HookState<TPartial, TFinal> {
switch (action.type) {
case 'START_REQUEST':
return {
Expand Down Expand Up @@ -232,27 +231,27 @@ export function useBamlAction<Action>(

export function useBamlAction<Action>(
serverAction: Action,
props: BamlHookProps<Action> = {},
props: HookProps<Action> = {},
): StreamingHookResult<Action> | NonStreamingHookResult<Action> {
const { onFinal, onError, onPartial } = props;
const isStreaming = isStreamingProps(props);
const [isPending, startTransition] = useTransition();

const [state, dispatch] = useReducer(bamlHookReducer<RecursivePartialNull<Awaited<ReturnType<Action>>>, Awaited<ReturnType<Action>>>, {
const [state, dispatch] = useReducer(hookReducer<RecursivePartialNull<Awaited<ReturnType<Action>>>, Awaited<ReturnType<Action>>>, {
isSuccess: false,
error: null,
data: null,
partialData: null,
});

const mutate = useCallback(
async (input: Parameters<typeof serverAction>[0]) => {
async (...input: Parameters<typeof serverAction>) => {
dispatch({ type: 'START_REQUEST' });

try {
let response;
let response: Awaited<ReturnType<Action>>;
await startTransition(async () => {
response = await serverAction(input, { stream: props.stream });
response = await serverAction(...input);

if (isStreaming && response instanceof ReadableStream) {
const reader = response.getReader();
Expand Down Expand Up @@ -455,7 +454,7 @@ export function use{{ func.name }}<Action extends typeof ServerActions.{{ func.n
): NonStreamingHookResult<Action>;

export function use{{ func.name }}<Action extends typeof ServerActions.{{ func.name }}Action>(
props: BamlHookProps<Action> = {},
props: HookProps<Action> = {},
): StreamingHookResult<Action> | NonStreamingHookResult<Action> {
if (props.stream) {
return useBamlAction(
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,15 @@ export type FinalResponse<Output> = {
final: Output
}

export type ServerAction<Input, Output> = (...input: Input) => Promise<ReadableStream<Uint8Array> | Output>;
export type ServerAction<Input = any, Output = any> = (...args: Input extends any[] ? Input : [Input]) => Promise<ReadableStream<Uint8Array> | Output>;

/**
* Props for streaming mode, which provides incremental updates.
* Use this when you want to show partial results as they become available.
*
* @template Action The server action type
*/
export type StreamingProps<Action extends ServerAction<any, any> = ServerAction<any, any>> = {
export type StreamingProps<Action extends ServerAction = ServerAction> = {
stream: true
onPartial?: (response?: RecursivePartialNull<Awaited<ReturnType<Action>>>) => void
onFinal?: (response?: Awaited<ReturnType<Action>>) => void
Expand All @@ -39,7 +39,7 @@ export type StreamingProps<Action extends ServerAction<any, any> = ServerAction<
* Props for non-streaming mode.
* @template Action The server action type
*/
export type NonStreamingProps<Action extends ServerAction<any, any> = ServerAction<any, any>> = {
export type NonStreamingProps<Action extends ServerAction = ServerAction> = {
stream?: false
onPartial?: never
onFinal?: (response?: Awaited<ReturnType<Action>>) => void
Expand All @@ -51,12 +51,12 @@ export type NonStreamingProps<Action extends ServerAction<any, any> = ServerActi
* Union type of all possible props for a BAML hook.
* @template Action The server action type
*/
export type BamlHookProps<Action extends ServerAction<any, any> = ServerAction<any, any>> = StreamingProps<Action> | NonStreamingProps<Action>
export type HookProps<Action extends ServerAction = ServerAction> = StreamingProps<Action> | NonStreamingProps<Action>

/**
* Base return type for all BAML hooks
*/
export type BamlHookResult<Action extends ServerAction<any, any> = ServerAction<any, any>> = {
export type BaseHookResult<Action extends ServerAction = ServerAction> = {
/**
* The complete, final result of the operation.
* Only available after successful completion (when isSuccess is true).
Expand Down Expand Up @@ -96,7 +96,7 @@ export type BamlHookResult<Action extends ServerAction<any, any> = ServerAction<
/**
* Return type for streaming mode BAML hooks
*/
export type StreamingHookResult<Action extends ServerAction<any, any> = ServerAction<any, any>> = BamlHookResult<Action> & {
export type StreamingHookResult<Action extends ServerAction = ServerAction> = BaseHookResult<Action> & {
/**
* The most recent partial result from the stream.
* Updates continuously while streaming, showing interim progress.
Expand All @@ -109,24 +109,24 @@ export type StreamingHookResult<Action extends ServerAction<any, any> = ServerAc
* Call this function to start the operation.
* Returns a promise that resolves with the final result or null if it failed.
*/
mutate: (input: Parameters<Action>[0]) => Promise<ReadableStream<Uint8Array>>;
mutate: (...input: Parameters<Action>) => Promise<ReadableStream<Uint8Array>>;
};

/**
* Return type for non-streaming mode BAML hooks
*/
export type NonStreamingHookResult<Action extends ServerAction<any, any> = ServerAction<any, any>> = BamlHookResult<Action> & {
export type NonStreamingHookResult<Action extends ServerAction = ServerAction> = BaseHookResult<Action> & {
/** Not available in non-streaming mode */
partialData?: never;
mutate: (input: Parameters<Action>[0]) => Promise<Awaited<ReturnType<Action>>>;
mutate: (...input: Parameters<Action>) => Promise<Awaited<ReturnType<Action>>>;
};

/**
* Conditional return type for BAML hooks based on the provided props
*/
export type BamlHookResult<
Action extends ServerAction<any, any> = ServerAction<any, any>,
Props extends BamlHookProps<Action> = BamlHookProps<Action>
export type HookResult<
Action extends ServerAction = ServerAction,
Props extends HookProps<Action> = HookProps<Action>
> = Props extends { stream: true }
? StreamingHookResult<Action>
: NonStreamingHookResult<Action>;
15 changes: 14 additions & 1 deletion engine/language_client_typescript/package.json
Original file line number Diff line number Diff line change
Expand Up @@ -30,10 +30,23 @@
"./artifacts/stream.d.ts",
"./artifacts/stream.js",
"./artifacts/type_builder.d.ts",
"./artifacts/type_builder.js"
"./artifacts/type_builder.js",
"./artifacts/errors.d.ts",
"./artifacts/errors.js"
],
"main": "./artifacts/index.js",
"types": "./artifacts/index.d.ts",
"exports": {
".": {
"types": "./artifacts/index.d.ts",
"node": "./artifacts/index.js",
"default": "./artifacts/client.js"
},
"./errors": {
"types": "./artifacts/errors.d.ts",
"default": "./artifacts/errors.js"
}
},
"napi": {
"binaryName": "baml",
"targets": [
Expand Down
111 changes: 111 additions & 0 deletions engine/language_client_typescript/typescript_src/errors.ts
Original file line number Diff line number Diff line change
@@ -0,0 +1,111 @@

export class BamlClientFinishReasonError extends Error {
prompt: string;
raw_output: string;

constructor(prompt: string, raw_output: string, message: string) {
super(message);
this.name = "BamlClientFinishReasonError";
this.prompt = prompt;
this.raw_output = raw_output;

Object.setPrototypeOf(this, BamlClientFinishReasonError.prototype);
}

toJSON(): string {
return JSON.stringify(
{
name: this.name,
message: this.message,
raw_output: this.raw_output,
prompt: this.prompt,
},
null,
2
);
}

static from(error: Error): BamlClientFinishReasonError | undefined {
if (error.message.includes("BamlClientFinishReasonError")) {
try {
const errorData = JSON.parse(error.message);
if (errorData.type === "BamlClientFinishReasonError") {
return new BamlClientFinishReasonError(
errorData.prompt || "",
errorData.raw_output || "",
errorData.message || error.message
);
} else {
console.warn("Not a BamlClientFinishReasonError:", error);
}
} catch (parseError) {
// If JSON parsing fails, fall back to the original error
console.warn("Failed to parse BamlClientFinishReasonError:", parseError);
}
}
return undefined;
}
}

export class BamlValidationError extends Error {
prompt: string;
raw_output: string;

constructor(prompt: string, raw_output: string, message: string) {
super(message);
this.name = "BamlValidationError";
this.prompt = prompt;
this.raw_output = raw_output;

Object.setPrototypeOf(this, BamlValidationError.prototype);
}

toJSON(): string {
return JSON.stringify(
{
name: this.name,
message: this.message,
raw_output: this.raw_output,
prompt: this.prompt,
},
null,
2
);
}

static from(error: Error): BamlValidationError | undefined {
if (error.message.includes("BamlValidationError")) {
try {
const errorData = JSON.parse(error.message);
if (errorData.type === "BamlValidationError") {
return new BamlValidationError(
errorData.prompt || "",
errorData.raw_output || "",
errorData.message || error.message
);
}
} catch (parseError) {
console.warn("Failed to parse BamlValidationError:", parseError);
}
}
return undefined;
}
}

// Helper function to safely create a BamlValidationError
export function createBamlValidationError(
error: Error
): BamlValidationError | BamlClientFinishReasonError | Error {
const bamlValidationError = BamlValidationError.from(error);
if (bamlValidationError) {
return bamlValidationError;
}

const bamlClientFinishReasonError = BamlClientFinishReasonError.from(error);
if (bamlClientFinishReasonError) {
return bamlClientFinishReasonError;
}

// otherwise return the original error
return error;
}
Loading

0 comments on commit f367583

Please sign in to comment.