Skip to content

Custom storage

The Multi-Agent Orchestrator System provides flexibility in how conversation data is stored through its abstract ChatStorage class. This guide will walk you through the process of creating a custom storage solution by extending this class.

Understanding the ChatStorage Abstract Class

The ChatStorage class defines the interface for all storage solutions in the system. It includes three main methods and two helper methods:

import { ConversationMessage } from "../types";
export abstract class ChatStorage {
protected isConsecutiveMessage(conversation: ConversationMessage[], newMessage: ConversationMessage): boolean {
if (conversation.length === 0) return false;
const lastMessage = conversation[conversation.length - 1];
return lastMessage.role === newMessage.role;
}
protected trimConversation(conversation: ConversationMessage[], maxHistorySize?: number): ConversationMessage[] {
if (maxHistorySize === undefined) return conversation;
// Ensure maxHistorySize is even to maintain complete binoms
const adjustedMaxHistorySize = maxHistorySize % 2 === 0 ? maxHistorySize : maxHistorySize - 1;
return conversation.slice(-adjustedMaxHistorySize);
}
abstract saveChatMessage(
userId: string,
sessionId: string,
agentId: string,
newMessage: ConversationMessage,
maxHistorySize?: number
): Promise<ConversationMessage[]>;
abstract fetchChat(
userId: string,
sessionId: string,
agentId: string,
maxHistorySize?: number
): Promise<ConversationMessage[]>;
abstract fetchAllChats(
userId: string,
sessionId: string
): Promise<ConversationMessage[]>;
}

The ChatStorage class now includes two helper methods:

  1. isConsecutiveMessage (TypeScript) / is_consecutive_message (Python): Checks if a new message is consecutive to the last message in the conversation.
  2. trimConversation (TypeScript) / trim_conversation (Python): Trims the conversation history to the specified maximum size, ensuring an even number of messages.

The three main abstract methods are:

  1. saveChatMessage (TypeScript) / save_chat_message (Python): Saves a new message to the storage.
  2. fetchChat (TypeScript) / fetch_chat (Python): Retrieves messages for a specific conversation.
  3. fetchAllChats (TypeScript) / fetch_all_chats (Python): Retrieves all messages for a user’s session.

Creating a Custom Storage Solution

To create a custom storage solution, follow these steps:

  1. Create a new class that extends ChatStorage.
  2. Implement all the abstract methods.
  3. Utilize the helper methods isConsecutiveMessage and trimConversation in your implementation.
  4. Add any additional methods or properties specific to your storage solution.

Here’s an example of a simple custom storage solution using an in-memory dictionary:

import { ChatStorage, ConversationMessage } from 'multi-agent-orchestrator';
class SimpleInMemoryStorage extends ChatStorage {
private storage: { [key: string]: ConversationMessage[] } = {};
async saveChatMessage(
userId: string,
sessionId: string,
agentId: string,
newMessage: ConversationMessage,
maxHistorySize?: number
): Promise<ConversationMessage[]> {
const key = `${userId}:${sessionId}:${agentId}`;
if (!this.storage[key]) {
this.storage[key] = [];
}
if (!this.isConsecutiveMessage(this.storage[key], newMessage)) {
this.storage[key].push(newMessage);
}
this.storage[key] = this.trimConversation(this.storage[key], maxHistorySize);
return this.storage[key];
}
async fetchChat(
userId: string,
sessionId: string,
agentId: string,
maxHistorySize?: number
): Promise<ConversationMessage[]> {
const key = `${userId}:${sessionId}:${agentId}`;
const conversation = this.storage[key] || [];
return this.trimConversation(conversation, maxHistorySize);
}
async fetchAllChats(
userId: string,
sessionId: string
): Promise<ConversationMessage[]> {
const allMessages: ConversationMessage[] = [];
for (const key in this.storage) {
if (key.startsWith(`${userId}:${sessionId}`)) {
allMessages.push(...this.storage[key]);
}
}
return allMessages;
}
}

Using Your Custom Storage

To use your custom storage with the Multi-Agent Orchestrator:

const customStorage = new SimpleInMemoryStorage();
const orchestrator = new MultiAgentOrchestrator({
storage: customStorage
});

By extending the ChatStorage class, you can create custom storage solutions tailored to your specific needs, whether it’s integrating with a particular database system, implementing caching mechanisms, or adapting to unique architectural requirements.

Remember to consider factors such as scalability, persistence, and error handling when implementing your custom storage solution for production use. The helper methods isConsecutiveMessage and trimConversation can be particularly useful for managing conversation history effectively.