Support TiDB Vector Store

Makefile

@@ -13,7 +13,7 @@ install:
 install_all:
 	poetry install
 	poetry run pip install groq together boto3 litellm ollama chromadb weaviate weaviate-client sentence_transformers vertexai \
-	                        google-generativeai elasticsearch opensearch-py vecs
+	                        google-generativeai elasticsearch opensearch-py vecs pymysql
 
 # Format code with ruff
 format:

docs/components/vectordbs/dbs/tidb.mdx

@@ -0,0 +1,50 @@
+[TiDB](https://github.com/pingcap/tidb) is an open-source, MySQL-compatible database for modern GenAI application, which enables developers consolidate Vector Embeddings, Knowledge Graphs, and Operational Data into one database - with enterprise-grade reliability and simplicity.
+
+For more information, please check [ai.pingcap.com](https://ai.pingcap.com/).
+
+### Usage
+
+```python
+import os
+from mem0 import Memory
+
+os.environ["OPENAI_API_KEY"] = "sk-xx"
+
+config = {
+    "vector_store": {
+        "provider": "tidb",
+        "config": {
+            "user": "root",
+            "password": "",
+            "host": "localhost",
+            "port": "4000",
+        }
+    }
+}
+
+m = Memory.from_config(config)
+messages = [
+    {"role": "user", "content": "I'm planning to watch a movie tonight. Any recommendations?"},
+    {"role": "assistant", "content": "How about a 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's the parameters available for configuring TiDB:
+
+| Parameter | Description | Default Value |
+| --- | --- | --- |
+| `database` | The name of the  | `test` |
+| `collection_name` | The name of the collection | `mem0` |
+| `embedding_model_dims` | Dimensions of the embedding model | `1536` |
+| `user` | User name to connect to the database | `root` |
+| `password` | Password to connect to the database | |
+| `host` | The host where the TiDB is running | `lcoalhost` |
+| `port` | The port where the TiDB is running | `4000` |
+| `enable_ssl` | Enable SSL connection (for TiDB Cloud Serverless, need to turn it to `true`) | `false` |
+| `index_method` | Index method to use. | `HNSW` |
+| `distance_metric` | Index measure to use. | `COSINE` |

docs/components/vectordbs/overview.mdx

@@ -26,6 +26,7 @@ See the list of supported vector databases below.
   <Card title="Supabase" href="/components/vectordbs/dbs/supabase"></Card>
   <Card title="Vertex AI Vector Search" href="/components/vectordbs/dbs/vertex_ai_vector_search"></Card>
   <Card title="Weaviate" href="/components/vectordbs/dbs/weaviate"></Card>
+  <Card title="TiDB" href="/components/vectordbs/dbs/tidb"></Card>
 </CardGroup>
 
 ## Usage

mem0/configs/vector_stores/tidb.py

@@ -0,0 +1,62 @@
+from enum import Enum
+from typing import Any, Dict, Optional
+
+from pydantic import BaseModel, Field, model_validator
+
+
+class IndexMethod(str, Enum):
+    HNSW = "hnsw"
+
+
+class DistanceMetric(str, Enum):
+    """
+    An enumeration representing different types of distance metrics.
+
+    - `DistanceMetric.L2`: L2 (Euclidean) distance metric.
+    - `DistanceMetric.COSINE`: Cosine distance metric.
+    """
+
+    L2 = "L2"
+    COSINE = "COSINE"
+
+    def to_sql_func(self):
+        """
+        Converts the DistanceMetric to its corresponding SQL function name.
+
+        Returns:
+            str: The SQL function name.
+
+        Raises:
+            ValueError: If the DistanceMetric enum member is not supported.
+        """
+        if self == DistanceMetric.L2:
+            return "VEC_L2_DISTANCE"
+        elif self == DistanceMetric.COSINE:
+            return "VEC_COSINE_DISTANCE"
+        else:
+            raise ValueError("unsupported distance metric")
+
+
+class TiDBConfig(BaseModel):
+    database: str = Field("test", description="Default name for the database")
+    collection_name: str = Field("mem0", description="Default name for the collection")
+    embedding_model_dims: Optional[int] = Field(1536, description="Dimensions of the embedding model")
+    user: Optional[str] = Field("root", description="Database user. Default is root")
+    password: Optional[str] = Field("", description="Database password. Default is empty string")
+    host: Optional[str] = Field("localhost", description="Database host. Default is localhost")
+    port: Optional[int] = Field(4000, description="Database port. Default is 4000")
+    enable_ssl: Optional[bool] = Field(False, description="Enable SSL connection. Default is False")
+    index_method: Optional[IndexMethod] = Field(IndexMethod.HNSW, description="Index method to use. Default is HNSW")
+    distance_metric: Optional[DistanceMetric] = Field(DistanceMetric.COSINE, description="Distance metric to use. Default is COSINE")
+
+    @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

mem0/utils/factory.py

@@ -73,6 +73,7 @@ class VectorStoreFactory:
         "opensearch": "mem0.vector_stores.opensearch.OpenSearchDB",
         "supabase": "mem0.vector_stores.supabase.Supabase",
         "weaviate": "mem0.vector_stores.weaviate.Weaviate",
+        "tidb": "mem0.vector_stores.tidb.TiDB",
     }
 
     @classmethod

mem0/vector_stores/configs.py

@@ -22,6 +22,7 @@ class VectorStoreConfig(BaseModel):
         "opensearch": "OpenSearchConfig",
         "supabase": "SupabaseConfig",
         "weaviate": "WeaviateConfig",
+        "tidb": "TiDBConfig",
     }
 
     @model_validator(mode="after")