|
| 1 | +from typing import Any, Dict, List |
| 2 | + |
| 3 | +from pydantic import BaseModel, ConfigDict, Field, model_validator |
| 4 | + |
| 5 | + |
| 6 | +class TencentVectorDBConfig(BaseModel): |
| 7 | + url: str = Field(None, description="URL for TencentVectorDB instance") |
| 8 | + key: str = Field(None, description="API key for TencentVectorDB instance") |
| 9 | + username: str = Field("root", description="Username for TencentVectorDB instance") |
| 10 | + database_name: str = Field("mem0", description="Name of the database") |
| 11 | + collection_name: str = Field("mem0", description="Name of the collection") |
| 12 | + embedding_model_dims: int = Field(1536, description="Dimensions of the embedding model") |
| 13 | + metric_type: str = Field("COSINE", description="Metric type for similarity search") |
| 14 | + index_type: str = Field("HNSW", description="Index type for vectors") |
| 15 | + shard_num: int = Field(2, description="Number of shards in the collection") |
| 16 | + replica_num: int = Field(2, description="Number of replicas for the collection") |
| 17 | + field_type: str = Field("vector", description="Field type for the embedding vectors") |
| 18 | + params: Dict[str, Any] = Field({}, description="Parameters for the index") |
| 19 | + no_index_fields: List[str] = Field([], description="Fields that will not be indexed") |
| 20 | + |
| 21 | + @model_validator(mode="before") |
| 22 | + @classmethod |
| 23 | + def validate_extra_fields(cls, values: Dict[str, Any]) -> Dict[str, Any]: |
| 24 | + url, key = values.get("url"), values.get("key") |
| 25 | + if not url or not key: |
| 26 | + raise ValueError( |
| 27 | + "Both 'url' and 'key' must be provided." |
| 28 | + ) |
| 29 | + return values |
| 30 | + |
| 31 | + model_config = ConfigDict(arbitrary_types_allowed=True) |
0 commit comments