feat(vector-stores): add Databricks provider to TypeScript OSS SDK (#5824)

Co-authored-by: kartik-mem0 <kartik.labhshetwar@mem0.ai>

mem0-ts/src/oss/src/config/manager.ts

@@ -60,9 +60,14 @@ export class ConfigManager {
         })(),
       },
       vectorStore: {
-        provider:
+        // Every factory already matches the provider case-insensitively, so a capitalized
+        // name constructs the right store -- but the `provider === "memory"` comparisons that
+        // pick per-provider entity-store settings do not. Normalize once, here, so those
+        // comparisons cannot silently miss.
+        provider: (
           userConfig.vectorStore?.provider ||
-          DEFAULT_MEMORY_CONFIG.vectorStore.provider,
+          DEFAULT_MEMORY_CONFIG.vectorStore.provider
+        ).toLowerCase(),
         config: (() => {
           const defaultConf = DEFAULT_MEMORY_CONFIG.vectorStore.config;
           const userConf = userConfig.vectorStore?.config;

mem0-ts/src/oss/src/index.ts

@@ -36,6 +36,7 @@ export * from "./vector_stores/langchain";
 export * from "./vector_stores/vectorize";
 export * from "./vector_stores/azure_ai_search";
 export * from "./vector_stores/pgvector";
+export * from "./vector_stores/databricks";
 export * from "./vector_stores/neptune_analytics";
 export * from "./vector_stores/elasticsearch";
 export * from "./vector_stores/upstash_vector";

mem0-ts/src/oss/src/memory/index.ts

@@ -7,6 +7,7 @@ import {
   Message,
   SearchFilters,
   SearchResult,
+  VectorStoreConfig,
 } from "../types";
 import {
   EmbedderFactory,
@@ -286,18 +287,24 @@ export class Memory {
 
   private async getEntityStore(): Promise<VectorStore> {
     if (!this._entityStore) {
+      const entityProvider = this.config.vectorStore.provider;
       const entityCollectionName = `${this.collectionName}_entities`;
-      const entityConfig = {
+      const entityConfig: VectorStoreConfig = {
         ...this.config.vectorStore.config,
         collectionName: entityCollectionName,
       };
       // For file-based stores (memory/SQLite), always use a separate DB for entities
-      if (this.config.vectorStore.provider === "memory") {
+      if (entityProvider === "memory") {
         const basePath = entityConfig.dbPath || getDefaultVectorStoreDbPath();
         entityConfig.dbPath = basePath.replace(/\.db$/, "_entities.db");
       }
+      if (entityProvider === "databricks") {
+        entityConfig.tableName = entityConfig.tableName
+          ? `${entityConfig.tableName}_entities`
+          : entityCollectionName;
+      }
       this._entityStore = VectorStoreFactory.create(
-        this.config.vectorStore.provider,
+        entityProvider,
         entityConfig,
       );
       await this._entityStore.initialize();
@@ -1747,7 +1754,7 @@ export class Memory {
     await this.db.reset();
 
     // Check provider before attempting deleteCol
-    if (this.config.vectorStore.provider.toLowerCase() !== "langchain") {
+    if (this.config.vectorStore.provider !== "langchain") {
       try {
         await this.vectorStore.deleteCol();
       } catch (e) {

mem0-ts/src/oss/src/utils/factory.ts

@@ -55,6 +55,7 @@ import { HuggingFaceEmbedder } from "../embeddings/huggingface";
 import { LangchainVectorStore } from "../vector_stores/langchain";
 import { AzureAISearch } from "../vector_stores/azure_ai_search";
 import { PGVector } from "../vector_stores/pgvector";
+import { DatabricksVectorStore } from "../vector_stores/databricks";
 import { NeptuneAnalyticsVectorStore } from "../vector_stores/neptune_analytics";
 import { VertexAIEmbedder } from "../embeddings/vertexai";
 import { ElasticsearchDB } from "../vector_stores/elasticsearch";
@@ -173,6 +174,8 @@ export class VectorStoreFactory {
         return new VertexAIVectorSearch(config as any);
       case "pgvector":
         return new PGVector(config as any);
+      case "databricks":
+        return new DatabricksVectorStore(config as any);
       case "neptune":
       case "neptune-analytics":
         return new NeptuneAnalyticsVectorStore(config as any);

mem0-ts/src/oss/src/vector_stores/neptune_analytics.ts

@@ -1,10 +1,10 @@
-import {
-  ExecuteQueryCommand,
-  NeptuneGraphClient,
-} from "@aws-sdk/client-neptune-graph";
 import { VectorStore } from "./base";
 import { SearchFilters, VectorStoreConfig, VectorStoreResult } from "../types";
 
+/**
+ * The `@aws-sdk/client-neptune-graph` dependency is loaded on first use via dynamic
+ * `import()` so the package stays optional (mirrors `aws_bedrock.ts`).
+ */
 interface NeptuneAnalyticsConfig extends VectorStoreConfig {
   graphIdentifier?: string;
   endpoint?: string;
@@ -14,7 +14,14 @@ interface NeptuneAnalyticsConfig extends VectorStoreConfig {
 }
 
 interface NeptuneGraphClientLike {
-  send(command: ExecuteQueryCommand): Promise<NeptuneExecuteQueryOutput>;
+  send(command: any): Promise<NeptuneExecuteQueryOutput>;
+}
+
+interface NeptuneSDK {
+  NeptuneGraphClient: new (
+    config: Record<string, any>,
+  ) => NeptuneGraphClientLike;
+  ExecuteQueryCommand: new (input: Record<string, any>) => any;
 }
 
 interface NeptuneExecuteQueryOutput {
@@ -33,7 +40,10 @@ interface WhereClauseResult {
 }
 
 export class NeptuneAnalyticsVectorStore implements VectorStore {
-  private readonly client: NeptuneGraphClientLike;
+  private clientConfig: Record<string, any>;
+  private clientOverride?: NeptuneGraphClientLike;
+  private sdkPromise?: Promise<NeptuneSDK>;
+  private clientPromise?: Promise<NeptuneGraphClientLike>;
   private readonly graphIdentifier: string;
   private readonly collectionName: string;
   private readonly collectionLabel: string;
@@ -54,8 +64,8 @@ export class NeptuneAnalyticsVectorStore implements VectorStore {
     this.userLabelExpr = this.escapeLabel(this.userLabel);
     this.userNodeId = "mem0-user";
     this.dimension = config.dimension || 1536;
-    this.client =
-      config.client || new NeptuneGraphClient(this.buildClientConfig(config));
+    this.clientConfig = this.buildClientConfig(config);
+    this.clientOverride = config.client;
 
     void this.initialize().catch(console.error);
   }
@@ -171,6 +181,18 @@ export class NeptuneAnalyticsVectorStore implements VectorStore {
     const hasPayload = !!payload && Object.keys(payload).length > 0;
     const hasVector = vector.length > 0;
 
+    // ponytail: a combined update writes the payload before the embedding, and Neptune's vector
+    // index isn't transactional -- if the upsert below fails, the new payload would otherwise be
+    // left committed against the stale embedding (searches would match the old vector but return
+    // the new metadata). Capture the prior node so a failed upsert can be restored; this is
+    // best-effort compensation, not a rollback. Only needed when both writes happen -- a
+    // payload-only or vector-only update can't desync.
+    // The restore assumes a single writer per vectorId -- concurrent updates to the same node can
+    // interleave and clobber each other's compensation. AWS advises against concurrent same-vertex
+    // writes to the Neptune Analytics vector index for exactly this reason.
+    const priorResult =
+      hasPayload && hasVector ? await this.get(vectorId) : null;
+
     if (hasPayload) {
       const properties = this.buildStoredPayload(payload);
       await this.executeQuery(
@@ -187,20 +209,45 @@ export class NeptuneAnalyticsVectorStore implements VectorStore {
     }
 
     if (hasVector) {
-      const updateResults = await this.executeQuery(
-        `
-          MATCH (n:${this.collectionLabelExpr} {\`~id\`: $vectorId})
-          WITH n, $embedding AS embedding
-          CALL neptune.algo.vectors.upsert(n, embedding)
-          YIELD success
-          RETURN success
-        `,
-        {
-          vectorId,
-          embedding: vector,
-        },
-      );
-      this.assertSuccessfulResults(updateResults, "Update");
+      try {
+        const updateResults = await this.executeQuery(
+          `
+            MATCH (n:${this.collectionLabelExpr} {\`~id\`: $vectorId})
+            WITH n, $embedding AS embedding
+            CALL neptune.algo.vectors.upsert(n, embedding)
+            YIELD success
+            RETURN success
+          `,
+          {
+            vectorId,
+            embedding: vector,
+          },
+        );
+        this.assertSuccessfulResults(updateResults, "Update");
+      } catch (error) {
+        if (priorResult) {
+          try {
+            await this.executeQuery(
+              `
+                MATCH (n:${this.collectionLabelExpr} {\`~id\`: $vectorId})
+                SET n = $properties
+                RETURN n
+              `,
+              {
+                vectorId,
+                properties: priorResult.payload,
+              },
+            );
+          } catch (restoreError) {
+            // Do not mask the original failure with a compensation failure.
+            console.error(
+              "Neptune Analytics: failed to restore prior payload after a failed update upsert",
+              restoreError,
+            );
+          }
+        }
+        throw error;
+      }
     }
   }
 
@@ -985,12 +1032,58 @@ export class NeptuneAnalyticsVectorStore implements VectorStore {
     }
   }
 
+  /**
+   * Load the optional AWS SDK on first use.
+   *
+   * This MUST be a dynamic `import()`, never `require()`: tsup/esbuild rewrite
+   * `require()` in the published ESM bundle (`dist/oss/index.mjs`) into a
+   * `__require` shim that throws `Dynamic require of "..." is not supported`,
+   * so every ESM consumer would hit a dead provider even with the SDK installed.
+   */
+  private async getSDK(): Promise<NeptuneSDK> {
+    if (!this.sdkPromise) {
+      this.sdkPromise = import("@aws-sdk/client-neptune-graph").then(
+        (sdk) => sdk as unknown as NeptuneSDK,
+        (err) => {
+          // Let a later call retry rather than caching the rejection forever.
+          this.sdkPromise = undefined;
+          const detail = err instanceof Error ? err.message : String(err);
+          throw new Error(
+            "The '@aws-sdk/client-neptune-graph' package is required to use the Neptune Analytics vector store. " +
+              `Install it with: npm install @aws-sdk/client-neptune-graph (original error: ${detail})`,
+          );
+        },
+      );
+    }
+    return this.sdkPromise;
+  }
+
+  /** Memoized Neptune client; an injected `config.client` short-circuits the SDK. */
+  private async getClient(): Promise<NeptuneGraphClientLike> {
+    if (this.clientOverride) return this.clientOverride;
+    if (!this.clientPromise) {
+      this.clientPromise = this.getSDK()
+        .then(
+          ({ NeptuneGraphClient }) => new NeptuneGraphClient(this.clientConfig),
+        )
+        .catch((err) => {
+          // Mirror getSDK(): drop the rejected promise so a later call retries rather than
+          // replaying a cached rejection forever (a rejected promise is still truthy, so the
+          // `!this.clientPromise` guard above would otherwise never re-enter).
+          this.clientPromise = undefined;
+          throw err;
+        });
+    }
+    return this.clientPromise;
+  }
+
   private async executeQuery(
     queryString: string,
     parameters: Record<string, any> = {},
   ): Promise<NeptuneQueryRecord[]> {
-    const response = await this.client.send(
-      new ExecuteQueryCommand({
+    const [client, sdk] = await Promise.all([this.getClient(), this.getSDK()]);
+    const response = await client.send(
+      new sdk.ExecuteQueryCommand({
         graphIdentifier: this.graphIdentifier,
         language: "OPEN_CYPHER",
         queryString,

mem0-ts/src/oss/tests/config-manager.test.ts

@@ -522,6 +522,34 @@ describe("ConfigManager", () => {
       expect(cfg.embedder.config).not.toHaveProperty("embedding_dims");
     });
   });
+
+  describe("mergeConfig - vector store provider normalization", () => {
+    const baseLlm = { provider: "openai", config: { apiKey: "test-key" } };
+    const baseEmbedder = { provider: "openai", config: { apiKey: "test-key" } };
+
+    it.each(["Memory", "DATABRICKS", "QdRaNt"])(
+      "lowercases the vector store provider %p",
+      (provider) => {
+        const config = ConfigManager.mergeConfig({
+          embedder: baseEmbedder,
+          vectorStore: { provider, config: { collectionName: "test" } },
+          llm: baseLlm,
+        });
+
+        expect(config.vectorStore.provider).toBe(provider.toLowerCase());
+      },
+    );
+
+    it("still falls back to the default provider when none is given", () => {
+      const config = ConfigManager.mergeConfig({
+        embedder: baseEmbedder,
+        vectorStore: { config: { collectionName: "test" } } as any,
+        llm: baseLlm,
+      });
+
+      expect(config.vectorStore.provider).toBe("memory");
+    });
+  });
 });
 
 // ─────────────────────────────────────────────────────────────────────────

mem0-ts/src/oss/tests/factory.unit.test.ts

@@ -183,6 +183,11 @@ jest.mock("../src/vector_stores/pgvector", () => ({
     .fn()
     .mockImplementation((config) => ({ type: "pgvector", config })),
 }));
+jest.mock("../src/vector_stores/databricks", () => ({
+  DatabricksVectorStore: jest
+    .fn()
+    .mockImplementation((config) => ({ type: "databricks", config })),
+}));
 jest.mock("../src/vector_stores/neptune_analytics", () => ({
   NeptuneAnalyticsVectorStore: jest.fn().mockImplementation((config) => ({
     type: "neptune-analytics",
@@ -346,6 +351,7 @@ describe("VectorStoreFactory", () => {
     ["vectorize"],
     ["azure-ai-search"],
     ["pgvector"],
+    ["databricks"],
     ["neptune"],
     ["neptune-analytics"],
     ["upstash_vector"],