diff --git a/Makefile b/Makefile index 14edc7f2b..3c8287edd 100644 --- a/Makefile +++ b/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: diff --git a/docs/changelog.mdx b/docs/changelog.mdx index 28c3d066d..92607d77e 100644 --- a/docs/changelog.mdx +++ b/docs/changelog.mdx @@ -7,6 +7,24 @@ mode: "wide" + +**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 + + + **New Features & Updates:** diff --git a/docs/components/vectordbs/dbs/databricks.mdx b/docs/components/vectordbs/dbs/databricks.mdx new file mode 100644 index 000000000..add8ee517 --- /dev/null +++ b/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. diff --git a/docs/components/vectordbs/overview.mdx b/docs/components/vectordbs/overview.mdx index 332a0fdbd..922493eeb 100644 --- a/docs/components/vectordbs/overview.mdx +++ b/docs/components/vectordbs/overview.mdx @@ -23,6 +23,7 @@ See the list of supported vector databases below. + diff --git a/mem0/configs/vector_stores/databricks.py b/mem0/configs/vector_stores/databricks.py new file mode 100644 index 000000000..8729590a6 --- /dev/null +++ b/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, + } diff --git a/mem0/utils/factory.py b/mem0/utils/factory.py index aa4f44708..11fb14924 100644 --- a/mem0/utils/factory.py +++ b/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", diff --git a/mem0/vector_stores/configs.py b/mem0/vector_stores/configs.py index 7e03a6f83..51dcaa63c 100644 --- a/mem0/vector_stores/configs.py +++ b/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", diff --git a/mem0/vector_stores/databricks.py b/mem0/vector_stores/databricks.py new file mode 100644 index 000000000..0a5b06f49 --- /dev/null +++ b/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", + 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 diff --git a/pyproject.toml b/pyproject.toml index b7a49e008..944f8a682 100644 --- a/pyproject.toml +++ b/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", diff --git a/tests/vector_stores/test_databricks.py b/tests/vector_stores/test_databricks.py new file mode 100644 index 000000000..5b87def07 --- /dev/null +++ b/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