Spaces:
Sleeping
Sleeping
File size: 2,728 Bytes
1c8d1ba | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | from __future__ import annotations
import logging
from collections.abc import Iterable
from typing import Any
from src.config import ALL_CATEGORIES_LABEL, ConfigurationError, Settings, load_settings
logger = logging.getLogger(__name__)
def build_category_filter(category: str | None) -> str | None:
if not category or category == ALL_CATEGORIES_LABEL:
return None
escaped = category.replace("\\", "\\\\").replace("'", "\\'")
return f"category = '{escaped}'"
class VectorStore:
def __init__(self, settings: Settings | None = None) -> None:
self.settings = settings or load_settings()
self.settings.require_upstash()
try:
from upstash_vector import Index
except ImportError as exc:
raise ConfigurationError(
"The upstash-vector package is not installed. Run `pip install -r requirements.txt`."
) from exc
assert self.settings.upstash_url is not None
assert self.settings.upstash_token is not None
self.namespace = self.settings.upstash_namespace
self.index = Index(url=self.settings.upstash_url, token=self.settings.upstash_token)
logger.info("Connected to Upstash Vector namespace %s", self.namespace)
def upsert_image_vector(self, id: str, vector: list[float], metadata: dict[str, Any]) -> None:
self.upsert_many([{"id": id, "vector": vector, "metadata": metadata}])
def upsert_many(self, items: Iterable[dict[str, Any]]) -> None:
vectors = [
{"id": item["id"], "vector": item["vector"], "metadata": item["metadata"]}
for item in items
]
if not vectors:
return
logger.info("Upserting %d vector(s) to namespace %s", len(vectors), self.namespace)
self.index.upsert(vectors=vectors, namespace=self.namespace)
def query_vector(
self,
vector: list[float],
top_k: int,
category: str | None = None,
) -> list[Any]:
query_filter = build_category_filter(category)
kwargs: dict[str, Any] = {
"vector": vector,
"top_k": int(top_k),
"include_metadata": True,
"include_vectors": False,
"namespace": self.namespace,
}
if query_filter:
kwargs["filter"] = query_filter
logger.info(
"Querying Upstash Vector with top_k=%s category=%s",
top_k,
category or ALL_CATEGORIES_LABEL,
)
return list(self.index.query(**kwargs))
def clear_namespace(self) -> None:
logger.warning("Clearing Upstash Vector namespace %s", self.namespace)
self.index.reset(namespace=self.namespace)
|