Skip to content

Commit

Permalink
combine retry and structured transform, rename to StructuredOutput
Browse files Browse the repository at this point in the history
  • Loading branch information
@EvanBoyle
EvanBoyle committed Feb 1, 2025
1 parent 64a1db5 commit a04e366
Showing 1 changed file with 92 additions and 95 deletions.
187 changes: 92 additions & 95 deletions packages/gensx-openai/src/composition.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -61,21 +61,22 @@ type ToolTransformProps = Omit<

type ToolTransformReturn = RawCompletionReturn;

// Structured output transform component
type StructuredTransformProps<O = unknown> = Omit<
// Updated type to include retry options
type StructuredOutputProps<O = unknown> = Omit<
ChatCompletionCreateParamsNonStreaming,
"stream" | "tools"
> & {
structuredOutput: GSXStructuredOutput<O>;
tools?: GSXTool<any>[];
retry?: {
maxAttempts?: number;
backoff?: "exponential" | "linear";
onRetry?: (attempt: number, error: Error, lastResponse?: string) => void;
shouldRetry?: (error: Error, attempt: number) => boolean;
};
};

type StructuredTransformReturn<T> = T;

// Add RetryTransformProps type
type RetryTransformProps<O = unknown> = StructuredTransformProps<O> & {
maxAttempts?: number;
};
type StructuredOutputReturn<T> = T;

// Types for the composition-based implementation
type StreamingProps = Omit<
Expand Down Expand Up @@ -268,41 +269,75 @@ export const ToolTransform = gsx.Component<
);
});

// Structured output transform component
export const StructuredTransform = gsx.Component<
StructuredTransformProps,
StructuredTransformReturn<unknown>
>("StructuredTransform", async (props) => {
const { structuredOutput, tools, ...rest } = props;
// Combined structured output component
export const StructuredOutput = gsx.Component<
StructuredOutputProps,
StructuredOutputReturn<unknown>
>("StructuredOutput", async (props) => {
const { structuredOutput, tools, retry, ...rest } = props;
const maxAttempts = retry?.maxAttempts ?? 3;
let lastError: Error | undefined;
let lastResponse: string | undefined;

// Make initial completion
const completion = await gsx.execute<ChatCompletionOutput>(
<RawCompletion
{...rest}
tools={tools}
response_format={structuredOutput.toResponseFormat()}
/>,
);
for (let attempt = 1; attempt <= maxAttempts; attempt++) {
try {
// Add retry context to messages if not first attempt
const messages = [...rest.messages];
if (attempt > 1) {
messages.push({
role: "system",
content: `Previous attempt failed: ${lastError?.message}. Please fix the JSON structure and try again.`,
});
}

const toolCalls = completion.choices[0].message.tool_calls;
// If we have tool calls, execute them and make another completion
if (toolCalls?.length && tools) {
const toolResult = await gsx.execute<ChatCompletionOutput>(
<ToolExecutor
tools={tools}
toolCalls={toolCalls}
messages={[...rest.messages, completion.choices[0].message]}
model={rest.model}
/>,
);
// Make initial completion
const completion = await gsx.execute<ChatCompletionOutput>(
<RawCompletion
{...rest}
messages={messages}
tools={tools}
response_format={structuredOutput.toResponseFormat()}
/>,
);

// Parse and validate the final result
const content = toolResult.choices[0]?.message.content;
if (!content) {
throw new Error("No content returned from OpenAI after tool execution");
}
const toolCalls = completion.choices[0].message.tool_calls;
// If we have tool calls, execute them and make another completion
if (toolCalls?.length && tools) {
const toolResult = await gsx.execute<ChatCompletionOutput>(
<ToolExecutor
tools={tools}
toolCalls={toolCalls}
messages={[...messages, completion.choices[0].message]}
model={rest.model}
/>,
);

try {
// Parse and validate the final result
const content = toolResult.choices[0]?.message.content;
if (!content) {
throw new Error(
"No content returned from OpenAI after tool execution",
);
}

lastResponse = content;
const parsed = JSON.parse(content) as unknown;
const validated = structuredOutput.safeParse(parsed);
if (!validated.success) {
throw new Error(
`Invalid structured output: ${validated.error.message}`,
);
}
return validated.data;
}

// No tool calls, parse and validate the direct result
const content = completion.choices[0].message.content;
if (!content) {
throw new Error("No content returned from OpenAI");
}

lastResponse = content;
const parsed = JSON.parse(content) as unknown;
const validated = structuredOutput.safeParse(parsed);
if (!validated.success) {
Expand All @@ -312,70 +347,32 @@ export const StructuredTransform = gsx.Component<
}
return validated.data;
} catch (e) {
throw new Error(
`Failed to parse structured output after tool execution: ${e instanceof Error ? e.message : String(e)}`,
);
}
}

// No tool calls, parse and validate the direct result
const content = completion.choices[0].message.content;
if (!content) {
throw new Error("No content returned from OpenAI");
}

try {
const parsed = JSON.parse(content) as unknown;
const validated = structuredOutput.safeParse(parsed);
if (!validated.success) {
throw new Error(`Invalid structured output: ${validated.error.message}`);
}
return validated.data;
} catch (e) {
throw new Error(
`Failed to parse structured output: ${e instanceof Error ? e.message : String(e)}`,
);
}
});
lastError = e instanceof Error ? e : new Error(String(e));

// Retry transform component for structured output
export const RetryTransform = gsx.Component<
RetryTransformProps,
StructuredTransformReturn<unknown>
>("RetryTransform", async (props) => {
const { maxAttempts = 3, ...rest } = props;
let lastError: string | undefined;
// Call onRetry callback if provided
retry?.onRetry?.(attempt, lastError, lastResponse);

for (let attempt = 1; attempt <= maxAttempts; attempt++) {
try {
const retryMessage: ChatCompletionMessageParam | undefined =
attempt > 1
? {
role: "system",
content: `Previous attempt failed: ${lastError}. Please fix the JSON structure and try again.`,
}
: undefined;

const messages = [...rest.messages];
if (retryMessage) {
messages.push(retryMessage);
}

return await gsx.execute(
<StructuredTransform {...rest} messages={messages} />,
);
} catch (e) {
lastError = e instanceof Error ? e.message : String(e);
if (attempt === maxAttempts) {
// Check if we should retry
const shouldRetry = retry?.shouldRetry?.(lastError, attempt) ?? true;
if (!shouldRetry || attempt === maxAttempts) {
throw new Error(
`Failed to get valid structured output after ${maxAttempts} attempts. Last error: ${lastError}`,
`Failed to get valid structured output after ${attempt} attempts. Last error: ${lastError.message}`,
);
}

// Apply backoff if specified
if (retry?.backoff) {
const delay =
retry.backoff === "exponential"
? Math.pow(2, attempt - 1) * 1000
: attempt * 1000;
await new Promise((resolve) => setTimeout(resolve, delay));
}
}
}
});

// The composition-based implementation that matches ChatCompletion's functionality
// Update CompositionCompletion to use the renamed component
export const CompositionCompletion = gsx.Component<
CompositionCompletionProps,
CompositionCompletionReturn<CompositionCompletionProps>
Expand All @@ -390,7 +387,7 @@ export const CompositionCompletion = gsx.Component<
if ("structuredOutput" in props && props.structuredOutput) {
const { tools, structuredOutput, ...rest } = props;
return (
<RetryTransform
<StructuredOutput
{...rest}
tools={tools}
structuredOutput={structuredOutput}
Expand Down

0 comments on commit a04e366

Please sign in to comment.