refactor: enhance AxAI message handling and signature validation

This commit refines the message creation logic in the AxAIOpenAIChatRequest, improving the handling of user and assistant roles, including better management of tool calls. The AxSignature class is updated to accept a configuration object for inputs and outputs, enhancing flexibility. Additionally, the AxGen class constructor is modified to ensure non-nullable signature parameters, and various debug logging improvements are introduced across the AxGen and AxProgram classes.

Author: dosco

src/ax/ai/openai/api.ts

@@ -456,16 +456,19 @@ const mapFinishReason = (
 function createMessages<TModel>(
   req: Readonly<AxInternalChatRequest<TModel>>
 ): AxAIOpenAIChatRequest<TModel>['messages'] {
+  type UserContent = Extract<
+    AxAIOpenAIChatRequest<TModel>['messages'][number],
+    { role: 'user' }
+  >['content']
+
   return req.chatPrompt.map((msg) => {
     switch (msg.role) {
       case 'system':
         return { role: 'system' as const, content: msg.content }
+
       case 'user':
-        if (Array.isArray(msg.content)) {
-          return {
-            role: 'user' as const,
-            name: msg.name,
-            content: msg.content.map((c) => {
+        const content: UserContent = Array.isArray(msg.content)
+          ? msg.content.map((c) => {
               switch (c.type) {
                 case 'text':
                   return { type: 'text' as const, text: c.text }
@@ -486,27 +489,48 @@ function createMessages<TModel>(
                 default:
                   throw new Error('Invalid content type')
               }
-            }),
-          }
+            })
+          : msg.content
+        return {
+          role: 'user' as const,
+          ...(msg.name ? { name: msg.name } : {}),
+          content,
         }
-        return { role: 'user' as const, content: msg.content, name: msg.name }
+
       case 'assistant':
+        const toolCalls = msg.functionCalls?.map((v) => ({
+          id: v.id,
+          type: 'function' as const,
+          function: {
+            name: v.function.name,
+            arguments:
+              typeof v.function.params === 'object'
+                ? JSON.stringify(v.function.params)
+                : v.function.params,
+          },
+        }))
+
+        if (toolCalls && toolCalls.length > 0) {
+          return {
+            role: 'assistant' as const,
+            ...(msg.content ? { content: msg.content } : {}),
+            name: msg.name,
+            tool_calls: toolCalls,
+          }
+        }
+
+        if (!msg.content) {
+          throw new Error(
+            'Assistant content is required when no tool calls are provided'
+          )
+        }
+
         return {
           role: 'assistant' as const,
-          content: msg.content as string,
-          name: msg.name,
-          tool_calls: msg.functionCalls?.map((v) => ({
-            id: v.id,
-            type: 'function' as const,
-            function: {
-              name: v.function.name,
-              arguments:
-                typeof v.function.params === 'object'
-                  ? JSON.stringify(v.function.params)
-                  : v.function.params,
-            },
-          })),
+          content: msg.content,
+          ...(msg.name ? { name: msg.name } : {}),
         }
+
       case 'function':
         return {
           role: 'tool' as const,

src/ax/ai/openai/chat_types.ts

@@ -124,7 +124,17 @@ export type AxAIOpenAIChatRequest<TModel> = {
               text: string
             }
         name?: string
-        tool_calls?: {
+      }
+    | {
+        role: 'assistant'
+        content?:
+          | string
+          | {
+              type: string
+              text: string
+            }
+        name?: string
+        tool_calls: {
           type: 'function'
           function: {
             name: string

src/ax/dsp/generate.ts

@@ -14,7 +14,6 @@ import type {
   AxChatResponse,
   AxChatResponseResult,
   AxFunction,
-  AxLoggerFunction,
 } from '../ai/types.js'
 import { mergeFunctionCalls } from '../ai/util.js'
 import { AxMemory } from '../mem/memory.js'
@@ -105,16 +104,14 @@ export class AxGen<
   private values: AxGenOutType = {}
   private excludeContentFromTrace: boolean = false
   private thoughtFieldName: string
-  private logger?: AxLoggerFunction
 
   constructor(
-    signature: Readonly<AxSignature | string>,
+    signature: NonNullable<ConstructorParameters<typeof AxSignature>[0]>,
     options?: Readonly<AxProgramForwardOptions>
   ) {
     super(signature, { description: options?.description })
 
     this.options = options
-    this.logger = options?.logger
     this.thoughtFieldName = options?.thoughtFieldName ?? 'thought'
     const promptTemplateOptions = {
       functions: options?.functions,
@@ -247,7 +244,7 @@ export class AxGen<
         traceId,
         rateLimiter,
         stream,
-        debug: false,
+        debug: false, // we do our own debug logging
         thinkingTokenBudget,
         showThoughts,
         traceContext,
@@ -275,7 +272,6 @@ export class AxGen<
   }>) {
     const { sessionId, traceId, functions: _functions } = options ?? {}
     const fastFail = options?.fastFail ?? this.options?.fastFail
-
     const model = options.model
 
     // biome-ignore lint/complexity/useFlatMap: you cannot use flatMap here
@@ -303,6 +299,8 @@ export class AxGen<
         fastFail,
         span,
       })
+
+      this.getLogger(ai, options)?.('', { tags: ['responseEnd'] })
     } else {
       yield await this.processResponse({
         ai,
@@ -343,7 +341,6 @@ export class AxGen<
     mem.addResult(
       {
         content: '',
-        name: 'initial',
         functionCalls: [],
       },
       sessionId
@@ -492,11 +489,6 @@ export class AxGen<
         xstate
       )
     }
-
-    if (ai.getOptions().debug) {
-      const logger = ai.getLogger()
-      logger('', { tags: ['responseEnd'] })
-    }
   }
 
   private async processResponse({
@@ -589,9 +581,11 @@ 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 debugHideSystemPrompt = options.debugHideSystemPrompt
-    const memOptions = { debug, debugHideSystemPrompt }
+    const memOptions = {
+      debug: this.isDebug(ai, options),
+      debugHideSystemPrompt,
+    }
 
     const mem =
       options.mem ?? this.options?.mem ?? new AxMemory(10000, memOptions)
@@ -664,11 +658,7 @@ export class AxGen<
             continue multiStepLoop
           }
 
-          if (debug) {
-            const logger = options.logger ?? this.logger ?? ai.getLogger()
-            logger('', { tags: ['responseEnd'] })
-          }
-
+          this.getLogger(ai, options)?.('', { tags: ['responseEnd'] })
           return
         } catch (e) {
           let errorFields: AxIField[] | undefined
@@ -872,6 +862,22 @@ export class AxGen<
     super.setExamples(examples, options)
     // No need to update prompt template - all fields can be missing in examples
   }
+
+  private isDebug(
+    ai: Readonly<AxAIService>,
+    options?: Readonly<AxProgramForwardOptions>
+  ) {
+    return (
+      options?.debug ?? this.options?.debug ?? ai.getOptions().debug ?? false
+    )
+  }
+
+  private getLogger(
+    ai: Readonly<AxAIService>,
+    options?: Readonly<AxProgramForwardOptions>
+  ) {
+    return options?.logger ?? this.options?.logger ?? ai.getLogger()
+  }
 }
 
 export type AxGenerateErrorDetails = {

src/ax/dsp/program.ts

@@ -149,7 +149,7 @@ export class AxProgramWithSignature<IN extends AxGenIn, OUT extends AxGenOut>
   private children: AxInstanceRegistry<Readonly<AxTunable & AxUsable>>
 
   constructor(
-    signature: Readonly<AxSignature | string>,
+    signature: NonNullable<ConstructorParameters<typeof AxSignature>[0]>,
     options?: Readonly<AxProgramWithSignatureOptions>
   ) {
     this.signature = new AxSignature(signature)

src/ax/dsp/sig.ts

@@ -26,7 +26,7 @@ export interface AxField {
       | 'datetime'
       | 'class'
       | 'code'
-    isArray: boolean
+    isArray?: boolean
     options?: string[]
   }
   isOptional?: boolean
@@ -46,6 +46,12 @@ class AxSignatureValidationError extends Error {
   }
 }
 
+export interface AxSignatureConfig {
+  description?: string
+  inputs: readonly AxField[]
+  outputs: readonly AxField[]
+}
+
 export class AxSignature {
   private description?: string
   private inputFields: AxIField[]
@@ -57,7 +63,7 @@ export class AxSignature {
   // Validation caching - stores hash when validation last passed
   private validatedAtHash?: string
 
-  constructor(signature?: Readonly<AxSignature | string>) {
+  constructor(signature?: Readonly<AxSignature | string | AxSignatureConfig>) {
     if (!signature) {
       this.inputFields = []
       this.outputFields = []
@@ -108,11 +114,47 @@ export class AxSignature {
       if (signature.validatedAtHash === this.sigHash) {
         this.validatedAtHash = this.sigHash
       }
+    } else if (typeof signature === 'object' && signature !== null) {
+      // Handle AxSignatureConfig object
+      if (!('inputs' in signature) || !('outputs' in signature)) {
+        throw new AxSignatureValidationError(
+          'Invalid signature object: missing inputs or outputs',
+          undefined,
+          'Signature object must have "inputs" and "outputs" arrays. Example: { inputs: [...], outputs: [...] }'
+        )
+      }
+
+      if (
+        !Array.isArray(signature.inputs) ||
+        !Array.isArray(signature.outputs)
+      ) {
+        throw new AxSignatureValidationError(
+          'Invalid signature object: inputs and outputs must be arrays',
+          undefined,
+          'Both "inputs" and "outputs" must be arrays of AxField objects'
+        )
+      }
+
+      try {
+        this.description = signature.description
+        this.inputFields = signature.inputs.map((v) => this.parseField(v))
+        this.outputFields = signature.outputs.map((v) => this.parseField(v))
+        ;[this.sigHash, this.sigString] = this.updateHash()
+      } catch (error) {
+        if (error instanceof AxSignatureValidationError) {
+          throw error
+        }
+        throw new AxSignatureValidationError(
+          `Failed to create signature from object: ${error instanceof Error ? error.message : 'Unknown error'}`,
+          undefined,
+          'Check that all fields in inputs and outputs arrays are valid AxField objects'
+        )
+      }
     } else {
       throw new AxSignatureValidationError(
         'Invalid signature argument type',
         undefined,
-        'Signature must be a string or another AxSignature instance'
+        'Signature must be a string, another AxSignature instance, or an object with inputs and outputs arrays'
       )
     }
   }
@@ -166,7 +208,7 @@ export class AxSignature {
     }
     this.description = desc
     this.invalidateValidationCache()
-    this.updateHash()
+    this.updateHashLight()
   }
 
   public addInputField = (field: Readonly<AxField>) => {

src/ax/prompts/agent.ts

@@ -184,7 +184,7 @@ export class AxAgent<IN extends AxGenIn, OUT extends AxGenOut = AxGenOut>
       name: string
       description: string
       definition?: string
-      signature: AxSignature | string
+      signature: NonNullable<ConstructorParameters<typeof AxSignature>[0]>
       agents?: AxAgentic[]
       functions?: AxInputFunctionType
     }>,

src/examples/fibonacci.ts

@@ -1,13 +1,16 @@
-import { AxAI, AxGen, AxJSInterpreter, AxSignature } from '@ax-llm/ax'
+import { AxAI, AxGen, AxJSInterpreter } from '@ax-llm/ax'
 
-const sig = new AxSignature(
-  `numberSeriesTask:string  -> fibonacciSeries:number[]`
+const gen = new AxGen<{ numberSeriesTask: string }>(
+  {
+    inputs: [{ name: 'numberSeriesTask', type: { name: 'string' } }],
+    outputs: [{ name: 'fibonacciSeries', type: { name: 'number' } }],
+  },
+  {
+    functions: [new AxJSInterpreter()],
+    debug: true,
+  }
 )
 
-const gen = new AxGen<{ numberSeriesTask: string }>(sig, {
-  functions: [new AxJSInterpreter()],
-})
-
 const ai = new AxAI({
   name: 'openai',
   apiKey: process.env.OPENAI_APIKEY as string,