File size: 1,076 Bytes
e04dd70
 
8e018ae
 
2c295e5
e04dd70
8e018ae
 
e04dd70
8e018ae
e04dd70
 
 
aadb9c7
 
 
5f95336
2c295e5
e04dd70
 
 
 
 
 
 
 
 
 
 
 
 
d258ef6
e04dd70
 
 
 
 
 
 
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
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)