memories-lite 0.9.0

package/src/memory/index.ts

@@ -0,0 +1,723 @@
+import { v4 as uuidv4 } from "uuid";
+import { createHash } from "crypto";
+import {
+  MemoryConfig,
+  MemoryConfigSchema,
+  MemoryItem,
+  MemoryType,
+  Message,
+  SearchFilters,
+  SearchResult,
+  VectorStoreConfig,
+} from "../types";
+import {
+  EmbedderFactory,
+  LLMFactory,
+  HistoryManagerFactory,
+} from "../utils/factory";
+import {
+  FactRetrievalSchema_extended,
+  getFactRetrievalMessages,
+  getUpdateMemoryMessages,
+  MemoryUpdateSchema,
+  removeCodeBlocks,
+} from "../prompts";
+import { DummyHistoryManager } from "../storage/DummyHistoryManager";
+import { Embedder } from "../embeddings/base";
+import { LLM } from "../llms/base";
+import { VectorStore } from "../vectorstores/base";
+import { ConfigManager } from "../config/manager";
+import {
+  AddMemoryOptions,
+  SearchMemoryOptions,
+  DeleteAllMemoryOptions,
+  GetAllMemoryOptions,
+} from "./memory.types";
+import { parse_vision_messages } from "../utils/memory";
+import { HistoryManager } from "../storage/base";
+import { captureClientEvent } from "../utils/telemetry";
+import { LiteVectorStore } from "../vectorstores/lite";
+import { zodResponseFormat } from "openai/helpers/zod";
+
+
+export class MemoriesLite {
+  private config: MemoryConfig;
+  private customPrompt: string | undefined;
+  private embedder: Embedder;
+  private vectorStoreConfig: VectorStoreConfig;
+  private llm: LLM;
+  private db: HistoryManager;
+  private collectionName: string | undefined;
+  private apiVersion: string;
+  private graphMemory?: any;
+  private enableGraph: boolean;
+  telemetryId: string;
+
+  constructor(config: Partial<MemoryConfig> = {}) {
+    // Merge and validate config
+    this.config = ConfigManager.mergeConfig(config);
+
+    this.customPrompt = this.config.customPrompt;
+    this.embedder = EmbedderFactory.create(
+      this.config.embedder.provider,
+      this.config.embedder.config,
+    );
+    //
+    // vectorStore.provider is "lite" 
+    this.vectorStoreConfig = this.config.vectorStore.config;
+    this.llm = LLMFactory.create(
+      this.config.llm.provider,
+      this.config.llm.config,
+    );
+    if (this.config.disableHistory) {
+      this.db = new DummyHistoryManager();
+    } else {
+      const defaultConfig = {
+        provider: "sqlite",
+        config: {
+          historyDbPath: this.config.historyDbPath || ":memory:",
+        },
+      };
+
+      this.db =
+        this.config.historyStore && !this.config.disableHistory
+          ? HistoryManagerFactory.create(
+              this.config.historyStore.provider,
+              this.config.historyStore,
+            )
+          : HistoryManagerFactory.create("sqlite", defaultConfig);
+    }
+
+    this.collectionName = this.config.vectorStore.config.collectionName;
+    this.apiVersion = this.config.version || "v1.0";
+    this.enableGraph = this.config.enableGraph || false;
+    this.telemetryId = "anonymous";
+
+    // Initialize graph memory if configured
+    if (this.enableGraph && this.config.graphStore) {
+      // this.graphMemory = new MemoryGraph(this.config);
+    }
+
+    // Initialize telemetry if vector store is initialized
+    // this._initializeTelemetry();
+  }
+
+  private async _initializeTelemetry() {
+    try {
+      await this._getTelemetryId();
+
+      // Capture initialization event
+      await captureClientEvent("init", this, {
+        api_version: this.apiVersion,
+        client_type: "Memory",
+        collection_name: this.collectionName,
+        enable_graph: this.enableGraph,
+      });
+    } catch (error) {}
+  }
+
+  private async _getTelemetryId() {
+    try {
+      if (
+        !this.telemetryId ||
+        this.telemetryId === "anonymous" ||
+        this.telemetryId === "anonymous-supabase"
+      ) {
+        // this.telemetryId = await this.vectorStore.getUserId();
+        throw new Error("Telemetry ID not found");
+      }
+      return this.telemetryId;
+    } catch (error) {
+      this.telemetryId = "anonymous";
+      return this.telemetryId;
+    }
+  }
+
+  private async _captureEvent(methodName: string, additionalData = {}) {
+    try {
+      await this._getTelemetryId();
+      await captureClientEvent(methodName, this, {
+        ...additionalData,
+        api_version: this.apiVersion,
+        collection_name: this.collectionName,
+      });
+    } catch (error) {
+      console.error(`Failed to capture ${methodName} event:`, error);
+    }
+  }
+
+  private $t(text:string) {
+    // return text.replace(/<thinking>[\s\S]*?(?:<\/thinking>)/g,'').replace(/^<step.*$/g,'');
+    return text.replace(/<thinking>[\s\S]*?<\/thinking>/g,'').replace(/^<step.*$/gm,'').replace(/<memories>[\s\S]*?<\/memories>/g, '');
+  }
+
+
+  private async addToVectorStore(
+    messages: Message[],
+    metadata: Record<string, any>,
+    userId: string,
+    filters: SearchFilters,
+    customFacts?: string,
+  ): Promise<MemoryItem[]> {
+
+    const $t = this.$t;
+    const vectorStore = await this.getVectorStore(userId);
+    const parsedMessages = messages.filter((m) => typeof m.content === 'string').map((m) => `${m.role=='user' ? '**USER**: ' : '**ASSISTANT**: '}${$t(m.content as string)}\n`).join("\n");
+
+    const [systemPrompt, userPrompt] = getFactRetrievalMessages(parsedMessages, customFacts||this.customPrompt);
+
+    const response = await this.llm.generateResponse(
+      [
+        { role: "system", content: systemPrompt },
+        { role: "user", content: userPrompt },
+      ],
+      {...zodResponseFormat(FactRetrievalSchema_extended,"FactRetrieval")},[],false
+    );
+    const parsedResponse = (response:any) => {
+      try {
+        // structured output
+        if(typeof response === 'object') {
+          return response;
+        }
+        const cleanResponse = removeCodeBlocks(response as string);
+        return JSON.parse(cleanResponse);
+      } catch (e) {
+        console.error(
+          "Failed to parse facts from LLM response:",
+          response,
+          e,
+          response
+        );
+        return [];
+      }
+    }
+    // 
+    // can use native structured output
+    const facts = parsedResponse(response).facts?.filter((f:any) => !f.existing)||[];
+
+    // console.log("-- DBG extract:", userPrompt);
+    // console.log("-- DBG facts:", facts);
+
+    // Get embeddings for new facts
+    const newMessageEmbeddings: Record<string, number[]> = {};
+    const retrievedOldMemory: Array<{ id: string; text: string; type: string }> = [];
+
+    // Create embeddings and search for similar memories
+    for (const elem of facts) {
+      const fact = elem.fact;
+      const embedding = await this.embedder.embed(fact);
+      newMessageEmbeddings[fact] = embedding;
+
+      const existingMemories = await vectorStore.search(
+        embedding,
+        5,
+        filters,
+      );
+      for (const mem of existingMemories) {
+        retrievedOldMemory.push({ id: mem.id, text: mem.payload.data,type: mem.payload.type });
+      }
+    }
+
+    // console.log("-- DBG old memories:", retrievedOldMemory);
+    // Remove duplicates from old memories
+    const uniqueOldMemories = retrievedOldMemory.filter(
+      (mem, index) =>
+        retrievedOldMemory.findIndex((m) => m.id === mem.id) === index,
+    );
+
+    // Create UUID mapping for handling UUID hallucinations
+    const tempUuidMapping: Record<string, string> = {};
+    uniqueOldMemories.forEach((item, idx) => {
+      tempUuidMapping[String(idx)] = item.id;
+      uniqueOldMemories[idx].id = String(idx);
+    });
+
+    // Get memory update decisions
+    const updatePrompt = getUpdateMemoryMessages(uniqueOldMemories, facts);
+
+    const updateResponse = await this.llm.generateResponse(
+      [{ role: "user", content: updatePrompt }],
+      {...zodResponseFormat(MemoryUpdateSchema,"Memory")},[],false,
+    );
+    // console.log("-- DBG merge:", updatePrompt);
+
+    const memoryActions: any[] = parsedResponse(updateResponse).memory || [];
+
+    // Process memory actions
+    const results: MemoryItem[] = [];
+    for (const action of memoryActions) {
+      if(action.reason === "undefined") {
+        console.log(`-- ⛔ LLM Error: ${action.event}, ${action.type}, "${action.text}"`);
+        continue;
+      }
+      console.log(`-- DBG memory action: ${action.event}, ${action.type}, "${action.text}", why: "${action.reason}"`);
+      try {
+        switch (action.event) {
+          case "ADD": {
+            if(!action.type) {
+              // log error
+              console.error("Type is mandatory to manage memories:", action);
+              continue;
+            }
+            metadata.type = action.type;
+            const memoryId = await this.createMemory(
+              action.text,
+              newMessageEmbeddings,
+              metadata,
+              userId,
+            );
+            results.push({
+              id: memoryId,
+              memory: action.text,
+              type: action.type,
+              metadata: { event: action.event },
+            });
+            break;
+          }
+          case "UPDATE": {
+            const realMemoryId = tempUuidMapping[action.id];
+            const type = uniqueOldMemories[action.id].type as MemoryType;
+            await this.updateMemory(
+              realMemoryId,
+              action.text,
+              newMessageEmbeddings,
+              metadata,
+              userId,
+            );
+            results.push({
+              id: realMemoryId,
+              memory: action.text,
+              type,
+              metadata: {
+                event: action.event,
+                previousMemory: action.old_memory,
+              },
+            });
+            break;
+          }
+          case "DELETE": {
+            const realMemoryId = tempUuidMapping[action.id];
+            await this.deleteMemory(realMemoryId, userId);
+            results.push({
+              id: realMemoryId,
+              memory: action.text,
+              type: action.type,
+              metadata: { event: action.event },
+            });
+            break;
+          }
+        }
+      } catch (error) {
+        console.error(`Error processing memory action: ${error}`);
+      }
+    }
+    return results;
+  }  
+
+  static fromConfig(configDict: Record<string, any>): MemoriesLite {
+    try {
+      const config = MemoryConfigSchema.parse(configDict);
+      return new MemoriesLite(config);
+    } catch (e) {
+      console.error("Configuration validation error:", e);
+      throw e;
+    }
+  }
+
+  async getVectorStore(userId: string): Promise<VectorStore> {
+    return LiteVectorStore.from(userId, this.vectorStoreConfig);
+  }
+
+  async capture(
+    messages: string | Message[],
+    userId: string,
+    config: AddMemoryOptions,
+  ): Promise<SearchResult> {
+    // await this._captureEvent("add", {
+    //   message_count: Array.isArray(messages) ? messages.length : 1,
+    //   has_metadata: !!config.metadata,
+    //   has_filters: !!config.filters,
+    //   infer: config.infer,
+    // });
+    const {
+      agentId,
+      runId,
+      metadata = {},
+      filters = {},
+      infer = true,
+      customFacts
+    } = config;
+
+    if (agentId) filters.agentId = metadata.agentId = agentId;
+    if (runId) filters.runId = metadata.runId = runId;
+
+    if (!userId && !filters.agentId && !filters.runId) {
+      throw new Error(
+        "One of the filters: userId, agentId or runId is required!",
+      );
+    }
+
+    const parsedMessages = Array.isArray(messages)
+      ? (messages as Message[])
+      : [{ role: "user", content: messages }];
+
+    const final_parsedMessages = await parse_vision_messages(parsedMessages);
+
+    // Add to vector store
+    const vectorStoreResult = await this.addToVectorStore(
+      final_parsedMessages,
+      metadata,
+      userId,
+      filters,
+      customFacts,
+    );
+
+    // Add to graph store if available
+    let graphResult;
+    if (this.graphMemory) {
+      try {
+        graphResult = await this.graphMemory.add(
+          final_parsedMessages.map((m) => m.content).join("\n"),
+          filters,
+        );
+      } catch (error) {
+        console.error("Error adding to graph memory:", error);
+      }
+    }
+
+    return {
+      results: vectorStoreResult,
+      relations: graphResult?.relations,
+    };
+  }
+
+
+  async get(memoryId: string, userId: string): Promise<MemoryItem | null> {
+    const vectorStore = await this.getVectorStore(userId);
+
+    const memory = await vectorStore.get(memoryId);
+    if (!memory) return null;
+
+    const filters = {
+      ...(memory.payload.agentId && { agentId: memory.payload.agentId }),
+      ...(memory.payload.runId && { runId: memory.payload.runId }),
+    };
+
+    const memoryItem: MemoryItem = {
+      id: memory.id,
+      memory: memory.payload.data,
+      type: memory.payload.type,
+      hash: memory.payload.hash,
+      createdAt: memory.payload.createdAt,
+      updatedAt: memory.payload.updatedAt,
+      metadata: {},
+    };
+
+
+    // Add additional metadata
+    const excludedKeys = new Set([
+      "userId",
+      "agentId",
+      "runId",
+      "hash",
+      "data",
+      "createdAt",
+      "updatedAt",
+    ]);
+    for (const [key, value] of Object.entries(memory.payload)) {
+      if (!excludedKeys.has(key)) {
+        memoryItem.metadata![key] = value;
+      }
+    }
+
+    return { ...memoryItem, ...filters };
+  }
+
+  async retrieve(
+    query: string,
+    userId: string,
+    config: SearchMemoryOptions,
+  ): Promise<SearchResult> {
+    // await this._captureEvent("search", {
+    //   query_length: query.length,
+    //   limit: config.limit,
+    //   has_filters: !!config.filters,
+    // });
+    const { agentId, runId, limit = 100, filters = {} } = config;
+
+    if (agentId) filters.agentId = agentId;
+    if (runId) filters.runId = runId;
+
+    if (!userId && !filters.agentId && !filters.runId) {
+      throw new Error(
+        "One of the filters: userId, agentId or runId is required!",
+      );
+    }
+
+    const vectorStore = await this.getVectorStore(userId);
+
+
+    // Search vector store
+    const queryEmbedding = await this.embedder.embed(query);
+    const memories = await vectorStore.search(
+      queryEmbedding,
+      limit,
+      filters,
+    );
+
+
+    // Search graph store if available
+    let graphResults =[];
+    if (this.graphMemory) {
+      try {
+        graphResults = await this.graphMemory.search(query, filters);
+      } catch (error) {
+        console.error("Error searching graph memory:", error);
+      }
+    }
+
+    const excludedKeys = new Set([
+      "userId",
+      "agentId",
+      "runId",
+      "hash",
+      "data",
+      "createdAt",
+      "updatedAt",
+    ]);
+    const results = memories.map((mem) => ({
+      id: mem.id,
+      memory: mem.payload.data,
+      hash: mem.payload.hash,
+      type: mem.payload.type,
+      createdAt: mem.payload.createdAt,
+      updatedAt: mem.payload.updatedAt,
+      score: mem.score,
+      metadata: Object.entries(mem.payload)
+        .filter(([key]) => !excludedKeys.has(key))
+        .reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}),
+      ...(mem.payload.agentId && { agentId: mem.payload.agentId }),
+      ...(mem.payload.runId && { runId: mem.payload.runId }),
+    }));
+
+    return {
+      results,
+      relations: graphResults,
+    };
+  }
+
+  async update(memoryId: string, data: string, userId: string): Promise<{ message: string }> {
+    // await this._captureEvent("update", { memory_id: memoryId });
+    const embedding = await this.embedder.embed(data);
+    await this.updateMemory(memoryId, data, { [data]: embedding }, {}, userId);
+    return { message: "Memory updated successfully!" };
+  }
+
+  async delete(memoryId: string, userId: string): Promise<{ message: string }> {
+    // await this._captureEvent("delete", { memory_id: memoryId });
+    await this.deleteMemory(memoryId, userId);
+    return { message: "Memory deleted successfully!" };
+  }
+
+  async deleteAll(
+    userId: string,
+    config: DeleteAllMemoryOptions,
+  ): Promise<{ message: string }> {
+    // await this._captureEvent("delete_all", {
+    //   has_user_id: !!config.userId,
+    //   has_agent_id: !!config.agentId,
+    //   has_run_id: !!config.runId,
+    // });
+    const { agentId, runId } = config;
+    if (!userId) throw new Error("vector store instanceId is required");
+    const vectorStore = await this.getVectorStore(userId);
+
+    const filters: SearchFilters = {};
+    if (agentId) filters.agentId = agentId;
+    if (runId) filters.runId = runId;
+
+    if (!Object.keys(filters).length) {
+      throw new Error(
+        "At least one filter is required to delete all memories. If you want to delete all memories, use the `reset()` method.",
+      );
+    }
+
+    const [memories] = await vectorStore.list(filters);
+    for (const memory of memories) {
+      await this.deleteMemory(memory.id, userId!);
+    }
+
+    return { message: "Memories deleted successfully!" };
+  }
+
+  //
+  // DEPRECATED: history must be user specific
+  async history(memoryId: string): Promise<any[]> {
+    return this.db.getHistory(memoryId);
+  }
+
+  async reset(userId: string): Promise<void> {
+    // await this._captureEvent("reset");
+    const vectorStore = await this.getVectorStore(userId);
+    await this.db.reset();
+
+    // Check provider before attempting deleteCol
+    await vectorStore.deleteCol();
+
+    if (this.graphMemory) {
+      await this.graphMemory.deleteAll({ userId: "default" }); // Assuming this is okay, or needs similar check?
+    }
+
+    // Re-initialize factories/clients based on the original config
+    this.embedder = EmbedderFactory.create(
+      this.config.embedder.provider,
+      this.config.embedder.config,
+    );
+    this.llm = LLMFactory.create(
+      this.config.llm.provider,
+      this.config.llm.config,
+    );
+    // Re-init DB if needed (though db.reset() likely handles its state)
+    // Re-init Graph if needed
+
+    // Re-initialize telemetry
+    // this._initializeTelemetry();
+  }
+
+  async getAll(userId: string, config: GetAllMemoryOptions): Promise<SearchResult> {
+    // await this._captureEvent("get_all", {
+    //   limit: config.limit,
+    //   has_agent_id: !!config.agentId,
+    //   has_run_id: !!config.runId,
+    // });
+    const { agentId, runId, limit = 10, type } = config;
+    const vectorStore = await this.getVectorStore(userId!);
+
+    const filters: SearchFilters = {};
+    if (agentId) filters.agentId = agentId;
+    if (runId) filters.runId = runId;
+    if (type) filters.type = type;
+    const [memories] = await vectorStore.list(filters, limit);
+
+    const excludedKeys = new Set([
+      "userId",
+      "agentId",
+      "runId",
+      "hash",
+      "data",
+      "createdAt",
+      "updatedAt",
+    ]);
+    const results = memories.map((mem) => ({
+      id: mem.id,
+      memory: mem.payload.data,
+      hash: mem.payload.hash,
+      type: mem.payload.type,
+      createdAt: mem.payload.createdAt,
+      updatedAt: mem.payload.updatedAt,
+      metadata: Object.entries(mem.payload)
+        .filter(([key]) => !excludedKeys.has(key))
+        .reduce((acc, [key, value]) => ({ ...acc, [key]: value }), {}),
+      ...(mem.payload.agentId && { agentId: mem.payload.agentId }),
+      ...(mem.payload.runId && { runId: mem.payload.runId }),
+    }));
+
+    return { results };
+  }
+
+  private async createMemory(
+    data: string,
+    existingEmbeddings: Record<string, number[]>,
+    metadata: Record<string, any>,
+    userId: string,
+  ): Promise<string> {
+    const vectorStore = await this.getVectorStore(userId);
+    const memoryId = uuidv4();
+    const embedding =
+      existingEmbeddings[data] || (await this.embedder.embed(data));
+
+    const memoryMetadata = {
+      ...metadata,
+      data,
+      hash: createHash("md5").update(data).digest("hex"),
+      createdAt: new Date().toISOString(),
+    };
+
+    await vectorStore.insert([embedding], [memoryId], [memoryMetadata]);
+    await this.db.addHistory(
+      memoryId,
+      null,
+      data,
+      "ADD",
+      memoryMetadata.createdAt,
+    );
+
+    return memoryId;
+  }
+
+  private async updateMemory(
+    memoryId: string,
+    data: string,
+    existingEmbeddings: Record<string, number[]>,
+    metadata: Record<string, any> = {},
+    userId: string,
+  ): Promise<string> {
+    const vectorStore = await this.getVectorStore(userId);
+    const existingMemory = await vectorStore.get(memoryId);
+    if (!existingMemory) {
+      throw new Error(`Memory with ID ${memoryId} not found`);
+    }
+
+    const prevValue = existingMemory.payload.data;
+    const embedding =
+      existingEmbeddings[data] || (await this.embedder.embed(data));
+
+    const newMetadata = {
+      ...metadata,
+      data,
+      hash: createHash("md5").update(data).digest("hex"),
+      createdAt: existingMemory.payload.createdAt,
+      updatedAt: new Date().toISOString(),
+      ...(existingMemory.payload.agentId && {
+        agentId: existingMemory.payload.agentId,
+      }),
+      ...(existingMemory.payload.runId && {
+        runId: existingMemory.payload.runId,
+      }),
+    };
+
+    await vectorStore.update(memoryId, embedding, newMetadata);
+    await this.db.addHistory(
+      memoryId,
+      prevValue,
+      data,
+      "UPDATE",
+      newMetadata.createdAt,
+      newMetadata.updatedAt,
+    );
+
+    return memoryId;
+  }
+
+  private async deleteMemory(memoryId: string, userId: string): Promise<string> {
+    const vectorStore = await this.getVectorStore(userId);
+    const existingMemory = await vectorStore.get(memoryId);
+    if (!existingMemory) {
+      throw new Error(`Memory with ID ${memoryId} not found`);
+    }
+
+    const prevValue = existingMemory.payload.data;
+    await vectorStore.delete(memoryId);
+    await this.db.addHistory(
+      memoryId,
+      prevValue,
+      null,
+      "DELETE",
+      undefined,
+      undefined,
+      1,
+    );
+
+    return memoryId;
+  }
+}