35 lines
1.4 KiB
Python
35 lines
1.4 KiB
Python
import os
|
|
from typing import Any, ClassVar, Dict, Optional
|
|
|
|
from pydantic import BaseModel, ConfigDict, Field, model_validator
|
|
|
|
try:
|
|
from upstash_vector import Index
|
|
except ImportError:
|
|
raise ImportError("The 'upstash_vector' library is required. Please install it using 'pip install upstash_vector'.")
|
|
|
|
|
|
class UpstashVectorConfig(BaseModel):
|
|
Index: ClassVar[type] = Index
|
|
|
|
url: Optional[str] = Field(None, description="URL for Upstash Vector index")
|
|
token: Optional[str] = Field(None, description="Token for Upstash Vector index")
|
|
client: Optional[Index] = Field(None, description="Existing `upstash_vector.Index` client instance")
|
|
collection_name: str = Field("mem0", description="Namespace to use for the index")
|
|
enable_embeddings: bool = Field(
|
|
False, description="Whether to use built-in upstash embeddings or not. Default is True."
|
|
)
|
|
|
|
@model_validator(mode="before")
|
|
@classmethod
|
|
def check_credentials_or_client(cls, values: Dict[str, Any]) -> Dict[str, Any]:
|
|
client = values.get("client")
|
|
url = values.get("url") or os.environ.get("UPSTASH_VECTOR_REST_URL")
|
|
token = values.get("token") or os.environ.get("UPSTASH_VECTOR_REST_TOKEN")
|
|
|
|
if not client and not (url and token):
|
|
raise ValueError("Either a client or URL and token must be provided.")
|
|
return values
|
|
|
|
model_config = ConfigDict(arbitrary_types_allowed=True)
|