Spaces:
Sleeping
Sleeping
File size: 1,289 Bytes
25d4f70 f35d149 25d4f70 bd53034 25d4f70 2290123 bd53034 2290123 25d4f70 2290123 | 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 | from pathlib import Path
from chromadb import PersistentClient
import asyncio
class VectorStore:
def __init__(self):
base_dir = Path(__file__).resolve().parent.parent.parent
store_path = base_dir / "store"
self.client = PersistentClient(path=str(store_path))
def get_collection(self, name: str):
return self.client.get_or_create_collection(
name=name, metadata={"hnsw:space": "cosine"}
)
async def upsert(self, collection_name, ids, documents, embeddings, metadatas=None):
collection = self.get_collection(collection_name)
await asyncio.to_thread(
collection.upsert,
ids=ids,
documents=documents,
embeddings=embeddings,
metadatas=metadatas,
)
async def retrieve(self, collection_name, embeddings, where, n_results=3):
collection = self.get_collection(collection_name)
return await asyncio.to_thread(
collection.query,
where=where,
query_embeddings=embeddings,
n_results=n_results,
)
async def get_all_documents(self, collection_name: str):
collection = self.get_collection(collection_name)
return await asyncio.to_thread(collection.get)
|