fix: streaming parser overrun

Author: dosco

src/ax/ai/base.ts

@@ -456,7 +456,7 @@ export class AxBaseAI<
     }
 
     if (options?.debug ?? this.debug) {
-      logChatRequest(req.chatPrompt)
+      logChatRequest(req.chatPrompt, options?.debugHideSystemPrompt)
     }
 
     const rt = options?.rateLimiter ?? this.rt

src/ax/ai/debug.ts

@@ -6,10 +6,14 @@ const colorLog = new ColorLog()
 
 const formatChatMessage = (
   msg: AxChatRequest['chatPrompt'][number],
-  hideContent?: boolean
+  hideContent?: boolean,
+  hideSystemPrompt?: boolean
 ) => {
   switch (msg.role) {
     case 'system':
+      if (hideSystemPrompt) {
+        return ''
+      }
       return `\n${colorLog.blueBright('System:')}\n${colorLog.whiteBright(msg.content)}`
     case 'function':
       return `\n${colorLog.blueBright('Function Result:')}\n${colorLog.whiteBright(msg.result)}`
@@ -48,16 +52,20 @@ const formatChatMessage = (
 }
 
 export const logChatRequestMessage = (
-  msg: AxChatRequest['chatPrompt'][number]
+  msg: AxChatRequest['chatPrompt'][number],
+  hideSystemPrompt?: boolean
 ) => {
-  process.stdout.write(`${formatChatMessage(msg)}\n`)
+  process.stdout.write(`${formatChatMessage(msg, hideSystemPrompt)}\n`)
   process.stdout.write(colorLog.blueBright('\nAssistant:\n'))
 }
 
 export const logChatRequest = (
-  chatPrompt: Readonly<AxChatRequest['chatPrompt']>
+  chatPrompt: Readonly<AxChatRequest['chatPrompt']>,
+  hideSystemPrompt?: boolean
 ) => {
-  const items = chatPrompt?.map((msg) => formatChatMessage(msg))
+  const items = chatPrompt?.map((msg) =>
+    formatChatMessage(msg, hideSystemPrompt)
+  )
 
   if (items) {
     process.stdout.write(items.join('\n'))

src/ax/ai/types.ts

@@ -240,6 +240,7 @@ export type AxAIServiceActionOptions<
   traceId?: string
   rateLimiter?: AxRateLimiterFunction
   debug?: boolean
+  debugHideSystemPrompt?: boolean
 }
 
 export interface AxAIService<TModel = unknown, TEmbedModel = unknown> {

src/ax/dsp/datetime.test.ts

@@ -11,34 +11,34 @@ const field: AxField = {
 describe('datetime parsing', () => {
   it('should parse datetime with timezone abbreviation', () => {
     const dt = parseLLMFriendlyDateTime(field, '2022-01-01 12:00 EST')
-    expect(dt.toUTCString()).toBe('Sat, 01 Jan 2022 17:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 01 Jan 2022 17:00:00 GMT')
   })
 
   it('should parse datetime with seconds and timezone abbreviation', () => {
     const dt = parseLLMFriendlyDateTime(field, '2022-01-01 12:00:10 EST')
-    expect(dt.toUTCString()).toBe('Sat, 01 Jan 2022 17:00:10 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 01 Jan 2022 17:00:10 GMT')
   })
 
   it('should parse datetime with full timezone', () => {
     const dt = parseLLMFriendlyDateTime(
       field,
       '2022-01-01 12:00 America/New_York'
     )
-    expect(dt.toUTCString()).toBe('Sat, 01 Jan 2022 17:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 01 Jan 2022 17:00:00 GMT')
   })
 
   it('should parse datetime with another full timezone', () => {
     const dt = parseLLMFriendlyDateTime(
       field,
       '2022-01-01 12:00 America/Los_Angeles'
     )
-    expect(dt.toUTCString()).toBe('Sat, 01 Jan 2022 20:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 01 Jan 2022 20:00:00 GMT')
   })
 
   it('should parse datetime across DST boundary', () => {
     const summerDt = parseLLMFriendlyDateTime(field, '2022-07-01 12:00 EST')
     const winterDt = parseLLMFriendlyDateTime(field, '2022-01-01 12:00 EST')
-    expect(summerDt.getUTCHours()).toBe(winterDt.getUTCHours())
+    expect(summerDt?.getUTCHours()).toBe(winterDt?.getUTCHours())
   })
 
   it('should throw error for invalid datetime value', () => {
@@ -55,22 +55,22 @@ describe('datetime parsing', () => {
 describe('date parsing', () => {
   it('should parse valid date', () => {
     const dt = parseLLMFriendlyDate(field, '2022-01-01')
-    expect(dt.toUTCString()).toBe('Sat, 01 Jan 2022 00:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 01 Jan 2022 00:00:00 GMT')
   })
 
   it('should parse date with leading zeros', () => {
     const dt = parseLLMFriendlyDate(field, '2022-02-05')
-    expect(dt.toUTCString()).toBe('Sat, 05 Feb 2022 00:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 05 Feb 2022 00:00:00 GMT')
   })
 
   it('should parse date at year boundary', () => {
     const dt = parseLLMFriendlyDate(field, '2022-12-31')
-    expect(dt.toUTCString()).toBe('Sat, 31 Dec 2022 00:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Sat, 31 Dec 2022 00:00:00 GMT')
   })
 
   it('should parse date in leap year', () => {
     const dt = parseLLMFriendlyDate(field, '2024-02-29')
-    expect(dt.toUTCString()).toBe('Thu, 29 Feb 2024 00:00:00 GMT')
+    expect(dt?.toUTCString()).toBe('Thu, 29 Feb 2024 00:00:00 GMT')
   })
 
   it('should throw error for invalid date value', () => {

src/ax/dsp/datetime.ts

@@ -5,11 +5,15 @@ import { ValidationError } from './validate.js'
 
 export function parseLLMFriendlyDate(
   field: Readonly<AxField>,
-  dateStr: string
+  dateStr: string,
+  required: boolean = false
 ) {
   try {
     return _parseLLMFriendlyDate(dateStr)
   } catch (err) {
+    if (field.isOptional && !required) {
+      return
+    }
     const message = (err as Error).message
     throw new ValidationError({ fields: [field], message, value: dateStr })
   }
@@ -31,11 +35,15 @@ function _parseLLMFriendlyDate(dateStr: string) {
 
 export function parseLLMFriendlyDateTime(
   field: Readonly<AxField>,
-  dateStr: string
+  dateStr: string,
+  required: boolean = false
 ) {
   try {
     return _parseLLMFriendlyDateTime(dateStr)
   } catch (err) {
+    if (field.isOptional && !required) {
+      return
+    }
     const message = (err as Error).message
     throw new ValidationError({ fields: [field], message, value: dateStr })
   }

src/ax/dsp/extract.ts

@@ -16,7 +16,7 @@ export const extractValues = (
 }
 
 export interface extractionState {
-  prevField?: { field: AxField; s: number; e: number }
+  prevFields?: { field: AxField; s: number; e: number }[]
   currField?: AxField
   currFieldIndex?: number
   extractedFields: AxField[]
@@ -64,8 +64,9 @@ export const streamingExtractValues = (
       continue
     }
 
-    const prefix = field.title + ':'
-    let e = matchesContent(content, prefix, xstate.s + 1)
+    const isFirst = xstate.extractedFields.length === 0
+    const prefix = (isFirst ? '' : '\n') + field.title + ':'
+    let e = matchesContent(content, prefix, xstate.s === 0 ? 0 : xstate.s + 1)
 
     switch (e) {
       case -1:
@@ -95,7 +96,11 @@ export const streamingExtractValues = (
       if (parsedValue !== undefined) {
         values[xstate.currField.name] = parsedValue
       }
-      xstate.prevField = { field: xstate.currField, s: xstate.s, e }
+      if (xstate.prevFields) {
+        xstate.prevFields?.push({ field: xstate.currField, s: xstate.s, e })
+      } else {
+        xstate.prevFields = [{ field: xstate.currField, s: xstate.s, e }]
+      }
     }
 
     checkMissingRequiredFields(xstate, values, index)
@@ -136,7 +141,11 @@ export const streamingExtractFinalValue = (
   checkMissingRequiredFields(xstate, values, sigFields.length)
 }
 
-const convertValueToType = (field: Readonly<AxField>, val: string) => {
+const convertValueToType = (
+  field: Readonly<AxField>,
+  val: string,
+  required: boolean = false
+) => {
   switch (field.type?.name) {
     case 'code':
       return extractBlock(val)
@@ -147,6 +156,9 @@ const convertValueToType = (field: Readonly<AxField>, val: string) => {
     case 'number': {
       const v = Number(val)
       if (Number.isNaN(v)) {
+        if (field.isOptional && !required) {
+          return
+        }
         throw new Error('Invalid number')
       }
       return v
@@ -162,18 +174,24 @@ const convertValueToType = (field: Readonly<AxField>, val: string) => {
       } else if (v === 'false') {
         return false
       } else {
+        if (field.isOptional && !required) {
+          return
+        }
         throw new Error('Invalid boolean')
       }
     }
     case 'date':
-      return parseLLMFriendlyDate(field, val)
+      return parseLLMFriendlyDate(field, val, required)
 
     case 'datetime':
-      return parseLLMFriendlyDateTime(field, val)
+      return parseLLMFriendlyDateTime(field, val, required)
 
     case 'class':
       const className = val
       if (field.type.classes && !field.type.classes.includes(className)) {
+        if (field.isOptional) {
+          return
+        }
         throw new Error(
           `Invalid class '${val}', expected one of the following: ${field.type.classes.join(', ')}`
         )
@@ -249,11 +267,11 @@ export function* streamValues<OUT>(
   // eslint-disable-next-line functional/prefer-immutable-types
   xstate: extractionState
 ) {
-  if (xstate.prevField && !xstate.prevField.field.isInternal) {
-    const { field, s, e } = xstate.prevField
+  for (const prevField of xstate.prevFields ?? []) {
+    const { field, s, e } = prevField
     yield* yieldDelta<OUT>(content, field, s, e, xstate)
-    xstate.prevField = undefined
   }
+  xstate.prevFields = undefined
 
   if (!xstate.currField || xstate.currField.isInternal) {
     return
@@ -354,7 +372,7 @@ function validateAndParseFieldValue(
       for (const [index, item] of value.entries()) {
         if (item !== undefined) {
           const v = typeof item === 'string' ? item.trim() : item
-          value[index] = convertValueToType(field, v)
+          value[index] = convertValueToType(field, v, true)
         }
       }
     } else {

src/ax/dsp/generate.ts

@@ -517,7 +517,12 @@ export class AxGen<
     const maxRetries = options.maxRetries ?? this.options?.maxRetries ?? 10
     const maxSteps = options.maxSteps ?? this.options?.maxSteps ?? 10
     const debug = options.debug ?? ai.getOptions().debug
-    const mem = options.mem ?? this.options?.mem ?? new AxMemory(10000, debug)
+    const memOptions = {
+      debug: options.debug,
+      debugHideSystemPrompt: options.debugHideSystemPrompt,
+    }
+    const mem =
+      options.mem ?? this.options?.mem ?? new AxMemory(10000, memOptions)
 
     let err: ValidationError | AxAssertionError | undefined
 

src/ax/dsp/program.ts

@@ -62,6 +62,7 @@ export type AxProgramForwardOptions = {
   stopFunction?: string
   fastFail?: boolean
   debug?: boolean
+  debugHideSystemPrompt?: boolean
 }
 
 export type AxProgramStreamingForwardOptions = Omit<

src/ax/dsp/sig.test.ts

@@ -96,13 +96,7 @@ describe('extract values with signatures', () => {
     extractValues(
       sig,
       v1,
-      `
-      Title: Coastal Ecosystem Restoration
-
-      Key Points: Coastal regions prone to natural disasters, Selection criteria based on vulnerability indices and population density, Climate risk assessments conducted for sea-level rise and extreme weather events, Targeted ecosystems include mangrove forests, coral reefs, wetlands
-
-      Description: The project focuses on coastal regions vulnerable to natural disasters like hurricanes and flooding. Selection criteria included vulnerability indices, population density, and proximity to critical infrastructure. Climate risk assessments identified risks related to sea-level rise, storm surges, and extreme weather events. Targeted ecosystems encompass mangrove forests, coral reefs, and wetlands that provide coastal protection, biodiversity support, and livelihood opportunities for local communities.
-      `
+      `Title: Coastal Ecosystem Restoration\nKey Points: Coastal regions prone to natural disasters, Selection criteria based on vulnerability indices and population density, Climate risk assessments conducted for sea-level rise and extreme weather events, Targeted ecosystems include mangrove forests, coral reefs, wetlands\nDescription: The project focuses on coastal regions vulnerable to natural disasters like hurricanes and flooding. Selection criteria included vulnerability indices, population density, and proximity to critical infrastructure. Climate risk assessments identified risks related to sea-level rise, storm surges, and extreme weather events. Targeted ecosystems encompass mangrove forests, coral reefs, and wetlands that provide coastal protection, biodiversity support, and livelihood opportunities for local communities.`
     )
 
     expect(v1).toEqual({

src/ax/dsp/util.ts

@@ -308,6 +308,7 @@ export function matchesContent(
 
   // First check if the complete prefix exists anywhere after startIndex
   const exactMatchIndex = content.indexOf(prefix, startIndex)
+
   if (exactMatchIndex !== -1) {
     return exactMatchIndex
   }