from fastapi import APIRouter
from utils import get_chroma_client, get_embedding_function


router = APIRouter()
default_embedding_function = get_embedding_function()


def get_current_knowledge_bases(client):
    knowledge_boxes = client.list_collections()
    return knowledge_boxes


def get_knowledge_base_information(
    client, kb_name: str, embedding_function=default_embedding_function
):
    collection = client.get_collection(
        name=kb_name, embedding_function=embedding_function
    )

    collection_info = collection.get(
        include=["documents", "metadatas"]
    )  # you can add "embeddings", "metadatas",

    return collection_info, collection, client


if __name__ == "__main__":
    client = get_chroma_client()
    knowledge_boxes = get_current_knowledge_bases(client=client)
    for kb in knowledge_boxes:
        print(kb.name)