|
|
|
|
@ -38,14 +38,68 @@ class ChromaDB: |
|
|
|
|
for collection in collections: |
|
|
|
|
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 |
|
|
|
|
chroma = ChromaDB() |
|
|
|
|
|
|
|
|
|
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()) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|