Add methods to add persons to Chroma database

main
lasseedfast 2 years ago
parent b45077950b
commit aafcf6b0dd
  1. 60
      _chroma.py

@ -38,14 +38,68 @@ class ChromaDB:
for collection in collections: for collection in collections:
print(collection.name) print(collection.name)
def add_person_to_chroma(self, person):
"""
Adds a person to the Chroma database.
Args:
person (dict): A dictionary containing information about the person.
Returns:
None
"""
collection = self.client.get_collection("mala_persons")
# Lists to store the documents, metadatas and ids
documents = []
metadatas = []
ids = []
documents.append(person["name"])
metadata = {"name": person["name"], "_key": person["_key"], 'info': "\n".join(person["info"])}
metadatas.append(metadata)
ids.append(person["_key"])
collection.add(documents=documents, metadatas=metadatas, ids=ids)
def add_all_persons_to_chroma(self):
"""
Adds all persons to the Chroma collection.
This method deletes the existing 'mala_persons' collection, creates a new collection,
and then adds all persons from the database whose 'verified' field is set to True.
Args:
None
Returns:
None
"""
from _arango import arango
self.client.delete_collection('mala_persons')
col = self.client.get_or_create_collection('mala_persons')
db = arango.db
q = "for doc in persons filter doc.confirmed == true return doc"
persons = list(db.aql.execute(q))
for person in persons:
self.add_person_to_chroma(person)
print('Persons in chroma:', col.count())
# Initialize the ChromaDB object # Initialize the ChromaDB object
chroma = ChromaDB() chroma = ChromaDB()
if __name__ == '__main__': if __name__ == '__main__':
chroma = ChromaDB()
chroma.add_all_persons_to_chroma()
chroma.client.delete_collection('mala_persons')
col = chroma.client.get_or_create_collection('mala_persons')
print(col.count())

Loading…
Cancel
Save