File size: 1,230 Bytes
6e1a53e |
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 |
import chromadb
from chromadb.config import Settings
from fastapi import APIRouter
router = APIRouter()
chroma_settings = Settings(
anonymized_telemetry=False,
persist_directory="db",
allow_reset=False,
is_persistent=True,
)
client = chromadb.Client(chroma_settings)
@router.get("/api/v1/admin/collections")
async def get_all_collections():
# Currently only works for ChromaDB but can be extended easily
# for other vector stores as well
collections = client.list_collections()
responses = [c.dict() for c in collections]
return responses
# TODO(deshraj): Add pagination and make this endpoint agnostic to the vector store
@router.get("/api/v1/admin/collections/chromadb/{collection_name}")
async def get_collection_details(collection_name: str):
collection = client.get_collection(collection_name)
collection_data = collection.get()
metadatas, documents = collection_data['metadatas'], collection_data['documents']
collated_data = []
for i in zip(metadatas, documents):
collated_data.append({
"metadata": i[0],
"document": i[1]
})
response = {"details": collection.model_dump(), "data": collated_data}
return response
|