illuma-agents 1.0.2

package/src/splitStream.ts

@@ -0,0 +1,234 @@
+import { nanoid } from 'nanoid';
+import type * as t from '@/types';
+import { ContentTypes, GraphEvents, StepTypes } from '@/common';
+
+export const SEPARATORS = [
+  '. ',
+  '?',
+  '!',
+  '۔',
+  '- ',
+  '。',
+  '‥',
+  ';',
+  '¡',
+  '¿',
+  '\n',
+  '```',
+];
+
+export class SplitStreamHandler {
+  private inCodeBlock = false;
+  private inThinkBlock = false;
+  private accumulate: boolean;
+  tokens: string[] = [];
+  lastToken = '';
+  reasoningTokens: string[] = [];
+  currentStepId?: string;
+  currentMessageId?: string;
+  currentType?: ContentTypes.TEXT | ContentTypes.THINK;
+  currentLength = 0;
+  reasoningKey: 'reasoning_content' | 'reasoning' = 'reasoning_content';
+  currentIndex = -1;
+  blockThreshold = 4500;
+  /** The run ID AKA the Message ID associated with the complete generation */
+  runId: string;
+  handlers?: t.SplitStreamHandlers;
+  constructor({
+    runId,
+    handlers,
+    accumulate,
+    reasoningKey,
+    blockThreshold,
+  }: {
+    runId: string;
+    accumulate?: boolean;
+    handlers: t.SplitStreamHandlers;
+    blockThreshold?: number;
+    reasoningKey?: 'reasoning_content' | 'reasoning';
+  }) {
+    this.runId = runId;
+    this.handlers = handlers;
+    if (reasoningKey) {
+      this.reasoningKey = reasoningKey;
+    }
+    if (blockThreshold != null) {
+      this.blockThreshold = blockThreshold;
+    }
+    this.accumulate = accumulate ?? false;
+  }
+  getMessageId = (): string | undefined => {
+    const messageId = this.currentMessageId;
+    if (messageId != null && messageId) {
+      return messageId;
+    }
+    return undefined;
+  };
+  createMessageStep = (
+    type?: ContentTypes.TEXT | ContentTypes.THINK
+  ): [string, string] => {
+    if (type != null && this.currentType !== type) {
+      this.currentType = type;
+    }
+    this.currentLength = 0;
+    this.currentIndex += 1;
+    this.currentStepId = `step_${nanoid()}`;
+    this.currentMessageId = `msg_${nanoid()}`;
+    return [this.currentStepId, this.currentMessageId];
+  };
+  dispatchRunStep = (stepId: string, stepDetails: t.StepDetails): void => {
+    const runStep: t.RunStep = {
+      id: stepId,
+      runId: this.runId,
+      type: stepDetails.type,
+      index: this.currentIndex,
+      stepDetails,
+      // usage: null,
+    };
+    this.handlers?.[GraphEvents.ON_RUN_STEP]?.({
+      event: GraphEvents.ON_RUN_STEP,
+      data: runStep,
+    });
+  };
+  dispatchMessageDelta = (stepId: string, delta: t.MessageDelta): void => {
+    const messageDelta: t.MessageDeltaEvent = {
+      id: stepId,
+      delta,
+    };
+    this.handlers?.[GraphEvents.ON_MESSAGE_DELTA]?.({
+      event: GraphEvents.ON_MESSAGE_DELTA,
+      data: messageDelta,
+    });
+  };
+  dispatchReasoningDelta = (stepId: string, delta: t.ReasoningDelta): void => {
+    const reasoningDelta: t.ReasoningDeltaEvent = {
+      id: stepId,
+      delta,
+    };
+    this.handlers?.[GraphEvents.ON_REASONING_DELTA]?.({
+      event: GraphEvents.ON_REASONING_DELTA,
+      data: reasoningDelta,
+    });
+  };
+  handleContent = (
+    content: string,
+    _type: ContentTypes.TEXT | ContentTypes.THINK
+  ): void => {
+    let type = _type;
+    if (this.inThinkBlock && type === ContentTypes.TEXT) {
+      type = ContentTypes.THINK;
+    }
+    if (this.accumulate) {
+      if (type === ContentTypes.THINK) {
+        this.reasoningTokens.push(content);
+      } else {
+        this.tokens.push(content);
+      }
+    }
+
+    if (this.currentType !== type) {
+      const [newStepId, newMessageId] = this.createMessageStep(type);
+      this.dispatchRunStep(newStepId, {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: {
+          message_id: newMessageId,
+        },
+      });
+    }
+
+    const stepId = this.currentStepId ?? '';
+    if (type === ContentTypes.THINK) {
+      this.dispatchReasoningDelta(stepId, {
+        content: [
+          {
+            type: ContentTypes.THINK,
+            think: content,
+          },
+        ],
+      });
+    } else {
+      this.dispatchMessageDelta(stepId, {
+        content: [
+          {
+            type: ContentTypes.TEXT,
+            text: content,
+          },
+        ],
+      });
+    }
+
+    this.currentLength += content.length;
+    if (this.inCodeBlock) {
+      return;
+    }
+
+    if (
+      this.currentLength > this.blockThreshold &&
+      SEPARATORS.some((sep) => content.includes(sep))
+    ) {
+      const [newStepId, newMessageId] = this.createMessageStep(type);
+      this.dispatchRunStep(newStepId, {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: {
+          message_id: newMessageId,
+        },
+      });
+    }
+  };
+  getDeltaContent(chunk?: t.CustomChunk): string {
+    return (chunk?.choices?.[0]?.delta as t.CustomChunkDelta)?.content ?? '';
+  }
+  getReasoningDelta(chunk?: t.CustomChunk): string {
+    return (
+      (chunk?.choices?.[0]?.delta as t.CustomChunkDelta)?.[this.reasoningKey] ??
+      ''
+    );
+  }
+  handle(chunk?: t.CustomChunk): void {
+    if (!chunk) {
+      return;
+    }
+
+    const content = this.getDeltaContent(chunk);
+    const reasoning_content = this.getReasoningDelta(chunk);
+    if (!content.length && !reasoning_content.length) {
+      return;
+    }
+
+    if (content.includes('```')) {
+      this.inCodeBlock = !this.inCodeBlock;
+    }
+
+    if (content.includes('<think>') && !this.inCodeBlock) {
+      this.inThinkBlock = true;
+    } else if (this.lastToken.includes('</think>') && !this.inCodeBlock) {
+      this.inThinkBlock = false;
+    }
+
+    this.lastToken = content;
+
+    const message_id = this.getMessageId() ?? '';
+
+    if (!message_id) {
+      const initialContentType = this.inThinkBlock
+        ? ContentTypes.THINK
+        : ContentTypes.TEXT;
+      const initialType = reasoning_content
+        ? ContentTypes.THINK
+        : initialContentType;
+      const [stepId, message_id] = this.createMessageStep(initialType);
+      this.dispatchRunStep(stepId, {
+        type: StepTypes.MESSAGE_CREATION,
+        message_creation: {
+          message_id,
+        },
+      });
+    }
+
+    if (reasoning_content) {
+      this.handleContent(reasoning_content, ContentTypes.THINK);
+    } else {
+      this.handleContent(content, ContentTypes.TEXT);
+    }
+  }
+}

package/src/stream.test.ts

@@ -0,0 +1,94 @@
+import type { AIMessageChunk } from '@langchain/core/messages';
+import type { ChatOpenAIReasoningSummary } from '@langchain/openai';
+import { getChunkContent } from './stream';
+import { Providers } from '@/common';
+
+describe('getChunkContent', () => {
+  it('should handle reasoning content for OpenAI/Azure providers', () => {
+    const chunk: Partial<AIMessageChunk> = {
+      content: 'Regular content',
+      additional_kwargs: {
+        reasoning: {
+          summary: [{ text: 'Reasoning summary text' }],
+        } as Partial<ChatOpenAIReasoningSummary>,
+      },
+    };
+
+    const result = getChunkContent({
+      chunk,
+      provider: Providers.OPENAI,
+      reasoningKey: 'reasoning',
+    });
+
+    expect(result).toBe('Reasoning summary text');
+  });
+
+  it('should fallback to reasoningKey when no OpenAI reasoning summary', () => {
+    const chunk: Partial<AIMessageChunk> = {
+      content: 'Regular content',
+      additional_kwargs: {
+        reasoning_content: 'Reasoning from key',
+      },
+    };
+
+    const result = getChunkContent({
+      chunk,
+      reasoningKey: 'reasoning_content',
+    });
+
+    expect(result).toBe('Reasoning from key');
+  });
+
+  it('should fallback to chunk.content when reasoningKey value is null or undefined', () => {
+    const chunk: Partial<AIMessageChunk> = {
+      content: 'Fallback content',
+      additional_kwargs: {
+        reasoning_content: null,
+      },
+    };
+
+    const result = getChunkContent({
+      chunk,
+      reasoningKey: 'reasoning_content',
+    });
+
+    expect(result).toBe('Fallback content');
+  });
+
+  it('should fallback to chunk.content when reasoningKey value is empty string', () => {
+    const chunk: Partial<AIMessageChunk> = {
+      content: ' can',
+      additional_kwargs: {
+        reasoning_content: '',
+      },
+    };
+
+    const result = getChunkContent({
+      chunk,
+      reasoningKey: 'reasoning_content',
+    });
+
+    expect(result).toBe(' can');
+  });
+
+  it('should return undefined when no content is available', () => {
+    const chunk: Partial<AIMessageChunk> = {
+      additional_kwargs: {},
+    };
+
+    const result = getChunkContent({
+      chunk,
+      reasoningKey: 'reasoning',
+    });
+
+    expect(result).toBeUndefined();
+  });
+
+  it('should handle missing chunk gracefully', () => {
+    const result = getChunkContent({
+      reasoningKey: 'reasoning',
+    });
+
+    expect(result).toBeUndefined();
+  });
+});