import { Memory } from "../clients/bedrock";
import { Memory as GenericMemory } from "../llm-invoker";

export interface MemoryBuilder {
  addMemory(text: string): MemoryBuilder;
  addMemories(texts: string[]): MemoryBuilder;
  build(): Memory[] | GenericMemory[];
}

export class BedrockMemoryBuilder {
  private memories: Memory[] = [];

  addMemory(text: string): BedrockMemoryBuilder {
    this.memories.push({text});
    return this;
  }

  addMemories(texts: string[]): BedrockMemoryBuilder {
    texts.forEach(text => this.addMemory(text));
    return this;
  }

  build(): Memory[] {
    return this.memories;
  }
}

export class GenericMemoryBuilder implements MemoryBuilder {
  private memories: GenericMemory[] = [];

  addMemory(text: string): MemoryBuilder {
    this.memories.push({ content: text });
    return this;
  }

  addMemories(texts: string[]): MemoryBuilder {
    texts.forEach(text => this.addMemory(text));
    return this;
  }

  build(): GenericMemory[] {
    return this.memories;
  }
}