feat(vector-store): Add Databricks Mosaic AI vector store support (#3325)

Makefile

@@ -13,7 +13,7 @@ install:
 install_all:
 	pip install ruff==0.6.9 groq together boto3 litellm ollama chromadb weaviate weaviate-client sentence_transformers vertexai \
 	            google-generativeai elasticsearch opensearch-py vecs "pinecone<7.0.0" pinecone-text faiss-cpu langchain-community \
-							upstash-vector azure-search-documents langchain-memgraph langchain-neo4j langchain-aws rank-bm25 pymochow pymongo psycopg kuzu
+							upstash-vector azure-search-documents langchain-memgraph langchain-neo4j langchain-aws rank-bm25 pymochow pymongo psycopg kuzu databricks-sdk
 
 # Format code with ruff
 format:

docs/changelog.mdx

@@ -7,6 +7,24 @@ mode: "wide"
 <Tabs>
 <Tab title="Python">
 
+
+**New Features:**
+- **Databricks Vector Search:** Added comprehensive support for Databricks Vector Search as a vector store provider
+  - Delta Sync Index integration with automatic synchronization from Delta tables
+  - Dual authentication support (Service Principal and Personal Access Token)
+  - Both STANDARD and STORAGE_OPTIMIZED endpoint types supported
+  - Self-managed and Databricks-computed embedding options
+  - Auto-creation of endpoints and indexes when not present
+  - Unity Catalog integration for secure data governance
+
+**Improvements:**
+- **Documentation:**
+  - Added detailed Databricks Vector Search configuration guide with authentication methods
+  - Updated vector databases overview to include Databricks option
+  - Enhanced vector store documentation with embedding configuration examples
+
+</Update>
+
 <Update label="2025-07-24" description="v0.1.115">
 
 **New Features & Updates:**

docs/components/vectordbs/dbs/databricks.mdx

@@ -0,0 +1,130 @@
+[Databricks Vector Search](https://docs.databricks.com/en/generative-ai/vector-search.html) is a serverless similarity search engine that allows you to store a vector representation of your data, including metadata, in a vector database. With Vector Search, you can create auto-updating vector search indexes from Delta tables managed by Unity Catalog and query them with a simple API to return the most similar vectors.
+
+### Usage
+
+```python
+import os
+from mem0 import Memory
+
+config = {
+    "vector_store": {
+        "provider": "databricks",
+        "config": {
+            "workspace_url": "https://your-workspace.databricks.com",
+            "access_token": "your-access-token",
+            "endpoint_name": "your-vector-search-endpoint",
+            "index_name": "catalog.schema.index_name",
+            "source_table_name": "catalog.schema.source_table",
+            "embedding_dimension": 1536
+        }
+    }
+}
+
+m = Memory.from_config(config)
+messages = [
+    {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
+    {"role": "assistant", "content": "How about thriller movies? They can be quite engaging."},
+    {"role": "user", "content": "I'm not a big fan of thriller movies but I love sci-fi movies."},
+    {"role": "assistant", "content": "Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future."}
+]
+m.add(messages, user_id="alice", metadata={"category": "movies"})
+```
+
+### Config
+
+Here are the parameters available for configuring Databricks Vector Search:
+
+| Parameter | Description | Default Value |
+| --- | --- | --- |
+| `workspace_url` | The URL of your Databricks workspace | **Required** |
+| `access_token` | Personal Access Token for authentication | `None` |
+| `service_principal_client_id` | Service principal client ID (alternative to access_token) | `None` |
+| `service_principal_client_secret` | Service principal client secret (required with client_id) | `None` |
+| `endpoint_name` | Name of the Vector Search endpoint | **Required** |
+| `index_name` | Name of the vector index (Unity Catalog format: catalog.schema.index) | **Required** |
+| `source_table_name` | Name of the source Delta table (Unity Catalog format: catalog.schema.table) | **Required** |
+| `embedding_dimension` | Dimension of self-managed embeddings | `1536` |
+| `embedding_source_column` | Column name for text when using Databricks-computed embeddings | `None` |
+| `embedding_model_endpoint_name` | Databricks serving endpoint for embeddings | `None` |
+| `embedding_vector_column` | Column name for self-managed embedding vectors | `embedding` |
+| `endpoint_type` | Type of endpoint (`STANDARD` or `STORAGE_OPTIMIZED`) | `STANDARD` |
+| `sync_computed_embeddings` | Whether to sync computed embeddings automatically | `True` |
+
+### Authentication
+
+Databricks Vector Search supports two authentication methods:
+
+#### Service Principal (Recommended for Production)
+```python
+config = {
+    "vector_store": {
+        "provider": "databricks",
+        "config": {
+            "workspace_url": "https://your-workspace.databricks.com",
+            "service_principal_client_id": "your-service-principal-id",
+            "service_principal_client_secret": "your-service-principal-secret",
+            "endpoint_name": "your-endpoint",
+            "index_name": "catalog.schema.index_name",
+            "source_table_name": "catalog.schema.source_table"
+        }
+    }
+}
+```
+
+#### Personal Access Token (for Development)
+```python
+config = {
+    "vector_store": {
+        "provider": "databricks",
+        "config": {
+            "workspace_url": "https://your-workspace.databricks.com",
+            "access_token": "your-personal-access-token",
+            "endpoint_name": "your-endpoint",
+            "index_name": "catalog.schema.index_name",
+            "source_table_name": "catalog.schema.source_table"
+        }
+    }
+}
+```
+
+### Embedding Options
+
+#### Self-Managed Embeddings (Default)
+Use your own embedding model and provide vectors directly:
+
+```python
+config = {
+    "vector_store": {
+        "provider": "databricks",
+        "config": {
+            # ... authentication config ...
+            "embedding_dimension": 768,  # Match your embedding model
+            "embedding_vector_column": "embedding"
+        }
+    }
+}
+```
+
+#### Databricks-Computed Embeddings
+Let Databricks compute embeddings from text using a serving endpoint:
+
+```python
+config = {
+    "vector_store": {
+        "provider": "databricks",
+        "config": {
+            # ... authentication config ...
+            "embedding_source_column": "text",
+            "embedding_model_endpoint_name": "e5-small-v2"
+        }
+    }
+}
+```
+
+### Important Notes
+
+- **Delta Sync Index**: This implementation uses Delta Sync Index, which automatically syncs with your source Delta table. Direct vector insertion/deletion/update operations will log warnings as they're not supported with Delta Sync.
+- **Unity Catalog**: Both the source table and index must be in Unity Catalog format (`catalog.schema.table_name`).
+- **Endpoint Auto-Creation**: If the specified endpoint doesn't exist, it will be created automatically.
+- **Index Auto-Creation**: If the specified index doesn't exist, it will be created automatically with the provided configuration.
+- **Filter Support**: Supports filtering by metadata fields, with different syntax for STANDARD vs STORAGE_OPTIMIZED endpoints.

docs/components/vectordbs/overview.mdx

@@ -23,6 +23,7 @@ See the list of supported vector databases below.
   <Card title="Pinecone" href="/components/vectordbs/dbs/pinecone"></Card>
   <Card title="MongoDB" href="/components/vectordbs/dbs/mongodb"></Card>
   <Card title="Azure" href="/components/vectordbs/dbs/azure"></Card>
+  <Card title="Databricks" href="/components/vectordbs/dbs/databricks"></Card>
   <Card title="Redis" href="/components/vectordbs/dbs/redis"></Card>
   <Card title="Elasticsearch" href="/components/vectordbs/dbs/elasticsearch"></Card>
   <Card title="OpenSearch" href="/components/vectordbs/dbs/opensearch"></Card>

mem0/configs/vector_stores/databricks.py

@@ -0,0 +1,63 @@
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+from databricks.sdk.service.vectorsearch import EndpointType, VectorIndexType, PipelineType
+
+
+class DatabricksConfig(BaseModel):
+    """Configuration for Databricks Vector Search vector store."""
+
+    workspace_url: str = Field(..., description="Databricks workspace URL")
+    access_token: Optional[str] = Field(None, description="Personal access token for authentication")
+    client_id: Optional[str] = Field(None, description="Databricks Service principal client ID")
+    client_secret: Optional[str] = Field(None, description="Databricks Service principal client secret")
+    azure_client_id: Optional[str] = Field(None, description="Azure AD application client ID (for Azure Databricks)")
+    azure_client_secret: Optional[str] = Field(
+        None, description="Azure AD application client secret (for Azure Databricks)"
+    )
+    endpoint_name: str = Field(..., description="Vector search endpoint name")
+    catalog: str = Field(..., description="The Unity Catalog catalog name")
+    schema: str = Field(..., description="The Unity Catalog schama name")
+    table_name: str = Field(..., description="Source Delta table name")
+    collection_name: str = Field("mem0", description="Vector search index name")
+    index_type: VectorIndexType = Field("DELTA_SYNC", description="Index type: DELTA_SYNC or DIRECT_ACCESS")
+    embedding_model_endpoint_name: Optional[str] = Field(
+        None, description="Embedding model endpoint for Databricks-computed embeddings"
+    )
+    embedding_dimension: int = Field(1536, description="Vector embedding dimensions")
+    endpoint_type: EndpointType = Field("STANDARD", description="Endpoint type: STANDARD or STORAGE_OPTIMIZED")
+    pipeline_type: PipelineType = Field("TRIGGERED", description="Sync pipeline type: TRIGGERED or CONTINUOUS")
+    warehouse_name: Optional[str] = Field(None, description="Databricks SQL warehouse Name")
+    query_type: str = Field("ANN", description="Query type: `ANN` and `HYBRID`")
+
+    @model_validator(mode="before")
+    @classmethod
+    def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]:
+        allowed_fields = set(cls.model_fields.keys())
+        input_fields = set(values.keys())
+        extra_fields = input_fields - allowed_fields
+        if extra_fields:
+            raise ValueError(
+                f"Extra fields not allowed: {', '.join(extra_fields)}. Please input only the following fields: {', '.join(allowed_fields)}"
+            )
+        return values
+
+    @model_validator(mode="after")
+    def validate_authentication(self):
+        """Validate that either access_token or service principal credentials are provided."""
+        has_token = self.access_token is not None
+        has_service_principal = (self.client_id is not None and self.client_secret is not None) or (
+            self.azure_client_id is not None and self.azure_client_secret is not None
+        )
+
+        if not has_token and not has_service_principal:
+            raise ValueError(
+                "Either access_token or both client_id/client_secret or azure_client_id/azure_client_secret must be provided"
+            )
+
+        return self
+
+    model_config = {
+        "arbitrary_types_allowed": True,
+    }

mem0/utils/factory.py

@@ -165,6 +165,7 @@ class VectorStoreFactory:
         "pinecone": "mem0.vector_stores.pinecone.PineconeDB",
         "mongodb": "mem0.vector_stores.mongodb.MongoDB",
         "redis": "mem0.vector_stores.redis.RedisDB",
+        "databricks": "mem0.vector_stores.databricks.Databricks",
         "elasticsearch": "mem0.vector_stores.elasticsearch.ElasticsearchDB",
         "vertex_ai_vector_search": "mem0.vector_stores.vertex_ai_vector_search.GoogleMatchingEngine",
         "opensearch": "mem0.vector_stores.opensearch.OpenSearchDB",

mem0/vector_stores/configs.py

@@ -21,6 +21,7 @@ class VectorStoreConfig(BaseModel):
         "upstash_vector": "UpstashVectorConfig",
         "azure_ai_search": "AzureAISearchConfig",
         "redis": "RedisDBConfig",
+        "databricks": "DatabricksConfig",
         "elasticsearch": "ElasticsearchConfig",
         "vertex_ai_vector_search": "GoogleMatchingEngineConfig",
         "opensearch": "OpenSearchConfig",

mem0/vector_stores/databricks.py

@@ -0,0 +1,759 @@
+import json
+import logging
+import uuid
+from typing import Optional, List
+from datetime import datetime, date
+from databricks.sdk.service.catalog import ColumnInfo, ColumnTypeName, TableType, DataSourceFormat
+from databricks.sdk.service.catalog import TableConstraint, PrimaryKeyConstraint
+from databricks.sdk import WorkspaceClient
+from databricks.sdk.service.vectorsearch import (
+    VectorIndexType,
+    DeltaSyncVectorIndexSpecRequest,
+    DirectAccessVectorIndexSpec,
+    EmbeddingSourceColumn,
+    EmbeddingVectorColumn,
+)
+from pydantic import BaseModel
+from mem0.memory.utils import extract_json
+from mem0.vector_stores.base import VectorStoreBase
+
+logger = logging.getLogger(__name__)
+
+
+class MemoryResult(BaseModel):
+    id: Optional[str] = None
+    score: Optional[float] = None
+    payload: Optional[dict] = None
+
+
+excluded_keys = {"user_id", "agent_id", "run_id", "hash", "data", "created_at", "updated_at"}
+
+
+class Databricks(VectorStoreBase):
+    def __init__(
+        self,
+        workspace_url: str,
+        access_token: Optional[str] = None,
+        client_id: Optional[str] = None,
+        client_secret: Optional[str] = None,
+        azure_client_id: Optional[str] = None,
+        azure_client_secret: Optional[str] = None,
+        endpoint_name: str = None,
+        catalog: str = None,
+        schema: str = None,
+        table_name: str = None,
+        collection_name: str = "mem0",
+        index_type: str = "DELTA_SYNC",
+        embedding_model_endpoint_name: Optional[str] = None,
+        embedding_dimension: int = 1536,
+        endpoint_type: str = "STANDARD",
+        pipeline_type: str = "TRIGGERED",
+        warehouse_name: Optional[str] = None,
+        query_type: str = "ANN",
+    ):
+        """
+        Initialize the Databricks Vector Search vector store.
+
+        Args:
+            workspace_url (str): Databricks workspace URL.
+            access_token (str, optional): Personal access token for authentication.
+            client_id (str, optional): Service principal client ID for authentication.
+            client_secret (str, optional): Service principal client secret for authentication.
+            azure_client_id (str, optional): Azure AD application client ID (for Azure Databricks).
+            azure_client_secret (str, optional): Azure AD application client secret (for Azure Databricks).
+            endpoint_name (str): Vector search endpoint name.
+            catalog (str): Unity Catalog catalog name.
+            schema (str): Unity Catalog schema name.
+            table_name (str): Source Delta table name.
+            index_name (str, optional): Vector search index name (default: "mem0").
+            index_type (str, optional): Index type, either "DELTA_SYNC" or "DIRECT_ACCESS" (default: "DELTA_SYNC").
+            embedding_model_endpoint_name (str, optional): Embedding model endpoint for Databricks-computed embeddings.
+            embedding_dimension (int, optional): Vector embedding dimensions (default: 1536).
+            endpoint_type (str, optional): Endpoint type, either "STANDARD" or "STORAGE_OPTIMIZED" (default: "STANDARD").
+            pipeline_type (str, optional): Sync pipeline type, either "TRIGGERED" or "CONTINUOUS" (default: "TRIGGERED").
+            warehouse_name (str, optional): Databricks SQL warehouse Name (if using SQL warehouse).
+            query_type (str, optional): Query type, either "ANN" or "HYBRID" (default: "ANN").
+        """
+        # Basic identifiers
+        self.workspace_url = workspace_url
+        self.endpoint_name = endpoint_name
+        self.catalog = catalog
+        self.schema = schema
+        self.table_name = table_name
+        self.fully_qualified_table_name = f"{self.catalog}.{self.schema}.{self.table_name}"
+        self.index_name = collection_name
+        self.fully_qualified_index_name = f"{self.catalog}.{self.schema}.{self.index_name}"
+
+        # Configuration
+        self.index_type = index_type
+        self.embedding_model_endpoint_name = embedding_model_endpoint_name
+        self.embedding_dimension = embedding_dimension
+        self.endpoint_type = endpoint_type
+        self.pipeline_type = pipeline_type
+        self.query_type = query_type
+
+        # Schema
+        self.columns = [
+            ColumnInfo(
+                name="memory_id",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                nullable=False,
+                comment="Primary key",
+                position=0,
+            ),
+            ColumnInfo(
+                name="hash",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                comment="Hash of the memory content",
+                position=1,
+            ),
+            ColumnInfo(
+                name="agent_id",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                comment="ID of the agent",
+                position=2,
+            ),
+            ColumnInfo(
+                name="run_id",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                comment="ID of the run",
+                position=3,
+            ),
+            ColumnInfo(
+                name="user_id",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                comment="ID of the user",
+                position=4,
+            ),
+            ColumnInfo(
+                name="memory",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                comment="Memory content",
+                position=5,
+            ),
+            ColumnInfo(
+                name="metadata",
+                type_name=ColumnTypeName.STRING,
+                type_text="string",
+                type_json='{"type":"string"}',
+                comment="Additional metadata",
+                position=6,
+            ),
+            ColumnInfo(
+                name="created_at",
+                type_name=ColumnTypeName.TIMESTAMP,
+                type_text="timestamp",
+                type_json='{"type":"timestamp"}',
+                comment="Creation timestamp",
+                position=7,
+            ),
+            ColumnInfo(
+                name="updated_at",
+                type_name=ColumnTypeName.TIMESTAMP,
+                type_text="timestamp",
+                type_json='{"type":"timestamp"}',
+                comment="Last update timestamp",
+                position=8,
+            ),
+        ]
+        if self.index_type == VectorIndexType.DIRECT_ACCESS:
+            self.columns.append(
+                ColumnInfo(
+                    name="embedding",
+                    type_name=ColumnTypeName.ARRAY,
+                    type_text="array<float>",
+                    type_json='{"type":"array","element":"float","element_nullable":false}',
+                    nullable=True,
+                    comment="Embedding vector",
+                    position=9,
+                )
+            )
+        self.column_names = [col.name for col in self.columns]
+
+        # Initialize Databricks workspace client
+        client_config = {}
+        if client_id and client_secret:
+            client_config.update(
+                {
+                    "host": workspace_url,
+                    "client_id": client_id,
+                    "client_secret": client_secret,
+                }
+            )
+        elif azure_client_id and azure_client_secret:
+            client_config.update(
+                {
+                    "host": workspace_url,
+                    "azure_client_id": azure_client_id,
+                    "azure_client_secret": azure_client_secret,
+                }
+            )
+        elif access_token:
+            client_config.update({"host": workspace_url, "token": access_token})
+        else:
+            # Try automatic authentication
+            client_config["host"] = workspace_url
+
+        try:
+            self.client = WorkspaceClient(**client_config)
+            logger.info("Initialized Databricks workspace client")
+        except Exception as e:
+            logger.error(f"Failed to initialize Databricks workspace client: {e}")
+            raise
+
+        # Get the warehouse ID by name
+        self.warehouse_id = next((w.id for w in self.client.warehouses.list() if w.name == warehouse_name), None)
+
+        # Initialize endpoint (required in Databricks)
+        self._ensure_endpoint_exists()
+
+        # Check if index exists and create if needed
+        collections = self.list_cols()
+        if self.fully_qualified_index_name not in collections:
+            self.create_col()
+
+    def _ensure_endpoint_exists(self):
+        """Ensure the vector search endpoint exists, create if it doesn't."""
+        try:
+            self.client.vector_search_endpoints.get_endpoint(endpoint_name=self.endpoint_name)
+            logger.info(f"Vector search endpoint '{self.endpoint_name}' already exists")
+        except Exception:
+            # Endpoint doesn't exist, create it
+            try:
+                logger.info(f"Creating vector search endpoint '{self.endpoint_name}' with type '{self.endpoint_type}'")
+                self.client.vector_search_endpoints.create_endpoint_and_wait(
+                    name=self.endpoint_name, endpoint_type=self.endpoint_type
+                )
+                logger.info(f"Successfully created vector search endpoint '{self.endpoint_name}'")
+            except Exception as e:
+                logger.error(f"Failed to create vector search endpoint '{self.endpoint_name}': {e}")
+                raise
+
+    def _ensure_source_table_exists(self):
+        """Ensure the source Delta table exists with the proper schema."""
+        check = self.client.tables.exists(self.fully_qualified_table_name)
+
+        if check.table_exists:
+            logger.info(f"Source table '{self.fully_qualified_table_name}' already exists")
+        else:
+            logger.info(f"Source table '{self.fully_qualified_table_name}' does not exist, creating it...")
+            self.client.tables.create(
+                name=self.table_name,
+                catalog_name=self.catalog,
+                schema_name=self.schema,
+                table_type=TableType.MANAGED,
+                data_source_format=DataSourceFormat.DELTA,
+                storage_location=None,  # Use default storage location
+                columns=self.columns,
+                properties={"delta.enableChangeDataFeed": "true"},
+            )
+            logger.info(f"Successfully created source table '{self.fully_qualified_table_name}'")
+            self.client.table_constraints.create(
+                full_name_arg="logistics_dev.ai.dev_memory",
+                constraint=TableConstraint(
+                    primary_key_constraint=PrimaryKeyConstraint(
+                        name="pk_dev_memory",  # Name of the primary key constraint
+                        child_columns=["memory_id"],  # Columns that make up the primary key
+                    )
+                ),
+            )
+            logger.info(
+                f"Successfully created primary key constraint on 'memory_id' for table '{self.fully_qualified_table_name}'"
+            )
+
+    def create_col(self, name=None, vector_size=None, distance=None):
+        """
+        Create a new collection (index).
+
+        Args:
+            name (str, optional): Index name. If provided, will create a new index using the provided source_table_name.
+            vector_size (int, optional): Vector dimension size.
+            distance (str, optional): Distance metric (not directly applicable for Databricks).
+
+        Returns:
+            The index object.
+        """
+        # Determine index configuration
+        embedding_dims = vector_size or self.embedding_dimension
+        embedding_source_columns = [
+            EmbeddingSourceColumn(
+                name="memory",
+                embedding_model_endpoint_name=self.embedding_model_endpoint_name,
+            )
+        ]
+
+        logger.info(f"Creating vector search index '{self.fully_qualified_index_name}'")
+
+        # First, ensure the source Delta table exists
+        self._ensure_source_table_exists()
+
+        if self.index_type not in [VectorIndexType.DELTA_SYNC, VectorIndexType.DIRECT_ACCESS]:
+            raise ValueError("index_type must be either 'DELTA_SYNC' or 'DIRECT_ACCESS'")
+
+        try:
+            if self.index_type == VectorIndexType.DELTA_SYNC:
+                index = self.client.vector_search_indexes.create_index(
+                    name=self.fully_qualified_index_name,
+                    endpoint_name=self.endpoint_name,
+                    primary_key="memory_id",
+                    index_type=self.index_type,
+                    delta_sync_index_spec=DeltaSyncVectorIndexSpecRequest(
+                        source_table=self.fully_qualified_table_name,
+                        pipeline_type=self.pipeline_type,
+                        columns_to_sync=self.column_names,
+                        embedding_source_columns=embedding_source_columns,
+                    ),
+                )
+                logger.info(
+                    f"Successfully created vector search index '{self.fully_qualified_index_name}' with DELTA_SYNC type"
+                )
+                return index
+
+            elif self.index_type == VectorIndexType.DIRECT_ACCESS:
+                index = self.client.vector_search_indexes.create_index(
+                    name=self.fully_qualified_index_name,
+                    endpoint_name=self.endpoint_name,
+                    primary_key="memory_id",
+                    index_type=self.index_type,
+                    direct_access_index_spec=DirectAccessVectorIndexSpec(
+                        embedding_source_columns=embedding_source_columns,
+                        embedding_vector_columns=[
+                            EmbeddingVectorColumn(name="embedding", embedding_dimension=embedding_dims)
+                        ],
+                    ),
+                )
+                logger.info(
+                    f"Successfully created vector search index '{self.fully_qualified_index_name}' with DIRECT_ACCESS type"
+                )
+                return index
+        except Exception as e:
+            logger.error(f"Error making index_type: {self.index_type} for index {self.fully_qualified_index_name}: {e}")
+
+    def _format_sql_value(self, v):
+        """
+        Format a Python value into a safe SQL literal for Databricks.
+        """
+        if v is None:
+            return "NULL"
+        if isinstance(v, bool):
+            return "TRUE" if v else "FALSE"
+        if isinstance(v, (int, float)):
+            return str(v)
+        if isinstance(v, (datetime, date)):
+            return f"'{v.isoformat()}'"
+        if isinstance(v, list):
+            # Render arrays (assume numeric or string elements)
+            elems = []
+            for x in v:
+                if x is None:
+                    elems.append("NULL")
+                elif isinstance(x, (int, float)):
+                    elems.append(str(x))
+                else:
+                    s = str(x).replace("'", "''")
+                    elems.append(f"'{s}'")
+            return f"array({', '.join(elems)})"
+        if isinstance(v, dict):
+            try:
+                s = json.dumps(v)
+            except Exception:
+                s = str(v)
+            s = s.replace("'", "''")
+            return f"'{s}'"
+        # Fallback: treat as string
+        s = str(v).replace("'", "''")
+        return f"'{s}'"
+
+    def insert(self, vectors: list, payloads: list = None, ids: list = None):
+        """
+        Insert vectors into the index.
+
+        Args:
+            vectors (List[List[float]]): List of vectors to insert.
+            payloads (List[Dict], optional): List of payloads corresponding to vectors.
+            ids (List[str], optional): List of IDs corresponding to vectors.
+        """
+        # Determine the number of items to process
+        num_items = len(payloads) if payloads else len(vectors) if vectors else 0
+
+        value_tuples = []
+        for i in range(num_items):
+            values = []
+            for col in self.columns:
+                if col.name == "memory_id":
+                    val = ids[i] if ids and i < len(ids) else str(uuid.uuid4())
+                elif col.name == "embedding":
+                    val = vectors[i] if vectors and i < len(vectors) else []
+                elif col.name == "memory":
+                    val = payloads[i].get("data") if payloads and i < len(payloads) else None
+                else:
+                    val = payloads[i].get(col.name) if payloads and i < len(payloads) else None
+                values.append(val)
+            formatted = [self._format_sql_value(v) for v in values]
+            value_tuples.append(f"({', '.join(formatted)})")
+
+        insert_sql = f"INSERT INTO {self.fully_qualified_table_name} ({', '.join(self.column_names)}) VALUES {', '.join(value_tuples)}"
+
+        # Execute the insert
+        try:
+            response = self.client.statement_execution.execute_statement(
+                statement=insert_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
+            )
+            if response.status.state.value == "SUCCEEDED":
+                logger.info(
+                    f"Successfully inserted {num_items} items into Delta table {self.fully_qualified_table_name}"
+                )
+                return
+            else:
+                logger.error(f"Failed to insert items: {response.status.error}")
+                raise Exception(f"Insert operation failed: {response.status.error}")
+        except Exception as e:
+            logger.error(f"Insert operation failed: {e}")
+            raise
+
+    def search(self, query: str, vectors: list, limit: int = 5, filters: dict = None) -> List[MemoryResult]:
+        """
+        Search for similar vectors or text using the Databricks Vector Search index.
+
+        Args:
+            query (str): Search query text (for text-based search).
+            vectors (list): Query vector (for vector-based search).
+            limit (int): Maximum number of results.
+            filters (dict): Filters to apply.
+
+        Returns:
+            List of MemoryResult objects.
+        """
+        try:
+            filters_json = json.dumps(filters) if filters else None
+
+            # Choose query type
+            if self.index_type == VectorIndexType.DELTA_SYNC and query:
+                # Text-based search
+                sdk_results = self.client.vector_search_indexes.query_index(
+                    index_name=self.fully_qualified_index_name,
+                    columns=self.column_names,
+                    query_text=query,
+                    num_results=limit,
+                    query_type=self.query_type,
+                    filters_json=filters_json,
+                )
+            elif self.index_type == VectorIndexType.DIRECT_ACCESS and vectors:
+                # Vector-based search
+                sdk_results = self.client.vector_search_indexes.query_index(
+                    index_name=self.fully_qualified_index_name,
+                    columns=self.column_names,
+                    query_vector=vectors,
+                    num_results=limit,
+                    query_type=self.query_type,
+                    filters_json=filters_json,
+                )
+            else:
+                raise ValueError("Must provide query text for DELTA_SYNC or vectors for DIRECT_ACCESS.")
+
+            # Parse results
+            result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
+            data_array = result_data.data_array if hasattr(result_data, "data_array") else []
+
+            memory_results = []
+            for row in data_array:
+                # Map columns to values
+                row_dict = dict(zip(self.column_names, row)) if isinstance(row, (list, tuple)) else row
+                score = row_dict.get("score") or (
+                    row[-1] if isinstance(row, (list, tuple)) and len(row) > len(self.column_names) else None
+                )
+                payload = {k: row_dict.get(k) for k in self.column_names}
+                payload["data"] = payload.get("memory", "")
+                memory_id = row_dict.get("memory_id") or row_dict.get("id")
+                memory_results.append(MemoryResult(id=memory_id, score=score, payload=payload))
+            return memory_results
+
+        except Exception as e:
+            logger.error(f"Search failed: {e}")
+            raise
+
+    def delete(self, vector_id):
+        """
+        Delete a vector by ID from the Delta table.
+
+        Args:
+            vector_id (str): ID of the vector to delete.
+        """
+        try:
+            logger.info(f"Deleting vector with ID {vector_id} from Delta table {self.fully_qualified_table_name}")
+
+            delete_sql = f"DELETE FROM {self.fully_qualified_table_name} WHERE memory_id = '{vector_id}'"
+
+            response = self.client.statement_execution.execute_statement(
+                statement=delete_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
+            )
+
+            if response.status.state.value == "SUCCEEDED":
+                logger.info(f"Successfully deleted vector with ID {vector_id}")
+            else:
+                logger.error(f"Failed to delete vector with ID {vector_id}: {response.status.error}")
+
+        except Exception as e:
+            logger.error(f"Delete operation failed for vector ID {vector_id}: {e}")
+            raise
+
+    def update(self, vector_id=None, vector=None, payload=None):
+        """
+        Update a vector and its payload in the Delta table.
+
+        Args:
+            vector_id (str): ID of the vector to update.
+            vector (list, optional): New vector values.
+            payload (dict, optional): New payload data.
+        """
+
+        update_sql = f"UPDATE {self.fully_qualified_table_name} SET "
+        set_clauses = []
+        if not vector_id:
+            logger.error("vector_id is required for update operation")
+            return
+        if vector is not None:
+            if not isinstance(vector, list):
+                logger.error("vector must be a list of float values")
+                return
+            set_clauses.append(f"embedding = {vector}")
+        if payload:
+            if not isinstance(payload, dict):
+                logger.error("payload must be a dictionary")
+                return
+            for key, value in payload.items():
+                if key not in excluded_keys:
+                    set_clauses.append(f"{key} = '{value}'")
+
+        if not set_clauses:
+            logger.error("No fields to update")
+            return
+        update_sql += ", ".join(set_clauses)
+        update_sql += f" WHERE memory_id = '{vector_id}'"
+        try:
+            logger.info(f"Updating vector with ID {vector_id} in Delta table {self.fully_qualified_table_name}")
+
+            response = self.client.statement_execution.execute_statement(
+                statement=update_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
+            )
+
+            if response.status.state.value == "SUCCEEDED":
+                logger.info(f"Successfully updated vector with ID {vector_id}")
+            else:
+                logger.error(f"Failed to update vector with ID {vector_id}: {response.status.error}")
+        except Exception as e:
+            logger.error(f"Update operation failed for vector ID {vector_id}: {e}")
+            raise
+
+    def get(self, vector_id) -> MemoryResult:
+        """
+        Retrieve a vector by ID.
+
+        Args:
+            vector_id (str): ID of the vector to retrieve.
+
+        Returns:
+            MemoryResult: The retrieved vector.
+        """
+        try:
+            # Use query with ID filter to retrieve the specific vector
+            filters = {"memory_id": vector_id}
+            filters_json = json.dumps(filters)
+
+            results = self.client.vector_search_indexes.query_index(
+                index_name=self.fully_qualified_index_name,
+                columns=self.column_names,
+                query_text=" ",  # Empty query, rely on filters
+                num_results=1,
+                query_type=self.query_type,
+                filters_json=filters_json,
+            )
+
+            # Process results
+            result_data = results.result if hasattr(results, "result") else results
+            data_array = result_data.data_array if hasattr(result_data, "data_array") else []
+
+            if not data_array:
+                raise KeyError(f"Vector with ID {vector_id} not found")
+
+            result = data_array[0]
+            row_data = result if isinstance(result, dict) else result.__dict__
+
+            # Build payload following the standard schema
+            payload = {
+                "hash": row_data.get("hash", "unknown"),
+                "data": row_data.get("memory", row_data.get("data", "unknown")),
+                "created_at": row_data.get("created_at"),
+            }
+
+            # Add updated_at if available
+            if "updated_at" in row_data:
+                payload["updated_at"] = row_data.get("updated_at")
+
+            # Add optional fields
+            for field in ["agent_id", "run_id", "user_id"]:
+                if field in row_data:
+                    payload[field] = row_data[field]
+
+            # Add metadata
+            if "metadata" in row_data:
+                try:
+                    metadata = json.loads(extract_json(row_data["metadata"]))
+                    payload.update(metadata)
+                except (json.JSONDecodeError, TypeError):
+                    logger.warning(f"Failed to parse metadata: {row_data.get('metadata')}")
+
+            memory_id = row_data.get("memory_id", row_data.get("memory_id", vector_id))
+            return MemoryResult(id=memory_id, payload=payload)
+
+        except Exception as e:
+            logger.error(f"Failed to get vector with ID {vector_id}: {e}")
+            raise
+
+    def list_cols(self) -> List[str]:
+        """
+        List all collections (indexes).
+
+        Returns:
+            List of index names.
+        """
+        try:
+            indexes = self.client.vector_search_indexes.list_indexes(endpoint_name=self.endpoint_name)
+            return [idx.name for idx in indexes]
+        except Exception as e:
+            logger.error(f"Failed to list collections: {e}")
+            raise
+
+    def delete_col(self):
+        """
+        Delete the current collection (index).
+        """
+        try:
+            # Try fully qualified first
+            try:
+                self.client.vector_search_indexes.delete_index(index_name=self.fully_qualified_index_name)
+                logger.info(f"Successfully deleted index '{self.fully_qualified_index_name}'")
+            except Exception:
+                self.client.vector_search_indexes.delete_index(index_name=self.index_name)
+                logger.info(f"Successfully deleted index '{self.index_name}' (short name)")
+        except Exception as e:
+            logger.error(f"Failed to delete index '{self.index_name}': {e}")
+            raise
+
+    def col_info(self, name=None):
+        """
+        Get information about a collection (index).
+
+        Args:
+            name (str, optional): Index name. Defaults to current index.
+
+        Returns:
+            Dict: Index information.
+        """
+        try:
+            index_name = name or self.index_name
+            index = self.client.vector_search_indexes.get_index(index_name=index_name)
+            return {"name": index.name, "fields": self.columns}
+        except Exception as e:
+            logger.error(f"Failed to get info for index '{name or self.index_name}': {e}")
+            raise
+
+    def list(self, filters: dict = None, limit: int = None) -> list[MemoryResult]:
+        """
+        List all recent created memories from the vector store.
+
+        Args:
+            filters (dict, optional): Filters to apply.
+            limit (int, optional): Maximum number of results.
+
+        Returns:
+            List containing list of MemoryResult objects.
+        """
+        try:
+            filters_json = json.dumps(filters) if filters else None
+            num_results = limit or 100
+            columns = self.column_names
+            sdk_results = self.client.vector_search_indexes.query_index(
+                index_name=self.fully_qualified_index_name,
+                columns=columns,
+                query_text=" ",
+                num_results=num_results,
+                query_type=self.query_type,
+                filters_json=filters_json,
+            )
+            result_data = sdk_results.result if hasattr(sdk_results, "result") else sdk_results
+            data_array = result_data.data_array if hasattr(result_data, "data_array") else []
+
+            memory_results = []
+            for row in data_array:
+                row_dict = dict(zip(columns, row)) if isinstance(row, (list, tuple)) else row
+                payload = {k: row_dict.get(k) for k in columns}
+                # Parse metadata if present
+                if "metadata" in payload and payload["metadata"]:
+                    try:
+                        payload.update(json.loads(payload["metadata"]))
+                    except Exception:
+                        pass
+                memory_id = row_dict.get("memory_id") or row_dict.get("id")
+                memory_results.append(MemoryResult(id=memory_id, payload=payload))
+            return memory_results
+        except Exception as e:
+            logger.error(f"Failed to list memories: {e}")
+            return []
+
+    def reset(self):
+        """Reset the vector search index and underlying source table.
+
+        This will attempt to delete the existing index (both fully qualified and short name forms
+        for robustness), drop the backing Delta table, recreate the table with the expected schema,
+        and finally recreate the index. Use with caution as all existing data will be removed.
+        """
+        fq_index = self.fully_qualified_index_name
+        logger.warning(f"Resetting Databricks vector search index '{fq_index}'...")
+        try:
+            # Try deleting via fully qualified name first
+            try:
+                self.client.vector_search_indexes.delete_index(index_name=fq_index)
+                logger.info(f"Deleted index '{fq_index}'")
+            except Exception as e_fq:
+                logger.debug(f"Failed deleting fully qualified index name '{fq_index}': {e_fq}. Trying short name...")
+                try:
+                    # Fallback to existing helper which may use short name
+                    self.delete_col()
+                except Exception as e_short:
+                    logger.debug(f"Failed deleting short index name '{self.index_name}': {e_short}")
+
+            # Drop the backing table (if it exists)
+            try:
+                drop_sql = f"DROP TABLE IF EXISTS {self.fully_qualified_table_name}"
+                resp = self.client.statement_execution.execute_statement(
+                    statement=drop_sql, warehouse_id=self.warehouse_id, wait_timeout="30s"
+                )
+                if getattr(resp.status, "state", None) == "SUCCEEDED":
+                    logger.info(f"Dropped table '{self.fully_qualified_table_name}'")
+                else:
+                    logger.warning(
+                        f"Attempted to drop table '{self.fully_qualified_table_name}' but state was {getattr(resp.status, 'state', 'UNKNOWN')}: {getattr(resp.status, 'error', None)}"
+                    )
+            except Exception as e_drop:
+                logger.warning(f"Failed to drop table '{self.fully_qualified_table_name}': {e_drop}")
+
+            # Recreate table & index
+            self._ensure_source_table_exists()
+            self.create_col()
+            logger.info(f"Successfully reset index '{fq_index}'")
+        except Exception as e:
+            logger.error(f"Error resetting index '{fq_index}': {e}")
+            raise

pyproject.toml

@@ -41,6 +41,7 @@ vector_stores = [
     "psycopg>=3.2.8",
     "pymongo>=4.13.2",
     "pymochow>=2.2.9",
+    "databricks-sdk>=0.63.0",
 ]
 llms = [
     "groq>=0.3.0",

tests/vector_stores/test_databricks.py

@@ -0,0 +1,315 @@
+from types import SimpleNamespace
+from unittest.mock import MagicMock, patch
+from databricks.sdk.service.vectorsearch import VectorIndexType
+from mem0.vector_stores.databricks import Databricks
+import pytest
+
+
+# ---------------------- Fixtures ---------------------- #
+
+
+def _make_status(state="SUCCEEDED", error=None):
+    return SimpleNamespace(state=SimpleNamespace(value=state), error=error)
+
+
+def _make_exec_response(state="SUCCEEDED", error=None):
+    return SimpleNamespace(status=_make_status(state, error))
+
+
+@pytest.fixture
+def mock_workspace_client():
+    """Patch WorkspaceClient and provide a fully mocked client with required sub-clients."""
+    with patch("mem0.vector_stores.databricks.WorkspaceClient") as mock_wc_cls:
+        mock_wc = MagicMock(name="WorkspaceClient")
+
+        # warehouses.list -> iterable of objects with name/id
+        warehouse_obj = SimpleNamespace(name="test-warehouse", id="wh-123")
+        mock_wc.warehouses.list.return_value = [warehouse_obj]
+
+        # vector search endpoints
+        mock_wc.vector_search_endpoints.get_endpoint.side_effect = [Exception("not found"), MagicMock()]
+        mock_wc.vector_search_endpoints.create_endpoint_and_wait.return_value = None
+
+        # tables.exists
+        exists_obj = SimpleNamespace(table_exists=False)
+        mock_wc.tables.exists.return_value = exists_obj
+        mock_wc.tables.create.return_value = None
+        mock_wc.table_constraints.create.return_value = None
+
+        # vector_search_indexes list/create/query/delete
+        mock_wc.vector_search_indexes.list_indexes.return_value = []
+        mock_wc.vector_search_indexes.create_index.return_value = SimpleNamespace(name="catalog.schema.mem0")
+        mock_wc.vector_search_indexes.query_index.return_value = SimpleNamespace(result=SimpleNamespace(data_array=[]))
+        mock_wc.vector_search_indexes.delete_index.return_value = None
+        mock_wc.vector_search_indexes.get_index.return_value = SimpleNamespace(name="mem0")
+
+        # statement execution
+        mock_wc.statement_execution.execute_statement.return_value = _make_exec_response()
+
+        mock_wc_cls.return_value = mock_wc
+        yield mock_wc
+
+
+@pytest.fixture
+def db_instance_delta(mock_workspace_client):
+    return Databricks(
+        workspace_url="https://test",
+        access_token="tok",
+        endpoint_name="vs-endpoint",
+        catalog="catalog",
+        schema="schema",
+        table_name="table",
+        collection_name="mem0",
+        warehouse_name="test-warehouse",
+        index_type=VectorIndexType.DELTA_SYNC,
+        embedding_model_endpoint_name="embedding-endpoint",
+    )
+
+
+@pytest.fixture
+def db_instance_direct(mock_workspace_client):
+    # For DIRECT_ACCESS we want table exists path to skip creation; adjust mock first
+    mock_workspace_client.tables.exists.return_value = SimpleNamespace(table_exists=True)
+    return Databricks(
+        workspace_url="https://test",
+        access_token="tok",
+        endpoint_name="vs-endpoint",
+        catalog="catalog",
+        schema="schema",
+        table_name="table",
+        collection_name="mem0",
+        warehouse_name="test-warehouse",
+        index_type=VectorIndexType.DIRECT_ACCESS,
+        embedding_dimension=4,
+        embedding_model_endpoint_name="embedding-endpoint",
+    )
+
+
+# ---------------------- Initialization Tests ---------------------- #
+
+
+def test_initialization_delta_sync(db_instance_delta, mock_workspace_client):
+    # Endpoint ensure called (first attempt get_endpoint fails then create)
+    mock_workspace_client.vector_search_endpoints.create_endpoint_and_wait.assert_called_once()
+    # Table creation sequence
+    mock_workspace_client.tables.create.assert_called_once()
+    # Index created with expected args
+    assert (
+        mock_workspace_client.vector_search_indexes.create_index.call_args.kwargs["index_type"]
+        == VectorIndexType.DELTA_SYNC
+    )
+    assert mock_workspace_client.vector_search_indexes.create_index.call_args.kwargs["primary_key"] == "memory_id"
+
+
+def test_initialization_direct_access(db_instance_direct, mock_workspace_client):
+    # DIRECT_ACCESS should include embedding column
+    assert "embedding" in db_instance_direct.column_names
+    assert (
+        mock_workspace_client.vector_search_indexes.create_index.call_args.kwargs["index_type"]
+        == VectorIndexType.DIRECT_ACCESS
+    )
+
+
+def test_create_col_invalid_type(mock_workspace_client):
+    # Force invalid type by manually constructing and calling create_col after monkeypatching index_type
+    inst = Databricks(
+        workspace_url="https://test",
+        access_token="tok",
+        endpoint_name="vs-endpoint",
+        catalog="catalog",
+        schema="schema",
+        table_name="table",
+        collection_name="mem0",
+        warehouse_name="test-warehouse",
+        index_type=VectorIndexType.DELTA_SYNC,
+    )
+    inst.index_type = "BAD_TYPE"
+    with pytest.raises(ValueError):
+        inst.create_col()
+
+
+# ---------------------- Insert Tests ---------------------- #
+
+
+def test_insert_generates_sql(db_instance_direct, mock_workspace_client):
+    vectors = [[0.1, 0.2, 0.3, 0.4]]
+    payloads = [
+        {
+            "data": "hello world",
+            "user_id": "u1",
+            "agent_id": "a1",
+            "run_id": "r1",
+            "metadata": '{"topic":"greeting"}',
+            "hash": "h1",
+        }
+    ]
+    ids = ["id1"]
+    db_instance_direct.insert(vectors=vectors, payloads=payloads, ids=ids)
+    args, kwargs = mock_workspace_client.statement_execution.execute_statement.call_args
+    sql = kwargs["statement"] if "statement" in kwargs else args[0]
+    assert "INSERT INTO" in sql
+    assert "catalog.schema.table" in sql
+    assert "id1" in sql
+    # Embedding list rendered
+    assert "array(0.1, 0.2, 0.3, 0.4)" in sql
+
+
+# ---------------------- Search Tests ---------------------- #
+
+
+def test_search_delta_sync_text(db_instance_delta, mock_workspace_client):
+    # Simulate query results
+    row = [
+        "id1",
+        "hash1",
+        "agent1",
+        "run1",
+        "user1",
+        "memory text",
+        '{"topic":"greeting"}',
+        "2024-01-01T00:00:00",
+        "2024-01-01T00:00:00",
+        0.42,
+    ]
+    mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
+        result=SimpleNamespace(data_array=[row])
+    )
+    results = db_instance_delta.search(query="hello", vectors=None, limit=1)
+    mock_workspace_client.vector_search_indexes.query_index.assert_called_once()
+    assert len(results) == 1
+    assert results[0].id == "id1"
+    assert results[0].score == 0.42
+    assert results[0].payload["data"] == "memory text"
+
+
+def test_search_direct_access_vector(db_instance_direct, mock_workspace_client):
+    row = [
+        "id2",
+        "hash2",
+        "agent2",
+        "run2",
+        "user2",
+        "memory two",
+        '{"topic":"info"}',
+        "2024-01-02T00:00:00",
+        "2024-01-02T00:00:00",
+        [0.1, 0.2, 0.3, 0.4],
+        0.77,
+    ]
+    mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
+        result=SimpleNamespace(data_array=[row])
+    )
+    results = db_instance_direct.search(query="", vectors=[0.1, 0.2, 0.3, 0.4], limit=1)
+    assert len(results) == 1
+    assert results[0].id == "id2"
+    assert results[0].score == 0.77
+
+
+def test_search_missing_params_raises(db_instance_delta):
+    with pytest.raises(ValueError):
+        db_instance_delta.search(query="", vectors=[0.1, 0.2])  # DELTA_SYNC requires query text
+
+
+# ---------------------- Delete Tests ---------------------- #
+
+
+def test_delete_vector(db_instance_delta, mock_workspace_client):
+    db_instance_delta.delete("id-delete")
+    args, kwargs = mock_workspace_client.statement_execution.execute_statement.call_args
+    sql = kwargs.get("statement") or args[0]
+    assert "DELETE FROM" in sql and "id-delete" in sql
+
+
+# ---------------------- Update Tests ---------------------- #
+
+
+def test_update_vector(db_instance_direct, mock_workspace_client):
+    db_instance_direct.update(
+        vector_id="id-upd",
+        vector=[0.4, 0.5, 0.6, 0.7],
+        payload={"custom": "val", "user_id": "skip"},  # user_id should be excluded
+    )
+    args, kwargs = mock_workspace_client.statement_execution.execute_statement.call_args
+    sql = kwargs.get("statement") or args[0]
+    assert "UPDATE" in sql and "id-upd" in sql
+    assert "embedding = [0.4, 0.5, 0.6, 0.7]" in sql
+    assert "custom = 'val'" in sql
+    assert "user_id" not in sql  # excluded
+
+
+# ---------------------- Get Tests ---------------------- #
+
+
+def test_get_vector(db_instance_delta, mock_workspace_client):
+    mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
+        result=SimpleNamespace(
+            data_array=[
+                {
+                    "memory_id": "id-get",
+                    "hash": "h",
+                    "agent_id": "a",
+                    "run_id": "r",
+                    "user_id": "u",
+                    "memory": "some memory",
+                    "metadata": '{"tag":"x"}',
+                    "created_at": "2024-01-01T00:00:00",
+                    "updated_at": "2024-01-01T00:00:00",
+                    "score": 0.99,
+                }
+            ]
+        )
+    )
+    res = db_instance_delta.get("id-get")
+    assert res.id == "id-get"
+    assert res.payload["data"] == "some memory"
+    assert res.payload["tag"] == "x"
+
+
+# ---------------------- Collection Info / Listing Tests ---------------------- #
+
+
+def test_list_cols(db_instance_delta, mock_workspace_client):
+    mock_workspace_client.vector_search_indexes.list_indexes.return_value = [
+        SimpleNamespace(name="catalog.schema.mem0"),
+        SimpleNamespace(name="catalog.schema.other"),
+    ]
+    cols = db_instance_delta.list_cols()
+    assert "catalog.schema.mem0" in cols and "catalog.schema.other" in cols
+
+
+def test_col_info(db_instance_delta):
+    info = db_instance_delta.col_info()
+    assert info["name"] == "mem0"
+    assert any(col.name == "memory_id" for col in info["fields"])
+
+
+def test_list_memories(db_instance_delta, mock_workspace_client):
+    row = {
+        "memory_id": "id3",
+        "hash": "hash3",
+        "agent_id": "agent3",
+        "run_id": "run3",
+        "user_id": "user3",
+        "memory": "memory three",
+        "metadata": '{"topic":"misc"}',
+        "created_at": "2024-01-03T00:00:00",
+        "updated_at": "2024-01-03T00:00:00",
+        "score": 0.33,
+    }
+    mock_workspace_client.vector_search_indexes.query_index.return_value = SimpleNamespace(
+        result=SimpleNamespace(data_array=[row])
+    )
+    res = db_instance_delta.list(limit=1)
+    assert len(res) == 1
+    assert res[0].id == "id3"
+
+
+# ---------------------- Reset Tests ---------------------- #
+
+
+def test_reset(db_instance_delta, mock_workspace_client):
+    # Make delete raise to exercise fallback path then allow recreation
+    mock_workspace_client.vector_search_indexes.delete_index.side_effect = [Exception("fail fq"), None, None]
+    with patch.object(db_instance_delta, "create_col", wraps=db_instance_delta.create_col) as create_spy:
+        db_instance_delta.reset()
+        assert create_spy.called