s3thing.ts

import { S3Client, PutObjectCommand, GetObjectCommand, ListObjectsV2Command } from '@aws-sdk/client-s3';
import { createHash } from 'crypto';

// Types for our chat system
interface ChatMessage {
  id: string;
  timestamp: number;
  role: 'user' | 'assistant';
  content: string;
  metadata?: Record<string, any>;
}

interface ChatSession {
  chatId: string;
  userId: string;
  documentId?: string;
  workspaceId: string;
  messages: ChatMessage[];
  startTime: number;
  endTime?: number;
  metadata?: Record<string, any>;
}

interface S3StorageConfig {
  region: string;
  bucketName: string;
  accessKeyId?: string;
  secretAccessKey?: string;
}

class ChatHistoryStorage {
  private s3Client: S3Client;
  private bucketName: string;

  constructor(config: S3StorageConfig) {
    this.s3Client = new S3Client({
      region: config.region,
      credentials: config.accessKeyId && config.secretAccessKey 
        ? { accessKeyId: config.accessKeyId, secretAccessKey: config.secretAccessKey }
        : undefined
    });
    this.bucketName = config.bucketName;
  }

  /**
   * Save a chat session with multiple indexing strategies
   */
  async saveChatSession(session: ChatSession): Promise<void> {
    const sessionData = JSON.stringify(session);
    
    // Create primary and secondary indices
    const savePromises = [
      // Primary storage location - indexed by chatId
      this.saveObject(`chats/${session.chatId}/data.json`, sessionData),
      
      // Index by userId
      this.saveObject(`users/${session.userId}/chats/${session.chatId}/reference.json`, 
        JSON.stringify({ 
          type: 'reference',
          targetKey: `chats/${session.chatId}/data.json`,
          chatId: session.chatId,
          timestamp: Date.now()
        })
      ),
      
      // Index by workspaceId
      this.saveObject(`workspaces/${session.workspaceId}/chats/${session.chatId}/reference.json`,
        JSON.stringify({ 
          type: 'reference',
          targetKey: `chats/${session.chatId}/data.json`,
          chatId: session.chatId,
          timestamp: Date.now()
        })
      )
    ];
    
    // Add document index if document ID is provided
    if (session.documentId) {
      savePromises.push(
        this.saveObject(`documents/${session.documentId}/chats/${session.chatId}/reference.json`,
          JSON.stringify({ 
            type: 'reference',
            targetKey: `chats/${session.chatId}/data.json`,
            chatId: session.chatId,
            timestamp: Date.now()
          })
        )
      );
    }
    
    // Create compound indexes for more complex queries
    savePromises.push(
      this.saveObject(
        `compound/user-workspace/${session.userId}/${session.workspaceId}/chats/${session.chatId}/reference.json`,
        JSON.stringify({ 
          type: 'reference',
          targetKey: `chats/${session.chatId}/data.json`,
          chatId: session.chatId,
          timestamp: Date.now()
        })
      )
    );
    
    await Promise.all(savePromises);
  }

  /**
   * Update an existing chat session with new messages
   */
  async updateChatSession(chatId: string, newMessages: ChatMessage[]): Promise<void> {
    try {
      // Retrieve existing session
      const session = await this.getChatSessionById(chatId);
      if (!session) {
        throw new Error(`Chat session with ID ${chatId} not found`);
      }
      
      // Add new messages and update end time
      session.messages = [...session.messages, ...newMessages];
      session.endTime = Date.now();
      
      // Save updated session
      await this.saveChatSession(session);
    } catch (error) {
      console.error(`Error updating chat session ${chatId}:`, error);
      throw error;
    }
  }

  /**
   * Get a chat session by its ID
   */
  async getChatSessionById(chatId: string): Promise<ChatSession | null> {
    try {
      return await this.getObject<ChatSession>(`chats/${chatId}/data.json`);
    } catch (error) {
      console.error(`Error retrieving chat session ${chatId}:`, error);
      return null;
    }
  }

  /**
   * Get all chat sessions for a user
   */
  async getChatSessionsByUser(userId: string): Promise<ChatSession[]> {
    return await this.getSessionsByPrefix(`users/${userId}/chats/`);
  }

  /**
   * Get all chat sessions for a workspace
   */
  async getChatSessionsByWorkspace(workspaceId: string): Promise<ChatSession[]> {
    return await this.getSessionsByPrefix(`workspaces/${workspaceId}/chats/`);
  }

  /**
   * Get all chat sessions for a document
   */
  async getChatSessionsByDocument(documentId: string): Promise<ChatSession[]> {
    return await this.getSessionsByPrefix(`documents/${documentId}/chats/`);
  }

  /**
   * Get chat sessions by user in a specific workspace
   */
  async getChatSessionsByUserAndWorkspace(userId: string, workspaceId: string): Promise<ChatSession[]> {
    return await this.getSessionsByPrefix(`compound/user-workspace/${userId}/${workspaceId}/chats/`);
  }

  /**
   * Create a new chat session
   */
  createNewSession(
    userId: string,
    workspaceId: string,
    documentId?: string,
    metadata?: Record<string, any>
  ): ChatSession {
    return {
      chatId: this.generateUniqueId(),
      userId,
      workspaceId,
      documentId,
      messages: [],
      startTime: Date.now(),
      metadata
    };
  }

  /**
   * Add a message to a chat session
   */
  async addMessageToSession(
    chatId: string,
    message: Omit<ChatMessage, 'id' | 'timestamp'>
  ): Promise<void> {
    const fullMessage: ChatMessage = {
      ...message,
      id: this.generateUniqueId(),
      timestamp: Date.now()
    };
    
    await this.updateChatSession(chatId, [fullMessage]);
  }

  // Helper methods
  private async saveObject(key: string, data: string): Promise<void> {
    try {
      await this.s3Client.send(new PutObjectCommand({
        Bucket: this.bucketName,
        Key: key,
        Body: data,
        ContentType: 'application/json'
      }));
    } catch (error) {
      console.error(`Error saving object with key ${key}:`, error);
      throw error;
    }
  }

  private async getObject<T>(key: string): Promise<T | null> {
    try {
      const response = await this.s3Client.send(new GetObjectCommand({
        Bucket: this.bucketName,
        Key: key
      }));
      
      if (!response.Body) {
        return null;
      }
      
      const streamReader = response.Body.transformToString();
      const jsonData = await streamReader;
      return JSON.parse(jsonData) as T;
    } catch (error) {
      console.error(`Error retrieving object with key ${key}:`, error);
      return null;
    }
  }

  private async getSessionsByPrefix(prefix: string): Promise<ChatSession[]> {
    try {
      const response = await this.s3Client.send(new ListObjectsV2Command({
        Bucket: this.bucketName,
        Prefix: prefix
      }));
      
      if (!response.Contents || response.Contents.length === 0) {
        return [];
      }
      
      // Filter for reference files
      const referenceKeys = response.Contents
        .map(item => item.Key!)
        .filter(key => key.endsWith('reference.json'));
      
      // Get all references
      const references = await Promise.all(
        referenceKeys.map(key => this.getObject<{type: string; targetKey: string; chatId: string}>(key))
      );
      
      // Get all chat sessions from the references
      const chatSessions = await Promise.all(
        references
          .filter((ref): ref is {type: string; targetKey: string; chatId: string} => 
            ref !== null && ref.type === 'reference' && !!ref.targetKey
          )
          .map(ref => this.getObject<ChatSession>(ref.targetKey))
      );
      
      return chatSessions.filter((session): session is ChatSession => session !== null);
    } catch (error) {
      console.error(`Error retrieving chat sessions for prefix ${prefix}:`, error);
      return [];
    }
  }

  private generateUniqueId(): string {
    const timestamp = Date.now().toString();
    const random = Math.random().toString();
    return createHash('sha256').update(timestamp + random).digest('hex').substring(0, 16);
  }
}

// Example usage
async function example() {
  const storage = new ChatHistoryStorage({
    region: 'us-east-1',
    bucketName: 'my-chatbot-data'
  });
  
  // Create a new session
  const session = storage.createNewSession(
    'user123',
    'workspace456',
    'document789',
    { source: 'web', context: 'product support' }
  );
  
  // Add some messages
  await storage.addMessageToSession(session.chatId, {
    role: 'user',
    content: 'How do I use this feature?'
  });
  
  await storage.addMessageToSession(session.chatId, {
    role: 'assistant',
    content: 'To use this feature, first navigate to...'
  });
  
  // Retrieve chat by different indices
  const userChats = await storage.getChatSessionsByUser('user123');
  console.log(`Found ${userChats.length} chats for user123`);
  
  const workspaceChats = await storage.getChatSessionsByWorkspace('workspace456');
  console.log(`Found ${workspaceChats.length} chats in workspace456`);
  
  const documentChats = await storage.getChatSessionsByDocument('document789');
  console.log(`Found ${documentChats.length} chats related to document789`);
  
  const userWorkspaceChats = await storage.getChatSessionsByUserAndWorkspace('user123', 'workspace456');
  console.log(`Found ${userWorkspaceChats.length} chats for user123 in workspace456`);
}

export { ChatHistoryStorage, type ChatSession, type ChatMessage, type S3StorageConfig };