.

identify_person.py

@@ -36,64 +36,72 @@ def find_person(person):
         person (dict): A dictionary containing information about the person.
 
     Returns:
-        tuple: A tuple containing the following elements:
+        list: A list of tuples containing the following information:
-            - answer (str): The generated answer from the language model.
+            - generated answer (str): The generated answer from the language model.
-            - person_in_arango (dict): The information of the person retrieved from ArangoDB.
+            - person information (dict): Information about the matched person in the database.
-            - interrogation (str): The ID of the interrogation where the person was mentioned.
+            - interrogation document (dict): The document containing the interrogation text.
+            - mentioned person name (str): The name of the person mentioned in the interrogation.
+            - matched person name (str): The name of the person matched in the database.
+            - original person information (dict): The original information about the person.
+
     """
     db = arango.db
     llm = LLM()
 
-    found = False
     other_person = person['name']
 
     chroma = ChromaDB()
     col = chroma.client.get_or_create_collection('mala_persons')
     hits = col.query(query_texts=[other_person], n_results=1)
 
-    hit = hits['documents'][0][0]
+    found_person = hits['documents'][0][0]
-    _key = hits['metadatas'][0][0]['_key']
+    found_person_key = hits['metadatas'][0][0]['_key']
     distance = hits['distances'][0][0]
     
     #* Filter out hits with distance > 1
     if distance > 1:
         return None
-    person_in_arango = db.collection('persons').get(_key)
+    found_person_in_arango = db.collection('persons').get(found_person_key)
-    info = '\n'.join(person_in_arango['info'])
+    found_person_info = '\n'.join(found_person_in_arango['info'])
+
+    prompt = f'Nedan är olika bitar med information om en person:\n\n{found_person_info}\n\nSammanfatta dessa på ett detaljerat sätt, var noga med namn, platser, händelser och relationer. Använd bara sånt som finns i informationen. Svara ENBART med sammanfattningen, ingenting annat. '
+    person_in_arango_summary = llm.generate(prompt)
+    # Write summary about the person
 
-    print('Other:', other_person, 'Förslag:', hit, 'Distance:', distance)
     interrogations = person['mentioned_in_interrogation']
 
     output = []
     for interrogation in interrogations:
-        text = db.collection('interrogations').get(interrogation)['text']
+        interrogation_doc = db.collection('interrogations').get(interrogation)
+        text = interrogation_doc['text']
+
         prompt = f'''I texten nedan omnämns en "{other_person}" och jag försöker förstå om det kan vara exempelvis ett felstavat namn eller smeknamn för en annan person.\n
         TEXT:
         """{text}"""\n
 
-        På andra ställen i polisens förundersökning finns en person som heter "{hit}", och som beskrivs så här:
+        På andra ställen i polisens förundersökning finns en person som heter "{found_person}", och som beskrivs så här:
-        """{info}"""\n
+        """{person_in_arango_summary}"""\n
-        Verkar det troligt att personen som kallas {other_person} är samma person som {hit}? Svara bara JA eller NEJ, samt en kort förklaring till varför.
+        Verkar det troligt att personen som kallas {other_person} är samma person som {found_person}? Svara bara JA eller NEJ, samt en kort förklaring till varför.
         '''
         #  Om istället förnamnet eller efternamnet är helt olika så är det förmodligen inte samma person.Om det bara är ett namn (inget efternamn) kan det också handla om ett smeknamn eller en beskrivning.
         answer = llm.generate(prompt)
 
-        output.append((answer, person_in_arango, text))
+        output.append((answer, found_person_in_arango, interrogation_doc, other_person, found_person, found_person_info, person))
+    
     return output
 
-def verify(answer, person, person_in_arango, text):
+def verify(answer, person, person_in_arango, text, db):
-    found_input = input('Enter om det stämmer? >> ')
+    if answer == 'Yes':
-    if found_input == '':
-        found = True
         person['mentioned_in_interrogation'].remove(text)
         db.collection('persons').update(person)
-        person_in_arango['info'].append(answer)
+        person_in_arango['info'] += person['info']
-        person_in_arango['mentioned_in_interrogation'].append(text)
+        person_in_arango['mentioned_in_interrogation'] += ['mentioned_in_interrogation']
         db.collection('persons').update(person_in_arango)
-    if found:
+
         db.collection('other_persons').insert(person, overwrite=True)
         db.collection('persons').delete(person, check_rev=False)
         
+
 if __name__ == '__main__':
     db = arango.db
     persons = list(db.collection('persons').all())