feat: Remove zod-to-json-schema usage (#310)

* chore: Remove json-schema-to-zod * feat: Build model schema with JsonSchema * chore: Add agent result test
inferablehq · Dec 15, 2024 · 58e4fff · 58e4fff
1 parent 5c86cf7
commit 58e4fff
Show file tree

Hide file tree

Showing 8 changed files with 299 additions and 206 deletions.
diff --git a/control-plane/package-lock.json b/control-plane/package-lock.json
diff --git a/control-plane/package.json b/control-plane/package.json
@@ -46,7 +46,6 @@
     "inferable": "^0.30.59",
     "jest": "^29.6.4",
     "js-tiktoken": "^1.0.12",
-    "json-schema-to-zod": "^2.1.0",
     "jsonpath": "^1.1.1",
     "jsonschema": "^1.4.1",
     "jsonwebtoken": "^9.0.2",

diff --git a/control-plane/src/modules/service-definitions.test.ts b/control-plane/src/modules/service-definitions.test.ts
@@ -1,8 +1,6 @@
-import { dereferenceSync, JSONSchema } from "dereference-json-schema";
 import { InvalidJobArgumentsError, InvalidServiceRegistrationError } from "../utilities/errors";
 import { packer } from "./packer";
 import {
-  deserializeFunctionSchema,
   embeddableServiceFunction,
   parseJobArgs,
   serviceFunctionEmbeddingId,
@@ -214,76 +212,6 @@ describe("parseJobArgs", () => {
   });
 });
 
-describe("deserializeFunctionSchema", () => {
-  const jsonSchema = {
-    $schema: "http://json-schema.org/draft-04/schema#",
-    title: "ExtractResult",
-    type: "object",
-    additionalProperties: false,
-    properties: {
-      posts: {
-        type: "array",
-        items: {
-          $ref: "#/definitions/Post",
-        },
-      },
-    },
-    definitions: {
-      Post: {
-        type: "object",
-        additionalProperties: false,
-        properties: {
-          id: {
-            type: "string",
-          },
-          title: {
-            type: "string",
-          },
-          points: {
-            type: "string",
-          },
-          comments_url: {
-            type: "string",
-          },
-        },
-      },
-    },
-  };
-
-  it("should convert a JSON schema to a Zod schema", () => {
-    const zodSchema = deserializeFunctionSchema(
-      dereferenceSync(jsonSchema as any),
-    );
-    const jsonSchema2 = zodToJsonSchema(zodSchema);
-    const dereferenced = dereferenceSync(jsonSchema2 as JSONSchema);
-    expect(dereferenced).toMatchObject({
-      properties: {
-        posts: {
-          type: "array",
-          items: {
-            type: "object",
-            additionalProperties: false,
-            properties: {
-              id: {
-                type: "string",
-              },
-              title: {
-                type: "string",
-              },
-              points: {
-                type: "string",
-              },
-              comments_url: {
-                type: "string",
-              },
-            },
-          },
-        },
-      },
-    });
-  });
-});
-
 describe("validateServiceRegistration", () => {
   it("should reject invalid schema", () => {
     expect(() => {

diff --git a/control-plane/src/modules/service-definitions.ts b/control-plane/src/modules/service-definitions.ts
@@ -5,7 +5,6 @@ import {
   validateFunctionSchema,
   validateServiceName,
 } from "inferable";
-import jsonSchemaToZod, { JsonSchema } from "json-schema-to-zod";
 import { Validator } from "jsonschema";
 import { z } from "zod";
 import {
@@ -465,9 +464,8 @@ export const validateServiceRegistration = ({
       }
 
       // Check that the schema accepts and expected value
-      const zodSchema = deserializeFunctionSchema(fn.schema);
-      const schema = zodSchema.safeParse({ token: "test" });
-      if (!schema.success) {
+      const result = validator.validate({ token: "test" }, JSON.parse(fn.schema));
+      if (!result.valid) {
         throw new InvalidServiceRegistrationError(
           `${fn.name} schema is not valid`,
           "https://docs.inferable.ai/pages/auth#handlecustomerauth",
@@ -486,53 +484,6 @@ export const start = () =>
     },
   ); // 10 seconds
 
-/**
- * Convert a JSON schema (Object or String) to a Zod schema object
- */
-export const deserializeFunctionSchema = (schema: unknown) => {
-  if (typeof schema === "object") {
-    let zodSchema;
-
-    try {
-      zodSchema = jsonSchemaToZod(schema as JsonSchema);
-    } catch (e) {
-      logger.error("Failed to convert schema to Zod", { schema, error: e });
-      throw new Error("Failed to load the tool definition");
-    }
-
-    return eval(`
-const { z } = require("zod");
-${zodSchema}
-`);
-  } else if (typeof schema === "string") {
-    let parsed;
-
-    try {
-      parsed = JSON.parse(schema);
-    } catch (e) {
-      logger.error("Failed to parse schema", { schema, error: e });
-      throw new Error("Failed to parse the tool definition");
-    }
-
-    let zodSchema;
-
-    try {
-      zodSchema = jsonSchemaToZod(parsed);
-    } catch (e) {
-      logger.error("Failed to convert schema to Zod", { schema, error: e });
-      throw new Error("Failed to load the tool definition");
-    }
-
-    return eval(`
-const { z } = require("zod");
-${zodSchema}
-`);
-  } else {
-    logger.error("Invalid schema", { schema });
-    throw new Error("Invalid schema");
-  }
-};
-
 export const normalizeFunctionReference = (
   fn: string | { service: string; function: string },
 ) =>

diff --git a/control-plane/src/modules/workflows/agent/agent.ai.test.ts b/control-plane/src/modules/workflows/agent/agent.ai.test.ts
@@ -190,6 +190,56 @@ describe("Agent", () => {
     });
   });
 
+  describe("result schema", () => {
+    jest.setTimeout(120000);
+
+    it("should result result schema", async () => {
+
+      const app = await createWorkflowAgent({
+        workflow: {
+          ...workflow,
+          resultSchema: {
+            type: "object",
+            properties: {
+              word: {
+                type: "string"
+              }
+            }
+          }
+        },
+        findRelevantTools: async () => tools,
+        getTool: async () => tools[0],
+        postStepSave: async () => {},
+      });
+
+      const messages = [
+        {
+          type: "human",
+          data: {
+            message: "Return the word 'hello'",
+          },
+        },
+      ];
+
+      const outputState = await app.invoke({
+        workflow,
+        messages,
+      });
+
+      expect(outputState.messages).toHaveLength(2);
+      expect(outputState.messages[0]).toHaveProperty("type", "human");
+      expect(outputState.messages[1]).toHaveProperty("type", "agent");
+      expect(outputState.messages[1].data.result).toHaveProperty(
+        "word",
+        "hello",
+      );
+
+      expect(outputState.result).toEqual({
+        word: "hello"
+      });
+    });
+  });
+
   describe("early exit", () => {
     jest.setTimeout(120000);