from _llm import LLM
from _arango import db
from _chroma import chroma
from print_color import *
from identify_person import find_person


llm = LLM(small=True)


def check_from(relations):

    for relation in relations:
        interrogation = db.collection('interrogations').get(relation['interrogation'])
        text = f"Hörd person: {interrogation['person']}\n{interrogation['text']}"
        prompt = f"""Är "{relation['from']}" personen som förhörs i texten nedan?\n\n{text[:2000]}\n\nSvara enbart JA eller NEJ."""
        answer = llm.generate(prompt)
        if 'JA' in answer:
            relation['from_key'] = interrogation['person_id']
            db.collection('relations').update(relation, check_rev=False)
        print_rainbow(relation['from'], interrogation['person'], answer)


q = "for doc in relations filter doc.from_key == null limit 10 return doc" #! Limit 10
relations = list(db.aql.execute(q))

for relation in relations:
    desc = ''
    for r in relation['relations']:
        desc += r['description'] + '\n'
    desc = desc.strip()
    print_green(relation['to'])
    print(find_person(name=relation['to']))
    print()