feat: improving the checkpoint output for tools and structured outputs

examples/anthropic/index.tsx

@@ -375,7 +375,7 @@ async function main() {
     | "multiStepTools"
     | "toolsWithStructuredOutput";
 
-  const example: Example = "tools";
+  const example: Example = "toolsWithStructuredOutput";
 
   switch (example as Example) {
     case "basicCompletion":

packages/gensx-anthropic/src/gsx-completion.tsx

@@ -1,5 +1,8 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
+// Import Zod extensions for improved serialization
+import "./utils/zod-extensions.js";
+
 import {
   Message,
   MessageCreateParamsNonStreaming,

packages/gensx-anthropic/src/tools.tsx

@@ -65,6 +65,14 @@ export class GSXTool<TSchema extends z.ZodObject<z.ZodRawShape>> {
     return new GSXTool(params);
   }
 
+  toJSON(): Record<string, unknown> {
+    return {
+      name: this.name,
+      description: this.description,
+      schema: this.definition.input_schema,
+    };
+  }
+
   public readonly name: string;
   public readonly description: string;
   public readonly schema: TSchema;

packages/gensx-anthropic/src/utils/zod-extensions.ts

@@ -0,0 +1,30 @@
+/**
+ * Zod Schema Extensions
+ *
+ * This module provides extensions to Zod schemas that improve their
+ * serialization capabilities when used with GenSX's checkpoint system.
+ */
+
+import { z } from "zod";
+import { zodToJsonSchema } from "zod-to-json-schema";
+
+export function extendZodWithToJSON(): void {
+  // Use type assertion to check for toJSON method
+  const prototype = z.ZodSchema.prototype as unknown as {
+    toJSON?: () => Record<string, unknown>;
+  };
+
+  // Only add the method once to avoid overriding
+  if (!prototype.toJSON) {
+    Object.defineProperty(z.ZodSchema.prototype, "toJSON", {
+      value: function (this: z.ZodType) {
+        return zodToJsonSchema(this);
+      },
+      configurable: true,
+      writable: true,
+    });
+  }
+}
+
+// Execute the extension immediately when this module is imported
+extendZodWithToJSON();

packages/gensx-openai/src/gsx-completion.tsx

@@ -1,5 +1,8 @@
 /* eslint-disable @typescript-eslint/no-explicit-any */
 
+// Import Zod extensions for improved serialization
+import "./utils/zod-extensions.js";
+
 import { Args, gsx } from "gensx";
 import {
   ChatCompletion as ChatCompletionOutput,

packages/gensx-openai/src/tools.tsx

@@ -65,6 +65,14 @@ export class GSXTool<TSchema extends z.ZodObject<z.ZodRawShape>> {
     return new GSXTool(params);
   }
 
+  toJSON(): Record<string, unknown> {
+    return {
+      name: this.name,
+      description: this.description,
+      schema: zodToJsonSchema(this.schema),
+    };
+  }
+
   public readonly name: string;
   public readonly description: string;
   public readonly schema: TSchema;

packages/gensx-openai/src/utils/zod-extensions.ts

@@ -0,0 +1,30 @@
+/**
+ * Zod Schema Extensions
+ *
+ * This module provides extensions to Zod schemas that improve their
+ * serialization capabilities when used with GenSX's checkpoint system.
+ */
+
+import { z } from "zod";
+import { zodToJsonSchema } from "zod-to-json-schema";
+
+export function extendZodWithToJSON(): void {
+  // Use type assertion to check for toJSON method
+  const prototype = z.ZodSchema.prototype as unknown as {
+    toJSON?: () => Record<string, unknown>;
+  };
+
+  // Only add the method once to avoid overriding
+  if (!prototype.toJSON) {
+    Object.defineProperty(z.ZodSchema.prototype, "toJSON", {
+      value: function (this: z.ZodType) {
+        return zodToJsonSchema(this);
+      },
+      configurable: true,
+      writable: true,
+    });
+  }
+}
+
+// Execute the extension immediately when this module is imported
+extendZodWithToJSON();

packages/gensx/src/checkpoint.ts

@@ -591,6 +591,12 @@ export class CheckpointManager implements CheckpointWriter {
     if (Symbol.asyncIterator in value) return value;
     if (ArrayBuffer.isView(value)) return value;
 
+    // Check for toJSON method before doing regular object cloning
+    const objValue = value as { toJSON?: () => unknown };
+    if (typeof objValue.toJSON === "function") {
+      return this.cloneValue(objValue.toJSON());
+    }
+
     // For regular objects, clone each property
     return Object.fromEntries(
       Object.entries(value).map(([key, val]) => [key, this.cloneValue(val)]),