configs.py•2.34 kB
from pydantic import BaseModel, Field, model_validator
class VectorStoreConfig(BaseModel):
provider: str = Field(
description="Provider of the vector store (e.g., 'qdrant', 'chroma', 'upstash_vector')",
default="qdrant",
)
config: dict | None = Field(
description="Configuration for the specific vector store", default=None
)
_provider_configs: dict[str, str] = {
"qdrant": "QdrantConfig",
"chroma": "ChromaDbConfig",
"pgvector": "PGVectorConfig",
"pinecone": "PineconeConfig",
"mongodb": "MongoDBConfig",
"milvus": "MilvusDBConfig",
"baidu": "BaiduDBConfig",
"neptune": "NeptuneAnalyticsConfig",
"upstash_vector": "UpstashVectorConfig",
"azure_ai_search": "AzureAISearchConfig",
"redis": "RedisDBConfig",
"valkey": "ValkeyConfig",
"databricks": "DatabricksConfig",
"elasticsearch": "ElasticsearchConfig",
"vertex_ai_vector_search": "GoogleMatchingEngineConfig",
"opensearch": "OpenSearchConfig",
"supabase": "SupabaseConfig",
"weaviate": "WeaviateConfig",
"faiss": "FAISSConfig",
"langchain": "LangchainConfig",
"s3_vectors": "S3VectorsConfig",
}
@model_validator(mode="after")
def validate_and_create_config(self) -> "VectorStoreConfig":
provider = self.provider
config = self.config
if provider not in self._provider_configs:
raise ValueError(f"Unsupported vector store provider: {provider}")
module = __import__(
f"selfmemory.configs.vector_stores.{provider}",
fromlist=[self._provider_configs[provider]],
)
config_class = getattr(module, self._provider_configs[provider])
if config is None:
config = {}
if not isinstance(config, dict):
if not isinstance(config, config_class):
raise ValueError(f"Invalid config type for provider {provider}")
return self
# also check if path in allowed kays for pydantic model, and whether config extra fields are allowed
if "path" not in config and "path" in config_class.__annotations__:
config["path"] = f"/tmp/{provider}"
self.config = config_class(**config)
return self