{"version":3,"file":"entity_memory.cjs","names":["BaseChatMemory","LLMChain","ENTITY_EXTRACTION_PROMPT","ENTITY_SUMMARIZATION_PROMPT","InMemoryEntityStore"],"sources":["../../src/memory/entity_memory.ts"],"sourcesContent":["import type { BaseLanguageModelInterface } from \"@langchain/core/language_models/base\";\nimport { PromptTemplate } from \"@langchain/core/prompts\";\n\nimport {\n  InputValues,\n  MemoryVariables,\n  OutputValues,\n  getPromptInputKey,\n} from \"@langchain/core/memory\";\nimport { getBufferString } from \"@langchain/core/messages\";\nimport { InMemoryEntityStore } from \"./stores/entity/in_memory.js\";\nimport { LLMChain } from \"../chains/llm_chain.js\";\nimport {\n  ENTITY_EXTRACTION_PROMPT,\n  ENTITY_SUMMARIZATION_PROMPT,\n} from \"./prompt.js\";\nimport { BaseEntityStore } from \"./stores/entity/base.js\";\nimport { BaseChatMemory, BaseChatMemoryInput } from \"./chat_memory.js\";\n\n/**\n * Interface for the input parameters required by the EntityMemory class.\n */\nexport interface EntityMemoryInput extends BaseChatMemoryInput {\n  llm: BaseLanguageModelInterface;\n  humanPrefix?: string;\n  aiPrefix?: string;\n  entityExtractionPrompt?: PromptTemplate;\n  entitySummarizationPrompt?: PromptTemplate;\n  entityCache?: string[];\n  k?: number;\n  chatHistoryKey?: string;\n  entitiesKey?: string;\n  entityStore?: BaseEntityStore;\n}\n\n// Entity extractor & summarizer to memory.\n/**\n * Class for managing entity extraction and summarization to memory in\n * chatbot applications. Extends the BaseChatMemory class and implements\n * the EntityMemoryInput interface.\n * @example\n * ```typescript\n * const memory = new EntityMemory({\n *   llm: new ChatOpenAI({ model: \"gpt-4o-mini\", temperature: 0 }),\n *   chatHistoryKey: \"history\",\n *   entitiesKey: \"entities\",\n * });\n * const model = new ChatOpenAI({ model: \"gpt-4o-mini\", temperature: 0.9 });\n * const chain = new LLMChain({\n *   llm: model,\n *   prompt: ENTITY_MEMORY_CONVERSATION_TEMPLATE,\n *   memory,\n * });\n *\n * const res1 = await chain.call({ input: \"Hi! I'm Jim.\" });\n * console.log({\n *   res1,\n *   memory: await memory.loadMemoryVariables({ input: \"Who is Jim?\" }),\n * });\n *\n * const res2 = await chain.call({\n *   input: \"I work in construction. What about you?\",\n * });\n * console.log({\n *   res2,\n *   memory: await memory.loadMemoryVariables({ input: \"Who is Jim?\" }),\n * });\n *\n * ```\n */\nexport class EntityMemory extends BaseChatMemory implements EntityMemoryInput {\n  private entityExtractionChain: LLMChain;\n\n  private entitySummarizationChain: LLMChain;\n\n  entityStore: BaseEntityStore;\n\n  entityCache: string[] = [];\n\n  k = 3;\n\n  chatHistoryKey = \"history\";\n\n  llm: BaseLanguageModelInterface;\n\n  entitiesKey = \"entities\";\n\n  humanPrefix?: string;\n\n  aiPrefix?: string;\n\n  constructor(fields: EntityMemoryInput) {\n    super({\n      chatHistory: fields.chatHistory,\n      returnMessages: fields.returnMessages ?? false,\n      inputKey: fields.inputKey,\n      outputKey: fields.outputKey,\n    });\n    this.llm = fields.llm;\n    this.humanPrefix = fields.humanPrefix;\n    this.aiPrefix = fields.aiPrefix;\n    this.chatHistoryKey = fields.chatHistoryKey ?? this.chatHistoryKey;\n    this.entitiesKey = fields.entitiesKey ?? this.entitiesKey;\n    this.entityExtractionChain = new LLMChain({\n      llm: this.llm,\n      prompt: fields.entityExtractionPrompt ?? ENTITY_EXTRACTION_PROMPT,\n    });\n    this.entitySummarizationChain = new LLMChain({\n      llm: this.llm,\n      prompt: fields.entitySummarizationPrompt ?? ENTITY_SUMMARIZATION_PROMPT,\n    });\n    this.entityStore = fields.entityStore ?? new InMemoryEntityStore();\n    this.entityCache = fields.entityCache ?? this.entityCache;\n    this.k = fields.k ?? this.k;\n  }\n\n  get memoryKeys() {\n    return [this.chatHistoryKey];\n  }\n\n  // Will always return list of memory variables.\n  get memoryVariables(): string[] {\n    return [this.entitiesKey, this.chatHistoryKey];\n  }\n\n  // Return history buffer.\n  /**\n   * Method to load memory variables and perform entity extraction.\n   * @param inputs Input values for the method.\n   * @returns Promise resolving to an object containing memory variables.\n   */\n  async loadMemoryVariables(inputs: InputValues): Promise<MemoryVariables> {\n    const promptInputKey =\n      this.inputKey ?? getPromptInputKey(inputs, this.memoryVariables);\n    const messages = await this.chatHistory.getMessages();\n    const serializedMessages = getBufferString(\n      messages.slice(-this.k * 2),\n      this.humanPrefix,\n      this.aiPrefix\n    );\n    const output = await this.entityExtractionChain.predict({\n      history: serializedMessages,\n      input: inputs[promptInputKey],\n    });\n    const entities: string[] =\n      output.trim() === \"NONE\" ? [] : output.split(\",\").map((w) => w.trim());\n    const entitySummaries: { [key: string]: string | undefined } = {};\n\n    for (const entity of entities) {\n      entitySummaries[entity] = await this.entityStore.get(\n        entity,\n        \"No current information known.\"\n      );\n    }\n    this.entityCache = [...entities];\n    const buffer = this.returnMessages\n      ? messages.slice(-this.k * 2)\n      : serializedMessages;\n\n    return {\n      [this.chatHistoryKey]: buffer,\n      [this.entitiesKey]: entitySummaries,\n    };\n  }\n\n  // Save context from this conversation to buffer.\n  /**\n   * Method to save the context from a conversation to a buffer and perform\n   * entity summarization.\n   * @param inputs Input values for the method.\n   * @param outputs Output values from the method.\n   * @returns Promise resolving to void.\n   */\n  async saveContext(inputs: InputValues, outputs: OutputValues): Promise<void> {\n    await super.saveContext(inputs, outputs);\n\n    const promptInputKey =\n      this.inputKey ?? getPromptInputKey(inputs, this.memoryVariables);\n    const messages = await this.chatHistory.getMessages();\n    const serializedMessages = getBufferString(\n      messages.slice(-this.k * 2),\n      this.humanPrefix,\n      this.aiPrefix\n    );\n    const inputData = inputs[promptInputKey];\n\n    for (const entity of this.entityCache) {\n      const existingSummary = await this.entityStore.get(\n        entity,\n        \"No current information known.\"\n      );\n      const output = await this.entitySummarizationChain.predict({\n        summary: existingSummary,\n        entity,\n        history: serializedMessages,\n        input: inputData,\n      });\n      if (output.trim() !== \"UNCHANGED\") {\n        await this.entityStore.set(entity, output.trim());\n      }\n    }\n  }\n\n  // Clear memory contents.\n  /**\n   * Method to clear the memory contents.\n   * @returns Promise resolving to void.\n   */\n  async clear() {\n    await super.clear();\n    await this.entityStore.clear();\n  }\n}\n"],"mappings":";;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;AAsEA,IAAa,eAAb,cAAkCA,2BAAAA,eAA4C;CAC5E;CAEA;CAEA;CAEA,cAAwB,EAAE;CAE1B,IAAI;CAEJ,iBAAiB;CAEjB;CAEA,cAAc;CAEd;CAEA;CAEA,YAAY,QAA2B;AACrC,QAAM;GACJ,aAAa,OAAO;GACpB,gBAAgB,OAAO,kBAAkB;GACzC,UAAU,OAAO;GACjB,WAAW,OAAO;GACnB,CAAC;AACF,OAAK,MAAM,OAAO;AAClB,OAAK,cAAc,OAAO;AAC1B,OAAK,WAAW,OAAO;AACvB,OAAK,iBAAiB,OAAO,kBAAkB,KAAK;AACpD,OAAK,cAAc,OAAO,eAAe,KAAK;AAC9C,OAAK,wBAAwB,IAAIC,kBAAAA,SAAS;GACxC,KAAK,KAAK;GACV,QAAQ,OAAO,0BAA0BC,eAAAA;GAC1C,CAAC;AACF,OAAK,2BAA2B,IAAID,kBAAAA,SAAS;GAC3C,KAAK,KAAK;GACV,QAAQ,OAAO,6BAA6BE,eAAAA;GAC7C,CAAC;AACF,OAAK,cAAc,OAAO,eAAe,IAAIC,kBAAAA,qBAAqB;AAClE,OAAK,cAAc,OAAO,eAAe,KAAK;AAC9C,OAAK,IAAI,OAAO,KAAK,KAAK;;CAG5B,IAAI,aAAa;AACf,SAAO,CAAC,KAAK,eAAe;;CAI9B,IAAI,kBAA4B;AAC9B,SAAO,CAAC,KAAK,aAAa,KAAK,eAAe;;;;;;;CAShD,MAAM,oBAAoB,QAA+C;EACvE,MAAM,iBACJ,KAAK,aAAA,GAAA,uBAAA,mBAA8B,QAAQ,KAAK,gBAAgB;EAClE,MAAM,WAAW,MAAM,KAAK,YAAY,aAAa;EACrD,MAAM,sBAAA,GAAA,yBAAA,iBACJ,SAAS,MAAM,CAAC,KAAK,IAAI,EAAE,EAC3B,KAAK,aACL,KAAK,SACN;EACD,MAAM,SAAS,MAAM,KAAK,sBAAsB,QAAQ;GACtD,SAAS;GACT,OAAO,OAAO;GACf,CAAC;EACF,MAAM,WACJ,OAAO,MAAM,KAAK,SAAS,EAAE,GAAG,OAAO,MAAM,IAAI,CAAC,KAAK,MAAM,EAAE,MAAM,CAAC;EACxE,MAAM,kBAAyD,EAAE;AAEjE,OAAK,MAAM,UAAU,SACnB,iBAAgB,UAAU,MAAM,KAAK,YAAY,IAC/C,QACA,gCACD;AAEH,OAAK,cAAc,CAAC,GAAG,SAAS;EAChC,MAAM,SAAS,KAAK,iBAChB,SAAS,MAAM,CAAC,KAAK,IAAI,EAAE,GAC3B;AAEJ,SAAO;IACJ,KAAK,iBAAiB;IACtB,KAAK,cAAc;GACrB;;;;;;;;;CAWH,MAAM,YAAY,QAAqB,SAAsC;AAC3E,QAAM,MAAM,YAAY,QAAQ,QAAQ;EAExC,MAAM,iBACJ,KAAK,aAAA,GAAA,uBAAA,mBAA8B,QAAQ,KAAK,gBAAgB;EAElE,MAAM,sBAAA,GAAA,yBAAA,kBADW,MAAM,KAAK,YAAY,aAAa,EAE1C,MAAM,CAAC,KAAK,IAAI,EAAE,EAC3B,KAAK,aACL,KAAK,SACN;EACD,MAAM,YAAY,OAAO;AAEzB,OAAK,MAAM,UAAU,KAAK,aAAa;GACrC,MAAM,kBAAkB,MAAM,KAAK,YAAY,IAC7C,QACA,gCACD;GACD,MAAM,SAAS,MAAM,KAAK,yBAAyB,QAAQ;IACzD,SAAS;IACT;IACA,SAAS;IACT,OAAO;IACR,CAAC;AACF,OAAI,OAAO,MAAM,KAAK,YACpB,OAAM,KAAK,YAAY,IAAI,QAAQ,OAAO,MAAM,CAAC;;;;;;;CAUvD,MAAM,QAAQ;AACZ,QAAM,MAAM,OAAO;AACnB,QAAM,KAAK,YAAY,OAAO"}