@nomyx/assistant/src/middleware/MessageCollationMiddleware.ts

A powerful assistant library and cli for your AI projects. works with Vertex AI (Claude and Gemini)

import { AIProviderMiddleware } from '../types/provider';
import { ChatMessage } from '../types/chat';

export class MessageCollationMiddleware implements AIProviderMiddleware {
  async preProcess(messages: ChatMessage[]): Promise<ChatMessage[]> {
    const collatedMessages: ChatMessage[] = [];
    let currentRole: string | null = null;
    let currentContent: string[] = [];

    for (const message of messages) {
      if (message.role !== currentRole) {
        if (currentRole !== null) {
          collatedMessages.push({
            role: currentRole,
            content: currentContent.join('\n')
          });
        }
        currentRole = message.role;
        currentContent = [message.content];
      } else {
        currentContent.push(message.content);
      }
    }

    if (currentRole !== null) {
      collatedMessages.push({
        role: currentRole,
        content: currentContent.join('\n')
      });
    }

    // Ensure roles alternate between "user" and "assistant"
    const alternatingMessages: ChatMessage[] = [];
    let lastRole: string | null = null;

    for (const message of collatedMessages) {
      if (message.role === 'system') {
        // Prepend system messages to the user's first message
        if (alternatingMessages.length === 0) {
          alternatingMessages.push({
            role: 'system',
            content: `${message.content}\n\n`
          });
        } else {
          alternatingMessages[0].content = `${message.content}\n\n${alternatingMessages[0].content}`;
        }
      } else if (message.role !== lastRole) {
        alternatingMessages.push(message);
        lastRole = message.role;
      } else {
        // If the role is the same as the last one, insert a dummy message from the other role
        const dummyRole = message.role === 'user' ? 'assistant' : 'user';
        alternatingMessages.push({
          role: dummyRole,
          content: 'Continuing the conversation...'
        });
        alternatingMessages.push(message);
        lastRole = message.role;
      }
    }

    return alternatingMessages;
  }

  async postProcess(response: any): Promise<any> {
    // No post-processing needed for this middleware
    return response;
  }
}