From 37d631cadf25b7c778ea4731260c13eafae3e53f Mon Sep 17 00:00:00 2001
From: lasseedfast <>
Date: Fri, 3 May 2024 07:03:35 +0200
Subject: [PATCH] .

---
 identify_person.py | 52 ++++++++++++++++++++++++++--------------------
 1 file changed, 30 insertions(+), 22 deletions(-)

diff --git a/identify_person.py b/identify_person.py
index 3a4406c..75a60e5 100644
--- a/identify_person.py
+++ b/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:
-            - answer (str): The generated answer from the language model.
-            - person_in_arango (dict): The information of the person retrieved from ArangoDB.
-            - interrogation (str): The ID of the interrogation where the person was mentioned.
+        list: A list of tuples containing the following information:
+            - generated answer (str): The generated answer from the language model.
+            - person information (dict): Information about the matched person in the database.
+            - 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]
-    _key = hits['metadatas'][0][0]['_key']
+    found_person = hits['documents'][0][0]
+    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)
-    info = '\n'.join(person_in_arango['info'])
+    found_person_in_arango = db.collection('persons').get(found_person_key)
+    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:
-        """{info}"""\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.
+        På andra ställen i polisens förundersökning finns en person som heter "{found_person}", och som beskrivs så här:
+        """{person_in_arango_summary}"""\n
+        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):
-    found_input = input('Enter om det stämmer? >> ')
-    if found_input == '':
-        found = True
+def verify(answer, person, person_in_arango, text, db):
+    if answer == 'Yes':
         person['mentioned_in_interrogation'].remove(text)
         db.collection('persons').update(person)
-        person_in_arango['info'].append(answer)
-        person_in_arango['mentioned_in_interrogation'].append(text)
+        person_in_arango['info'] += person['info']
+        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())