"""ChromaDB 数据库层."""
import chromadb
from chromadb.api.models.Collection import Collection


class VectorDB:
    """向量数据库封装."""

    def __init__(self, persist_dir: str = "./data"):
        self.client = chromadb.PersistentClient(path=persist_dir)

    def get_or_create_collection(self, name: str) -> Collection:
        """获取或创建 collection."""
        return self.client.get_or_create_collection(name=name)

    def delete_collection(self, name: str) -> None:
        """删除 collection."""
        try:
            self.client.delete_collection(name=name)
        except ValueError:
            pass  # collection 不存在则忽略

    def close(self) -> None:
        """释放数据库连接."""
        self.client.close()