feat: Enhance complex object and JSON extraction, add validation tests, and improve error messages with LLM output.

Author: dosco

src/ax/dsp/complex_object_validation.test.ts

@@ -0,0 +1,226 @@
+import { describe, expect, it } from 'vitest';
+import { AxGen } from './generate.js';
+import { AxSignature, f } from './sig.js';
+
+describe('AxGen setExamples with complex signatures', () => {
+  it('should validate examples with complex object signatures', () => {
+    const sig = new AxSignature({
+      inputs: [
+        {
+          name: 'queryInput',
+          type: { name: 'string' },
+        },
+      ],
+      outputs: [
+        {
+          name: 'userProfile',
+          type: {
+            name: 'object',
+            fields: {
+              nestedString: { type: 'string' },
+              nestedNumber: { type: 'number' },
+              nestedObject: {
+                type: 'object',
+                fields: {
+                  deepString: { type: 'string' },
+                },
+              },
+            },
+          },
+        },
+      ],
+    });
+
+    const gen = new AxGen(sig);
+
+    const examples = [
+      {
+        queryInput: 'test input',
+        userProfile: {
+          nestedString: 'hello',
+          nestedNumber: 123,
+          nestedObject: {
+            deepString: 'deep',
+          },
+        },
+      },
+    ];
+
+    // This should not throw
+    expect(() => gen.setExamples(examples)).not.toThrow();
+  });
+
+  it('should validate examples with array of complex objects', () => {
+    const sig = new AxSignature({
+      inputs: [
+        {
+          name: 'queryInput',
+          type: { name: 'string' },
+        },
+      ],
+      outputs: [
+        {
+          name: 'itemList',
+          type: {
+            name: 'object',
+            isArray: true,
+            fields: {
+              id: { type: 'number' },
+              name: { type: 'string' },
+            },
+          },
+        },
+      ],
+    });
+
+    const gen = new AxGen(sig);
+
+    const examples = [
+      {
+        queryInput: 'list input',
+        itemList: [
+          { id: 1, name: 'Item 1' },
+          { id: 2, name: 'Item 2' },
+        ],
+      },
+    ];
+
+    // This should not throw
+    expect(() => gen.setExamples(examples)).not.toThrow();
+  });
+
+  it('should validate examples with fluent API signature', () => {
+    const sig = f()
+      .input('searchQuery', f.string())
+      .output(
+        'searchResult',
+        f.object({
+          title: f.string(),
+          meta: f.object({
+            score: f.number(),
+            tags: f.string().array(),
+          }),
+        })
+      )
+      .build();
+
+    const gen = new AxGen(sig);
+
+    const examples = [
+      {
+        searchQuery: 'search term',
+        searchResult: {
+          title: 'Result Title',
+          meta: {
+            score: 0.95,
+            tags: ['tag1', 'tag2'],
+          },
+        },
+      },
+    ];
+
+    // This should not throw
+    expect(() => gen.setExamples(examples)).not.toThrow();
+  });
+
+  it('should validate examples with user provided complex scenario', () => {
+    const signature = f()
+      .input(
+        'freeTimeRanges',
+        f
+          .json(
+            'List of available free time ranges with human-readable datetime strings, ISO strings, and timezone. You can pick any part of a time range to select a time slot of the proper duration'
+          )
+          .array()
+      )
+      .input(
+        'durationMinutes',
+        f.number('Required meeting duration in minutes')
+      )
+      .input('subject', f.string('Meeting subject').optional())
+      .input(
+        'previousRejectionContext',
+        f
+          .string(
+            "Natural language feedback from participants about what times don't work, WITH specific times"
+          )
+          .array()
+          .optional()
+      )
+      .input(
+        'participantCommunications',
+        f
+          .json(
+            'List of participant communications with participantId, from (participant name, email), content (message body), and receivedAt'
+          )
+          .array()
+          .optional()
+      )
+      .input(
+        'userPreferences',
+        f
+          .string(
+            "User's time preferences or scheduling constraints (e.g., 'tomorrow morning', 'Thursday afternoon', 'early next week', or general preferences like 'mornings preferred'). Use this to understand their timing preferences."
+          )
+          .optional()
+      )
+      .input(
+        'userPreferenceTimeRange',
+        f
+          .json(
+            "Structured datetime range representing the user's explicit time preference, parsed from natural language using englishToDatetimeRange. Contains precise start/end in the user's timezone."
+          )
+          .optional()
+      )
+      .output(
+        'selectedSlots',
+        f
+          .object({
+            startTimeISO: f.string(
+              "Start time as a complete ISO-8601 Instant string (e.g. '2025-05-20T16:00:00Z'). MUST include date, time, and timezone offset (Z or +HH:MM)."
+            ),
+            endTimeISO: f.string(
+              "End time as a complete ISO-8601 Instant string (e.g. '2025-05-20T17:00:00Z'). MUST include date, time, and timezone offset (Z or +HH:MM)."
+            ),
+            participantIds: f
+              .string(
+                'Participant ID that accepted this time slot (if matching)'
+              )
+              .array()
+              .optional(),
+          })
+          .array()
+      )
+      .build();
+
+    const gen = new AxGen(signature);
+
+    const examples = [
+      {
+        freeTimeRanges: [
+          {
+            startTime: 'Monday, June 2, 2025 at 9:00 AM PDT',
+            endTime: 'Monday, June 2, 2025 at 12:00 PM PDT',
+            durationMinutes: 180,
+            startTimeISO: '2025-06-02T16:00:00Z',
+            endTimeISO: '2025-06-02T19:00:00Z',
+          },
+        ],
+        durationMinutes: 60,
+        subject: 'Team Sync',
+        selectedSlots: [
+          {
+            startTimeISO: '2025-06-02T16:00:00Z',
+            endTimeISO: '2025-06-02T17:00:00Z',
+          },
+          {
+            startTimeISO: '2025-06-02T17:00:00Z',
+            endTimeISO: '2025-06-02T18:00:00Z',
+          },
+        ],
+      },
+    ];
+
+    expect(() => gen.setExamples(examples)).not.toThrow();
+  });
+});

src/ax/dsp/errors.ts

@@ -33,6 +33,8 @@ const toFieldType = (type: Readonly<AxField['type']>) => {
         return 'classification class';
       case 'code':
         return 'code';
+      case 'object':
+        return 'object';
       default:
         return 'string';
     }

src/ax/dsp/extract.ts

@@ -615,7 +615,7 @@ function validateAndParseFieldValue(
 
   let value: unknown | undefined;
 
-  if (field.type?.name === 'json') {
+  if (field.type?.name === 'json' && !field.type?.isArray) {
     try {
       const text = extractBlock(fieldValue);
       value = JSON.parse(text);
@@ -645,7 +645,24 @@ function validateAndParseFieldValue(
     if (Array.isArray(value)) {
       for (const [index, item] of value.entries()) {
         if (item !== undefined) {
-          const v = typeof item === 'string' ? item.trim() : item;
+          let v = typeof item === 'string' ? item.trim() : item;
+
+          // If we have a string item but expect an object/json, try to parse it as JSON
+          // This handles the case where the LLM outputs a markdown list of JSON strings
+          if (
+            typeof v === 'string' &&
+            (field.type?.name === 'object' ||
+              (field.type?.name as string) === 'json')
+          ) {
+            try {
+              // Try to extract a JSON block if present, otherwise parse directly
+              const jsonText = extractBlock(v);
+              v = JSON.parse(jsonText);
+            } catch {
+              // Ignore parsing errors here, let convertValueToType or validation handle it
+            }
+          }
+
           value[index] = convertValueToType(field, v, true);
         }
       }

src/ax/dsp/generate.ts

@@ -822,7 +822,7 @@ export class AxGen<IN = any, OUT extends AxGenOut = any>
             (err ?? lastError)?.message ??
             (err ?? lastError)?.toString() ??
             'unknown error'
-          }`
+          }\n\nLLM Output:\n${states.map((s) => s.content).join('\n---\n')}`
         ),
         ai,
         this.signature