@restnfeel/agentc-starter-kit/src/rag/chat/context-manager.ts

한국어 기업용 CMS 모듈 - Task Master AI와 함께 빠르게 웹사이트를 구현할 수 있는 재사용 가능한 컴포넌트 시스템

import { ChatMessage, ConversationContext } from "../types";

export interface ConversationConfig {
  maxMessages: number;
  maxTokens?: number;
  retainSystemMessages: boolean;
  contextTimeoutMs?: number;
}

export interface MessageSummary {
  totalMessages: number;
  userMessages: number;
  assistantMessages: number;
  systemMessages: number;
  totalTokens: number;
}

export class ConversationContextManager {
  private conversations: Map<string, ConversationContext> = new Map();
  private config: ConversationConfig;

  constructor(config: Partial<ConversationConfig> = {}) {
    this.config = {
      maxMessages: 20,
      maxTokens: 4000,
      retainSystemMessages: true,
      contextTimeoutMs: 30 * 60 * 1000, // 30 minutes
      ...config,
    };
  }

  createConversation(id: string, systemPrompt?: string): ConversationContext {
    const context: ConversationContext = {
      id,
      messages: [],
      createdAt: new Date(),
      updatedAt: new Date(),
      metadata: {},
    };

    if (systemPrompt) {
      context.messages.push({
        id: this.generateMessageId(),
        role: "system",
        content: systemPrompt,
        timestamp: new Date(),
      });
    }

    this.conversations.set(id, context);
    return context;
  }

  getConversation(id: string): ConversationContext | null {
    const conversation = this.conversations.get(id);
    if (!conversation) {
      return null;
    }

    // Check if conversation has expired
    if (this.config.contextTimeoutMs) {
      const now = Date.now();
      const lastUpdate = conversation.updatedAt.getTime();
      if (now - lastUpdate > this.config.contextTimeoutMs) {
        this.conversations.delete(id);
        return null;
      }
    }

    return conversation;
  }

  addMessage(
    conversationId: string,
    message: Omit<ChatMessage, "id" | "timestamp">
  ): ChatMessage {
    let conversation = this.getConversation(conversationId);

    if (!conversation) {
      // Create new conversation if it doesn't exist
      conversation = this.createConversation(conversationId);
    }

    const chatMessage: ChatMessage = {
      ...message,
      id: this.generateMessageId(),
      timestamp: new Date(),
    };

    conversation.messages.push(chatMessage);
    conversation.updatedAt = new Date();

    // Trim conversation if it exceeds limits
    this.trimConversation(conversation);

    return chatMessage;
  }

  updateMessage(
    conversationId: string,
    messageId: string,
    updates: Partial<ChatMessage>
  ): boolean {
    const conversation = this.getConversation(conversationId);
    if (!conversation) {
      return false;
    }

    const messageIndex = conversation.messages.findIndex(
      (msg) => msg.id === messageId
    );
    if (messageIndex === -1) {
      return false;
    }

    conversation.messages[messageIndex] = {
      ...conversation.messages[messageIndex],
      ...updates,
      updatedAt: new Date(),
    };

    conversation.updatedAt = new Date();
    return true;
  }

  getConversationHistory(
    conversationId: string,
    lastN?: number
  ): ChatMessage[] {
    const conversation = this.getConversation(conversationId);
    if (!conversation) {
      return [];
    }

    const messages = conversation.messages;
    if (lastN && lastN > 0) {
      return messages.slice(-lastN);
    }

    return messages;
  }

  getConversationSummary(conversationId: string): MessageSummary | null {
    const conversation = this.getConversation(conversationId);
    if (!conversation) {
      return null;
    }

    const summary: MessageSummary = {
      totalMessages: conversation.messages.length,
      userMessages: 0,
      assistantMessages: 0,
      systemMessages: 0,
      totalTokens: 0,
    };

    for (const message of conversation.messages) {
      switch (message.role) {
        case "user":
          summary.userMessages++;
          break;
        case "assistant":
          summary.assistantMessages++;
          break;
        case "system":
          summary.systemMessages++;
          break;
      }

      // Simple token estimation (4 chars ≈ 1 token)
      summary.totalTokens += Math.ceil(message.content.length / 4);
    }

    return summary;
  }

  clearConversation(conversationId: string): boolean {
    return this.conversations.delete(conversationId);
  }

  getAllConversations(): ConversationContext[] {
    return Array.from(this.conversations.values());
  }

  cleanupExpiredConversations(): number {
    if (!this.config.contextTimeoutMs) {
      return 0;
    }

    const now = Date.now();
    let cleaned = 0;

    for (const [id, conversation] of this.conversations) {
      const lastUpdate = conversation.updatedAt.getTime();
      if (now - lastUpdate > this.config.contextTimeoutMs) {
        this.conversations.delete(id);
        cleaned++;
      }
    }

    return cleaned;
  }

  private trimConversation(conversation: ConversationContext): void {
    const messages = conversation.messages;

    // Separate system messages from conversation messages
    const systemMessages = messages.filter((msg) => msg.role === "system");
    const conversationMessages = messages.filter(
      (msg) => msg.role !== "system"
    );

    // Trim by message count
    if (conversationMessages.length > this.config.maxMessages) {
      const keepCount = this.config.maxMessages;
      conversationMessages.splice(0, conversationMessages.length - keepCount);
    }

    // Trim by token count if specified
    if (this.config.maxTokens) {
      let totalTokens = 0;
      const trimmedMessages: ChatMessage[] = [];

      // Count tokens from the end (most recent messages)
      for (let i = conversationMessages.length - 1; i >= 0; i--) {
        const messageTokens = Math.ceil(
          conversationMessages[i].content.length / 4
        );

        if (totalTokens + messageTokens <= this.config.maxTokens) {
          trimmedMessages.unshift(conversationMessages[i]);
          totalTokens += messageTokens;
        } else {
          break;
        }
      }

      conversationMessages.splice(
        0,
        conversationMessages.length,
        ...trimmedMessages
      );
    }

    // Reconstruct conversation with system messages
    if (this.config.retainSystemMessages) {
      conversation.messages = [...systemMessages, ...conversationMessages];
    } else {
      conversation.messages = conversationMessages;
    }
  }

  private generateMessageId(): string {
    return `msg_${Date.now()}_${Math.random().toString(36).substr(2, 9)}`;
  }

  // Update conversation metadata
  updateConversationMetadata(
    conversationId: string,
    metadata: Record<string, any>
  ): boolean {
    const conversation = this.getConversation(conversationId);
    if (!conversation) {
      return false;
    }

    conversation.metadata = { ...conversation.metadata, ...metadata };
    conversation.updatedAt = new Date();
    return true;
  }

  // Get statistics
  getManagerStats(): {
    totalConversations: number;
    totalMessages: number;
    config: ConversationConfig;
  } {
    let totalMessages = 0;
    for (const conversation of this.conversations.values()) {
      totalMessages += conversation.messages.length;
    }

    return {
      totalConversations: this.conversations.size,
      totalMessages,
      config: this.config,
    };
  }
}