illuma-agents 1.0.38 → 1.0.40

package/src/graphs/Graph.ts

@@ -666,6 +666,116 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
     }
   }
 
+  /**
+   * Execute model invocation with structured output.
+   * Uses withStructuredOutput to force the model to return JSON conforming to the schema.
+   * Disables streaming and returns a validated JSON response.
+   */
+  private async attemptStructuredInvoke(
+    {
+      currentModel,
+      finalMessages,
+      schema,
+      structuredOutputConfig,
+    }: {
+      currentModel: t.ChatModelInstance;
+      finalMessages: BaseMessage[];
+      schema: Record<string, unknown>;
+      structuredOutputConfig: t.StructuredOutputConfig;
+    },
+    config?: RunnableConfig
+  ): Promise<{
+    structuredResponse: Record<string, unknown>;
+    rawMessage?: AIMessageChunk;
+  }> {
+    const model = this.overrideModel ?? currentModel;
+    if (!model) {
+      throw new Error('No model found');
+    }
+
+    // Check if model supports withStructuredOutput
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    if (typeof (model as any).withStructuredOutput !== 'function') {
+      throw new Error(
+        `The selected model does not support structured output. ` +
+        `Please use a model that supports JSON schema output (e.g., OpenAI GPT-4, Anthropic Claude, Google Gemini) ` +
+        `or disable structured output for this agent.`
+      );
+    }
+
+    const {
+      name = 'StructuredResponse',
+      includeRaw = false,
+      handleErrors = true,
+      maxRetries = 2,
+    } = structuredOutputConfig;
+
+    // Use withStructuredOutput to bind the schema
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    const structuredModel = (model as any).withStructuredOutput(schema, {
+      name,
+      includeRaw,
+      strict: structuredOutputConfig.strict !== false,
+    });
+
+    let lastError: Error | undefined;
+    let attempts = 0;
+
+    while (attempts <= maxRetries) {
+      try {
+        const result = await structuredModel.invoke(finalMessages, config);
+
+        // Handle includeRaw response format
+        if (includeRaw && result.raw && result.parsed) {
+          return {
+            structuredResponse: result.parsed as Record<string, unknown>,
+            rawMessage: result.raw as AIMessageChunk,
+          };
+        }
+
+        // Direct response
+        return {
+          structuredResponse: result as Record<string, unknown>,
+        };
+      } catch (error) {
+        lastError = error as Error;
+        attempts++;
+
+        // If error handling is disabled, throw immediately
+        if (handleErrors === false) {
+          throw error;
+        }
+
+        // If we've exhausted retries, throw
+        if (attempts > maxRetries) {
+          throw new Error(
+            `Structured output failed after ${maxRetries + 1} attempts: ${lastError.message}`
+          );
+        }
+
+        // Add error message to conversation for retry
+        const errorMessage =
+          typeof handleErrors === 'string'
+            ? handleErrors
+            : `The response did not match the expected schema. Error: ${lastError.message}. Please try again with a valid response.`;
+
+        console.warn(
+          `[Graph] Structured output attempt ${attempts} failed: ${lastError.message}. Retrying...`
+        );
+
+        // Add the error as a human message for context
+        finalMessages = [
+          ...finalMessages,
+          new HumanMessage({
+            content: `[VALIDATION ERROR]\n${errorMessage}`,
+          }),
+        ];
+      }
+    }
+
+    throw lastError ?? new Error('Structured output failed');
+  }
+
   cleanupSignalListener(currentModel?: t.ChatModel): void {
     if (!this.signal) {
       return;
@@ -732,14 +842,17 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
       }
 
       const toolsForBinding = agentContext.getToolsForBinding();
-      let model =
+      // Store reference to the raw model before piping for structured output support
+      const rawModel =
         this.overrideModel ??
         this.initializeModel({
           tools: toolsForBinding,
           provider: agentContext.provider,
           clientOptions: agentContext.clientOptions,
         });
-
+      
+      // Use piped model for regular invocations (includes system message handling)
+      let model = rawModel;
       if (agentContext.systemRunnable) {
         model = agentContext.systemRunnable.pipe(model as Runnable);
       }
@@ -948,6 +1061,57 @@ export class StandardGraph extends Graph<t.BaseGraphState, t.GraphNode> {
         config
       );
 
+      // Check if structured output mode is enabled
+      if (
+        agentContext.isStructuredOutputMode &&
+        agentContext.structuredOutput
+      ) {
+        const schema = agentContext.getStructuredOutputSchema();
+        if (!schema) {
+          throw new Error('Structured output schema is not configured');
+        }
+
+        try {
+          // Use structured output invocation (non-streaming)
+          // Use rawModel (not piped model) because withStructuredOutput is only on the chat model class
+          const { structuredResponse, rawMessage } =
+            await this.attemptStructuredInvoke(
+              {
+                currentModel: rawModel as t.ChatModelInstance,
+                finalMessages,
+                schema,
+                structuredOutputConfig: agentContext.structuredOutput,
+              },
+              config
+            );
+
+          // Emit structured output event
+          await safeDispatchCustomEvent(
+            GraphEvents.ON_STRUCTURED_OUTPUT,
+            {
+              structuredResponse,
+              schema,
+              raw: rawMessage,
+            },
+            config
+          );
+
+          agentContext.currentUsage = rawMessage
+            ? this.getUsageMetadata(rawMessage)
+            : undefined;
+          this.cleanupSignalListener();
+
+          // Return both the structured response and the raw message
+          return {
+            messages: rawMessage ? [rawMessage] : [],
+            structuredResponse,
+          };
+        } catch (structuredError) {
+          console.error('[Graph] Structured output failed:', structuredError);
+          throw structuredError;
+        }
+      }
+
       try {
         result = await this.attemptInvoke(
           {

package/src/index.ts

@@ -17,6 +17,9 @@ export * from './tools/ToolSearch';
 export * from './tools/handlers';
 export * from './tools/search';
 
+/* Schemas */
+export * from './schemas';
+
 /* Misc. */
 export * from './common';
 export * from './utils';

package/src/messages/cache.test.ts

@@ -920,6 +920,54 @@ describe('Immutability - addCacheControl does not mutate original messages', ()
 
     expect('cache_control' in originalFirstBlock).toBe(true);
   });
+
+  it('should remove lc_kwargs to prevent serialization mismatch for LangChain messages', () => {
+    type LangChainLikeMsg = TestMsg & {
+      lc_kwargs?: { content?: MessageContentComplex[] };
+    };
+
+    const messagesWithLcKwargs: LangChainLikeMsg[] = [
+      {
+        role: 'user',
+        content: [{ type: ContentTypes.TEXT, text: 'First user message' }],
+        lc_kwargs: {
+          content: [{ type: ContentTypes.TEXT, text: 'First user message' }],
+        },
+      },
+      {
+        role: 'assistant',
+        content: [{ type: ContentTypes.TEXT, text: 'Assistant response' }],
+        lc_kwargs: {
+          content: [{ type: ContentTypes.TEXT, text: 'Assistant response' }],
+        },
+      },
+      {
+        role: 'user',
+        content: [{ type: ContentTypes.TEXT, text: 'Second user message' }],
+        lc_kwargs: {
+          content: [{ type: ContentTypes.TEXT, text: 'Second user message' }],
+        },
+      },
+    ];
+
+    const result = addCacheControl(messagesWithLcKwargs as never);
+
+    const resultFirst = result[0] as LangChainLikeMsg;
+    const resultThird = result[2] as LangChainLikeMsg;
+
+    expect(resultFirst.lc_kwargs).toBeUndefined();
+    expect(resultThird.lc_kwargs).toBeUndefined();
+
+    const firstContent = resultFirst.content as MessageContentComplex[];
+    expect('cache_control' in firstContent[0]).toBe(true);
+
+    const originalFirst = messagesWithLcKwargs[0];
+    const originalContent = originalFirst.content as MessageContentComplex[];
+    const originalLcContent = originalFirst.lc_kwargs
+      ?.content as MessageContentComplex[];
+    expect('cache_control' in originalContent[0]).toBe(false);
+    expect('cache_control' in originalLcContent[0]).toBe(false);
+  });
 });
 
 describe('Immutability - addBedrockCacheControl does not mutate original messages', () => {
@@ -1049,7 +1097,7 @@ describe('Immutability - addBedrockCacheControl does not mutate original message
     expect('cache_control' in anthropicFirstContent[0]).toBe(true);
   });
 
-  it('should keep lc_kwargs.content in sync with content for LangChain messages', () => {
+  it('should remove lc_kwargs to prevent serialization mismatch for LangChain messages', () => {
     type LangChainLikeMsg = TestMsg & {
       lc_kwargs?: { content?: MessageContentComplex[] };
     };
@@ -1076,14 +1124,11 @@ describe('Immutability - addBedrockCacheControl does not mutate original message
     const resultFirst = bedrockResult[0] as LangChainLikeMsg;
     const resultSecond = bedrockResult[1] as LangChainLikeMsg;
 
-    expect(resultFirst.content).toEqual(resultFirst.lc_kwargs?.content);
-    expect(resultSecond.content).toEqual(resultSecond.lc_kwargs?.content);
+    expect(resultFirst.lc_kwargs).toBeUndefined();
+    expect(resultSecond.lc_kwargs).toBeUndefined();
 
     const firstContent = resultFirst.content as MessageContentComplex[];
-    const firstLcContent = resultFirst.lc_kwargs
-      ?.content as MessageContentComplex[];
     expect(firstContent.some((b) => 'cachePoint' in b)).toBe(true);
-    expect(firstLcContent.some((b) => 'cachePoint' in b)).toBe(true);
 
     const originalFirst = messagesWithLcKwargs[0];
     const originalContent = originalFirst.content as MessageContentComplex[];

package/src/messages/cache.ts

@@ -1,11 +1,4 @@
-import {
-  BaseMessage,
-  MessageContentComplex,
-  AIMessage,
-  HumanMessage,
-  SystemMessage,
-  ToolMessage,
-} from '@langchain/core/messages';
+import { BaseMessage, MessageContentComplex } from '@langchain/core/messages';
 import type { AnthropicMessage } from '@/types/messages';
 import type Anthropic from '@anthropic-ai/sdk';
 import { ContentTypes } from '@/common/enum';
@@ -41,84 +34,95 @@ function deepCloneContent<T extends string | MessageContentComplex[]>(
 }
 
 /**
- * Simple shallow clone with deep-cloned content.
- * Used for stripping cache control where we don't need proper LangChain instances.
+ * Clones a message with deep-cloned content, explicitly excluding LangChain
+ * serialization metadata to prevent coercion issues.
+ * Keeps lc_kwargs in sync with content to prevent LangChain serialization issues.
  */
-function _shallowCloneMessage<T extends MessageWithContent>(message: T): T {
-  const cloned = {
-    ...message,
-    content: deepCloneContent(message.content ?? ''),
-  } as T;
-  const lcKwargs = (cloned as Record<string, unknown>).lc_kwargs as
+function cloneMessage<T extends MessageWithContent>(
+  message: T,
+  content: string | MessageContentComplex[]
+): T {
+  const {
+    lc_kwargs: _lc_kwargs,
+    lc_serializable: _lc_serializable,
+    lc_namespace: _lc_namespace,
+    ...rest
+  } = message as T & {
+    lc_kwargs?: unknown;
+    lc_serializable?: unknown;
+    lc_namespace?: unknown;
+  };
+
+  const cloned = { ...rest, content } as T;
+
+  // Sync lc_kwargs.content with the new content to prevent LangChain coercion issues
+  const lcKwargs = (message as Record<string, unknown>).lc_kwargs as
     | Record<string, unknown>
     | undefined;
   if (lcKwargs != null) {
     (cloned as Record<string, unknown>).lc_kwargs = {
       ...lcKwargs,
-      content: cloned.content,
+      content: content,
+    };
+  }
+
+  // LangChain messages don't have a direct 'role' property - derive it from getType()
+  if (
+    'getType' in message &&
+    typeof message.getType === 'function' &&
+    !('role' in cloned)
+  ) {
+    const msgType = (message as unknown as BaseMessage).getType();
+    const roleMap: Record<string, string> = {
+      human: 'user',
+      ai: 'assistant',
+      system: 'system',
+      tool: 'tool',
     };
+    (cloned as Record<string, unknown>).role = roleMap[msgType] || msgType;
   }
+
   return cloned;
 }
 
 /**
- * Creates a new LangChain message instance with the given content.
- * Required when adding cache points to ensure proper serialization.
+ * Checks if a content block is a cache point
  */
-function _createNewMessage<T extends MessageWithContent>(
-  message: T,
-  content: MessageContentComplex[]
-): T {
-  if ('getType' in message && typeof message.getType === 'function') {
-    const baseMsg = message as unknown as BaseMessage;
-    const msgType = baseMsg.getType();
-
-    const baseFields = {
-      content,
-      name: baseMsg.name,
-      additional_kwargs: { ...baseMsg.additional_kwargs },
-      response_metadata: { ...baseMsg.response_metadata },
-      id: baseMsg.id,
-    };
+function isCachePoint(block: MessageContentComplex): boolean {
+  return 'cachePoint' in block && !('type' in block);
+}
 
-    switch (msgType) {
-    case 'human':
-      return new HumanMessage(baseFields) as unknown as T;
-    case 'ai': {
-      const aiMsg = baseMsg as AIMessage;
-      return new AIMessage({
-        ...baseFields,
-        tool_calls: aiMsg.tool_calls ? [...aiMsg.tool_calls] : [],
-        invalid_tool_calls: aiMsg.invalid_tool_calls
-          ? [...aiMsg.invalid_tool_calls]
-          : [],
-        usage_metadata: aiMsg.usage_metadata,
-      }) as unknown as T;
-    }
-    case 'system':
-      return new SystemMessage(baseFields) as unknown as T;
-    case 'tool': {
-      const toolMsg = baseMsg as ToolMessage;
-      return new ToolMessage({
-        ...baseFields,
-        tool_call_id: toolMsg.tool_call_id,
-        status: toolMsg.status,
-        artifact: toolMsg.artifact,
-      }) as unknown as T;
-    }
-    default:
-      break;
-    }
+/**
+ * Checks if a message's content needs cache control stripping.
+ * Returns true if content has cachePoint blocks or cache_control fields.
+ */
+function needsCacheStripping(content: MessageContentComplex[]): boolean {
+  for (let i = 0; i < content.length; i++) {
+    const block = content[i];
+    if (isCachePoint(block)) return true;
+    if ('cache_control' in block) return true;
   }
+  return false;
+}
 
-  const cloned = { ...message, content } as T;
-  const lcKwargs = (cloned as Record<string, unknown>).lc_kwargs as
-    | Record<string, unknown>
-    | undefined;
-  if (lcKwargs != null) {
-    (cloned as Record<string, unknown>).lc_kwargs = { ...lcKwargs, content };
+/**
+ * Checks if a message's content has Anthropic cache_control fields.
+ */
+function hasAnthropicCacheControl(content: MessageContentComplex[]): boolean {
+  for (let i = 0; i < content.length; i++) {
+    if ('cache_control' in content[i]) return true;
   }
-  return cloned;
+  return false;
+}
+
+/**
+ * Checks if a message's content has Bedrock cachePoint blocks.
+ */
+function hasBedrockCachePoint(content: MessageContentComplex[]): boolean {
+  for (let i = 0; i < content.length; i++) {
+    if (isCachePoint(content[i])) return true;
+  }
+  return false;
 }
 
 /**
@@ -126,8 +130,9 @@ function _createNewMessage<T extends MessageWithContent>(
  * Strips ALL existing cache control (both Anthropic and Bedrock formats) from all messages,
  * then adds fresh cache control to the last 2 user messages in a single backward pass.
  * This ensures we don't accumulate stale cache points across multiple turns.
+ * Returns a new array - only clones messages that require modification.
  * @param messages - The array of message objects.
- * @returns - The updated array of message objects with cache control added.
+ * @returns - A new array of message objects with cache control added.
  */
 export function addCacheControl<T extends AnthropicMessage | BaseMessage>(
   messages: T[]
@@ -136,68 +141,82 @@ export function addCacheControl<T extends AnthropicMessage | BaseMessage>(
     return messages;
   }
 
-  const updatedMessages = [...messages];
+  const updatedMessages: T[] = [...messages];
   let userMessagesModified = 0;
 
   for (let i = updatedMessages.length - 1; i >= 0; i--) {
-    const message = updatedMessages[i];
+    const originalMessage = updatedMessages[i];
+    const content = originalMessage.content;
     const isUserMessage =
-      ('getType' in message && message.getType() === 'human') ||
-      ('role' in message && message.role === 'user');
+      ('getType' in originalMessage && originalMessage.getType() === 'human') ||
+      ('role' in originalMessage && originalMessage.role === 'user');
+
+    const hasArrayContent = Array.isArray(content);
+    const needsStripping =
+      hasArrayContent &&
+      needsCacheStripping(content as MessageContentComplex[]);
+    const needsCacheAdd =
+      userMessagesModified < 2 &&
+      isUserMessage &&
+      (typeof content === 'string' || hasArrayContent);
+
+    if (!needsStripping && !needsCacheAdd) {
+      continue;
+    }
 
-    if (Array.isArray(message.content)) {
-      message.content = message.content.filter(
-        (block) => !isCachePoint(block as MessageContentComplex)
-      ) as typeof message.content;
+    let workingContent: MessageContentComplex[];
 
-      for (let j = 0; j < message.content.length; j++) {
-        const block = message.content[j] as Record<string, unknown>;
+    if (hasArrayContent) {
+      workingContent = deepCloneContent(
+        content as MessageContentComplex[]
+      ).filter((block) => !isCachePoint(block as MessageContentComplex));
+
+      for (let j = 0; j < workingContent.length; j++) {
+        const block = workingContent[j] as Record<string, unknown>;
         if ('cache_control' in block) {
           delete block.cache_control;
         }
       }
+    } else if (typeof content === 'string') {
+      workingContent = [
+        { type: 'text', text: content },
+      ] as MessageContentComplex[];
+    } else {
+      workingContent = [];
     }
 
     if (userMessagesModified >= 2 || !isUserMessage) {
+      updatedMessages[i] = cloneMessage(
+        originalMessage as MessageWithContent,
+        workingContent
+      ) as T;
       continue;
     }
 
-    if (typeof message.content === 'string') {
-      message.content = [
-        {
-          type: 'text',
-          text: message.content,
-          cache_control: { type: 'ephemeral' },
-        },
-      ];
-      userMessagesModified++;
-    } else if (Array.isArray(message.content)) {
-      for (let j = message.content.length - 1; j >= 0; j--) {
-        const contentPart = message.content[j];
-        if ('type' in contentPart && contentPart.type === 'text') {
-          (contentPart as Anthropic.TextBlockParam).cache_control = {
-            type: 'ephemeral',
-          };
-          userMessagesModified++;
-          break;
-        }
+    for (let j = workingContent.length - 1; j >= 0; j--) {
+      const contentPart = workingContent[j];
+      if ('type' in contentPart && contentPart.type === 'text') {
+        (contentPart as Anthropic.TextBlockParam).cache_control = {
+          type: 'ephemeral',
+        };
+        userMessagesModified++;
+        break;
       }
     }
+
+    updatedMessages[i] = cloneMessage(
+      originalMessage as MessageWithContent,
+      workingContent
+    ) as T;
   }
 
   return updatedMessages;
 }
 
-/**
- * Checks if a content block is a cache point
- */
-function isCachePoint(block: MessageContentComplex): boolean {
-  return 'cachePoint' in block && !('type' in block);
-}
-
 /**
  * Removes all Anthropic cache_control fields from messages
  * Used when switching from Anthropic to Bedrock provider
+ * Returns a new array - only clones messages that require modification.
  */
 export function stripAnthropicCacheControl<T extends MessageWithContent>(
   messages: T[]
@@ -206,20 +225,24 @@ export function stripAnthropicCacheControl<T extends MessageWithContent>(
     return messages;
   }
 
-  const updatedMessages = [...messages];
+  const updatedMessages: T[] = [...messages];
 
   for (let i = 0; i < updatedMessages.length; i++) {
-    const message = updatedMessages[i];
-    const content = message.content;
+    const originalMessage = updatedMessages[i];
+    const content = originalMessage.content;
 
-    if (Array.isArray(content)) {
-      for (let j = 0; j < content.length; j++) {
-        const block = content[j] as Record<string, unknown>;
-        if ('cache_control' in block) {
-          delete block.cache_control;
-        }
+    if (!Array.isArray(content) || !hasAnthropicCacheControl(content)) {
+      continue;
+    }
+
+    const clonedContent = deepCloneContent(content);
+    for (let j = 0; j < clonedContent.length; j++) {
+      const block = clonedContent[j] as Record<string, unknown>;
+      if ('cache_control' in block) {
+        delete block.cache_control;
       }
     }
+    updatedMessages[i] = cloneMessage(originalMessage, clonedContent) as T;
   }
 
   return updatedMessages;
@@ -228,6 +251,7 @@ export function stripAnthropicCacheControl<T extends MessageWithContent>(
 /**
  * Removes all Bedrock cachePoint blocks from messages
  * Used when switching from Bedrock to Anthropic provider
+ * Returns a new array - only clones messages that require modification.
  */
 export function stripBedrockCacheControl<T extends MessageWithContent>(
   messages: T[]
@@ -236,17 +260,20 @@ export function stripBedrockCacheControl<T extends MessageWithContent>(
     return messages;
   }
 
-  const updatedMessages = [...messages];
+  const updatedMessages: T[] = [...messages];
 
   for (let i = 0; i < updatedMessages.length; i++) {
-    const message = updatedMessages[i];
-    const content = message.content;
+    const originalMessage = updatedMessages[i];
+    const content = originalMessage.content;
 
-    if (Array.isArray(content)) {
-      message.content = content.filter(
-        (block) => !isCachePoint(block as MessageContentComplex)
-      ) as typeof content;
+    if (!Array.isArray(content) || !hasBedrockCachePoint(content)) {
+      continue;
     }
+
+    const clonedContent = deepCloneContent(content).filter(
+      (block) => !isCachePoint(block as MessageContentComplex)
+    );
+    updatedMessages[i] = cloneMessage(originalMessage, clonedContent) as T;
   }
 
   return updatedMessages;

package/src/schemas/index.ts

@@ -0,0 +1,2 @@
+// src/schemas/index.ts
+export * from './validate';