refactor: replace simplifySchema with direct schema question functions

sroussey · sroussey · commit fd09e7e805bd · 2025-05-27T10:01:35.000-07:00
- Updated multiple files to remove the use of simplifySchema, replacing it with direct access to schema properties.
- This change enhances performance and simplifies the codebase by eliminating unnecessary schema simplification calls.
- Adjusted relevant logic in AiTask, FsFolderKvRepository, KvViaTabularRepository, BaseSqlTabularRepository, Dataflow, and Workflow classes to reflect this update.
diff --git a/packages/ai/src/task/base/AiTask.ts b/packages/ai/src/task/base/AiTask.ts
@@ -17,7 +17,7 @@ import {
   type TaskInput,
   type TaskOutput,
 } from "@ellmers/task-graph";
-import { simplifySchema } from "@ellmers/util";
+import { schemaSemantic } from "@ellmers/util";
 import { type TSchema } from "@sinclair/typebox";
 import { AiJob } from "../../job/AiJob";
 import { getGlobalModelRepository } from "../../model/ModelRegistry";
@@ -87,9 +87,8 @@ export class AiTask<
    */
   async validateInput(input: Input): Promise<boolean> {
     // TODO(str): this is very inefficient, we should cache the results, including intermediate results
-    const inputSchemaProperties = simplifySchema(this.inputSchema).properties;
-    const modelTaskProperties = Object.entries<TSchema>(inputSchemaProperties).filter(
-      ([key, value]) => value.semantic?.startsWith("model:")
+    const modelTaskProperties = Object.entries<TSchema>(this.inputSchema.properties).filter(
+      ([key, schema]) => schemaSemantic(schema)?.startsWith("model:")
     );
     if (modelTaskProperties.length > 0) {
       const taskModels = await getGlobalModelRepository().findModelsByTask(this.type);
@@ -103,8 +102,8 @@ export class AiTask<
         }
       }
     }
-    const modelPlainProperties = Object.entries<TSchema>(inputSchemaProperties).filter(
-      ([key, value]) => value.semantic === "model"
+    const modelPlainProperties = Object.entries<TSchema>(this.inputSchema.properties).filter(
+      ([key, schema]) => schemaSemantic(schema) === "model"
     );
     if (modelPlainProperties.length > 0) {
       for (const [key, propSchema] of modelPlainProperties) {
@@ -125,9 +124,8 @@ export class AiTask<
   // if all of them are stripped, then the task will fail in validateInput
   async narrowInput(input: Input): Promise<Input> {
     // TODO(str): this is very inefficient, we should cache the results, including intermediate results
-    const inputSchemaProperties = simplifySchema(this.inputSchema).properties;
-    const modelTaskProperties = Object.entries<TSchema>(inputSchemaProperties).filter(
-      ([key, value]) => value.semantic?.startsWith("model:")
+    const modelTaskProperties = Object.entries<TSchema>(this.inputSchema.properties).filter(
+      ([key, schema]) => schemaSemantic(schema)?.startsWith("model:")
     );
     if (modelTaskProperties.length > 0) {
       const taskModels = await getGlobalModelRepository().findModelsByTask(this.type);
diff --git a/packages/storage/src/kv/FsFolderKvRepository.ts b/packages/storage/src/kv/FsFolderKvRepository.ts
@@ -5,7 +5,7 @@
 //    *   Licensed under the Apache License, Version 2.0 (the "License");           *
 //    *******************************************************************************
 
-import { createServiceToken, simplifySchema, TypeBlob } from "@ellmers/util";
+import { createServiceToken, TypeBlob } from "@ellmers/util";
 import { TSchema, Type } from "@sinclair/typebox";
 import { mkdir, readFile, rmdir, unlink, writeFile } from "fs/promises";
 import path from "path";
@@ -52,7 +52,7 @@ export class FsFolderKvRepository<
     let content: string;
     if (value === null) {
       content = "";
-    } else if (simplifySchema(this.valueSchema).type === "object") {
+    } else if (this.valueSchema.type === "object") {
       content = JSON.stringify(value);
     } else if (typeof value === "object") {
       // Handle 'json' type schema from tests
@@ -74,7 +74,7 @@ export class FsFolderKvRepository<
    */
   public async get(key: Key): Promise<Value | undefined> {
     const localPath = path.join(this.folderPath, this.pathWriter(key));
-    const typeDef = simplifySchema(this.valueSchema);
+    const typeDef = this.valueSchema;
     try {
       const encoding = typeDef.contentEncoding === "blob" ? "binary" : "utf-8";
       const content = (await readFile(localPath, { encoding })).trim();
diff --git a/packages/storage/src/kv/KvViaTabularRepository.ts b/packages/storage/src/kv/KvViaTabularRepository.ts
@@ -5,7 +5,6 @@
 //    *   Licensed under the Apache License, Version 2.0 (the "License");           *
 //    *******************************************************************************
 
-import { simplifySchema } from "@ellmers/util";
 import { Static } from "@sinclair/typebox";
 import type { TabularRepository } from "../tabular/TabularRepository";
 import { DefaultKeyValueKey, DefaultKeyValueSchema } from "./IKvRepository";
@@ -39,7 +38,7 @@ export abstract class KvViaTabularRepository<
   public async put(key: Key, value: Value): Promise<void> {
     // Handle objects that need to be JSON-stringified, TODO(str): should put in the type
     const shouldStringify = !["number", "boolean", "string", "blob"].includes(
-      simplifySchema(this.valueSchema).type
+      this.valueSchema.type
     );
 
     if (shouldStringify) {
@@ -58,9 +57,7 @@ export abstract class KvViaTabularRepository<
   public async get(key: Key): Promise<Value | undefined> {
     const result = await this.tabularRepository.get({ key });
     if (result) {
-      const shouldParse = !["number", "boolean", "string", "blob"].includes(
-        simplifySchema(this.valueSchema).type
-      );
+      const shouldParse = !["number", "boolean", "string", "blob"].includes(this.valueSchema.type);
 
       if (shouldParse) {
         try {
@@ -96,9 +93,7 @@ export abstract class KvViaTabularRepository<
           ({
             key: value.key,
             value: (() => {
-              const shouldParse = !["number", "boolean", "string"].includes(
-                simplifySchema(this.valueSchema).type
-              );
+              const shouldParse = !["number", "boolean", "string"].includes(this.valueSchema.type);
 
               if (shouldParse && typeof value.value === "string") {
                 try {
diff --git a/packages/storage/src/tabular/BaseSqlTabularRepository.ts b/packages/storage/src/tabular/BaseSqlTabularRepository.ts
@@ -5,7 +5,6 @@
 //    *   Licensed under the Apache License, Version 2.0 (the "License");           *
 //    *******************************************************************************
 
-import { simplifySchema } from "@ellmers/util";
 import { Static, TObject, TSchema } from "@sinclair/typebox";
 import { ExtractPrimaryKey, ExtractValue, ValueOptionType } from "./ITabularRepository";
 import { TabularRepository } from "./TabularRepository";
diff --git a/packages/task-graph/src/task-graph/Dataflow.ts b/packages/task-graph/src/task-graph/Dataflow.ts
@@ -5,7 +5,8 @@
 //    *   Licensed under the Apache License, Version 2.0 (the "License");           *
 //    *******************************************************************************
 
-import { areSemanticallyCompatible, EventEmitter, simplifySchema } from "@ellmers/util";
+import { areSemanticallyCompatible, EventEmitter } from "@ellmers/util";
+import { Type } from "@sinclair/typebox";
 import { TaskError } from "../task/TaskError";
 import { DataflowJson } from "../task/TaskJSON";
 import { Provenance, TaskIdType, TaskOutput, TaskStatus } from "../task/TaskTypes";
@@ -16,7 +17,6 @@ import {
   DataflowEvents,
 } from "./DataflowEvents";
 import { TaskGraph } from "./TaskGraph";
-import { Type } from "@sinclair/typebox";
 
 export type DataflowIdType = `${string}[${string}] ==> ${string}[${string}]`;
 
@@ -117,11 +117,11 @@ export class Dataflow {
     const targetSchemaProperty =
       DATAFLOW_ALL_PORTS === dataflow.targetTaskPortId
         ? Type.Any()
-        : simplifySchema(targetSchema.properties[dataflow.targetTaskPortId]);
+        : targetSchema.properties[dataflow.targetTaskPortId];
     const sourceSchemaProperty =
       DATAFLOW_ALL_PORTS === dataflow.sourceTaskPortId
         ? Type.Any()
-        : simplifySchema(sourceSchema.properties[dataflow.sourceTaskPortId]);
+        : sourceSchema.properties[dataflow.sourceTaskPortId];
 
     const semanticallyCompatible = areSemanticallyCompatible(
       sourceSchemaProperty,
diff --git a/packages/task-graph/src/task-graph/Workflow.ts b/packages/task-graph/src/task-graph/Workflow.ts
@@ -5,7 +5,7 @@
 //    *   Licensed under the Apache License, Version 2.0 (the "License");           *
 //    *******************************************************************************
 
-import { EventEmitter, simplifySchema, type EventParameters } from "@ellmers/util";
+import { EventEmitter, type EventParameters } from "@ellmers/util";
 import { TObject, TSchema } from "@sinclair/typebox";
 import { TaskOutputRepository } from "../storage/TaskOutputRepository";
 import { GraphAsTask } from "../task/GraphAsTask";
@@ -18,7 +18,7 @@ import { getLastTask, parallel, pipe, PipeFunction, Taskish } from "./Conversion
 import { Dataflow, DATAFLOW_ALL_PORTS } from "./Dataflow";
 import { IWorkflow } from "./IWorkflow";
 import { TaskGraph } from "./TaskGraph";
-import { CompoundMergeStrategy, GraphResultMap } from "./TaskGraphRunner";
+import { CompoundMergeStrategy } from "./TaskGraphRunner";
 
 // Type definitions for the workflow
 export type CreateWorkflow<I extends TaskIO, O extends TaskIO, C extends TaskConfig> = (
@@ -137,8 +137,8 @@ export class Workflow<Input extends TaskIO = TaskIO, Output extends TaskIO = Tas
       if (parent && this.graph.getTargetDataflows(parent.config.id).length === 0) {
         // Find matches between parent outputs and task inputs based on valueType
         const matches = new Map<string, string>();
-        const sourceSchema = simplifySchema(parent.outputSchema) as TObject;
-        const targetSchema = simplifySchema(task.inputSchema) as TObject;
+        const sourceSchema = parent.outputSchema as TObject;
+        const targetSchema = task.inputSchema as TObject;
 
         const makeMatch = (
           comparator: (
diff --git a/packages/test/src/test/util/TypeBox.test.ts b/packages/test/src/test/util/TypeBox.test.ts
@@ -0,0 +1,116 @@
+import { OptionalKind, Type } from "@sinclair/typebox";
+import { describe, expect, test } from "bun:test";
+import { simplifySchema } from "@ellmers/util";
+
+describe("simplifySchema", () => {
+  test("should throw error for undefined schema", () => {
+    expect(() => simplifySchema(undefined as any)).toThrow("Schema is undefined");
+  });
+
+  test("should return Any schema as is", () => {
+    const schema = Type.Any();
+    expect(simplifySchema(schema)).toEqual(schema);
+  });
+
+  test("should simplify union of base type and array of same type", () => {
+    const schema = Type.Union([Type.String(), Type.Array(Type.String())]);
+    const result = simplifySchema(schema);
+    expect(result).toEqual(Type.String({ isArray: true }));
+  });
+
+  test("should preserve annotations when simplifying union", () => {
+    const schema1 = Type.Union([Type.String(), Type.Array(Type.String())], {
+      title: "MyString",
+      description: "A string or array of strings",
+    });
+    const result1 = simplifySchema(schema1);
+    expect(result1).toEqual(
+      Type.String({
+        title: "MyString",
+        description: "A string or array of strings",
+        isArray: true,
+      })
+    );
+    const schema2 = Type.Union([
+      Type.String({
+        title: "MyString",
+        description: "A string or array of strings",
+      }),
+      Type.Array(Type.String()),
+    ]);
+    const result2 = simplifySchema(schema2);
+    expect(result2).toEqual(
+      Type.String({
+        title: "MyString",
+        description: "A string or array of strings",
+        isArray: true,
+      })
+    );
+  });
+
+  test("should handle nullable types", () => {
+    const schema = Type.Union([Type.String(), Type.Null()]);
+    const result = simplifySchema(schema);
+    expect(result).toEqual(
+      Type.String({
+        isNullable: true,
+        default: null,
+      })
+    );
+  });
+
+  test("should recursively simplify object properties", () => {
+    const schema = Type.Object({
+      name: Type.Union([Type.String(), Type.Array(Type.String())]),
+      age: Type.Union([Type.Number(), Type.Null()]),
+    });
+    const result = simplifySchema(schema);
+    expect(result).toEqual(
+      Type.Object({
+        name: Type.String({ isArray: true }),
+        age: Type.Number({ isNullable: true, default: null }),
+      })
+    );
+  });
+
+  test("should recursively simplify array items", () => {
+    const schema = Type.Array(Type.Union([Type.String(), Type.Array(Type.String())]));
+    const result = simplifySchema(schema);
+    expect(result).toEqual(Type.Array(Type.String({ isArray: true })));
+  });
+
+  test("should preserve optional flag and default values", () => {
+    const schema = Type.Object({
+      name: Type.Optional(Type.String({ default: "John" })),
+    });
+    const result = simplifySchema(schema);
+    expect(result).toEqual(
+      Type.Object({
+        name: Type.String({
+          [OptionalKind]: "Optional",
+          optional: true,
+          isNullable: true,
+          default: "John",
+        }),
+      })
+    );
+  });
+
+  test("should handle complex nested structures", () => {
+    const schema = Type.Object({
+      user: Type.Object({
+        name: Type.Union([Type.String(), Type.Array(Type.String())]),
+        addresses: Type.Array(Type.Union([Type.String(), Type.Array(Type.String())])),
+      }),
+    });
+    const result = simplifySchema(schema);
+    expect(result).toEqual(
+      Type.Object({
+        user: Type.Object({
+          name: Type.String({ isArray: true }),
+          addresses: Type.Array(Type.String({ isArray: true })),
+        }),
+      })
+    );
+  });
+});
diff --git a/packages/util/src/typebox/TypeBox.ts b/packages/util/src/typebox/TypeBox.ts
