retrieve_kb.py

from fastapi import APIRouter
from utils import get_chroma_client, get_embedding_function
import os
from dotenv import load_dotenv
import streamlit as st

load_dotenv()
openai_key = os.getenv("OPENAI_API_KEY")
router = APIRouter()
default_embedding_function = get_embedding_function(openai_key=openai_key)


def get_current_knowledge_bases(client):
    try:
        knowledge_boxes = client.list_collections()
        return knowledge_boxes
    except Exception as e:
        st.error(f"{str(e)}")


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)