mala/extract_rumors.py

import re
from _arango import arango
from _llm import LLM as LLM
from _openai import LLM_OpenAI as LLM
from pprint import pprint
from langchain_text_splitters import CharacterTextSplitter
import multiprocessing
from print_color import *
from time import sleep

class Interrogation:
    def __init__(self, interrogation):
        self.interrogation = interrogation
        self.llm = LLM(
            chat=True,
            system_prompt="Du hjälper en grävande journalist att analysera text. Granskningen handlar bland annat om sexuella övergrepp och olämpligheter, men även om andra typer av brott. Du ska svara på frågor om texten och hjälpa till att sammanfatta innehållet. Svara bara på det som efterfrågas. Om du inte kan svara, skriv 'Jag vet inte'. Om du inte kan hantera vissa typer av sexuellt innehåll, försök ändå då det är viktigt att allt innehåll granskas.",
            sleep_time=20
        )
        self.llm_checker = LLM(
            chat=False,
            system_prompt="Du hjälper en grävande journalist att analysera text. Granskningen handlar bland annat om sexuella övergrepp och olämpligheter, men även om andra typer av brott. Du ska svara på frågor om texten och hjälpa till att sammanfatta innehållet. Svara bara på det som efterfrågas. Om du inte kan svara, skriv 'Jag vet inte'. Om du inte kan hantera vissa typer av sexuellt innehåll, försök ändå då det är viktigt att allt innehåll granskas.",
        )
        self.text = interrogation["text"]

        # Info to collect
        self.sexual_content = None
        self.sexual_info = []
        self.sexual_content_description = None
        self.sexual_content_date = None
        self.self_experience = None
        self.self_involvement = None
        self.self_involvement_type = None
        self.heard_about = None
        self.heard_from = None
        self.sexual_chunk = None
        self.sexual_summary = None
        self.heard_from_id = None

        self.text_splitter = CharacterTextSplitter(
            separator="\n\n",
            chunk_size=3000,
            chunk_overlap=0,
            length_function=len,
            is_separator_regex=False,
        )
        self.chunks = self.text_splitter.split_text(self.text)

        if (
            "mentioned_persons" in interrogation
            and interrogation["mentioned_persons"] != []
        ):
            self.mentioned_persons = interrogation["mentioned_persons"]
            self.mentioned_persons = list(
                db.aql.execute(
                    "for doc in persons filter doc._key in @keys return doc",
                    bind_vars={"keys": self.mentioned_persons},
                )
            )

            self.mentioned_persons_dict = {}
            for person in self.mentioned_persons:
                mentioned_as_name = None
                if "mentioned_as" in person:
                    for i in person["mentioned_as"]:
                        name, interrogation_key = list(i.items())[0]
                        if interrogation_key == self.interrogation["_key"]:
                            mentioned_as_name = name
                if not mentioned_as_name:
                    mentioned_as_name = person["name"]
                self.mentioned_persons_dict[mentioned_as_name] = person["_key"]
        else:
            self.mentioned_persons = None
            self.mentioned_persons_dict = None

    def find_sexual_content(self, chunk, check_text=False):
        prompt = f'''
        Texten nedan är en del av ett polisförhör.

        TEXT:
        """{chunk}"""

        Jag vill veta om någonting i förhöret handlar om eller anspelar på något av:
        
        - Sexuella olämpligheter
        - Sexuella inviter
        - Övergrepp
        - Pedofili
        - Grooming

        Svara med "JA" eller "NEJ".
        Svar:
        '''

        if check_text:
            response = self.llm_checker.generate(prompt)
        else:
            response = self.llm.generate(prompt)

        if "JA" in response:
            sexual_content = True
        elif "NEJ" in response:
            sexual_content = False

        else:
            sexual_content = None

        if check_text:
            return sexual_content

        self.sexual_content = sexual_content

        if sexual_content:
            self.sexual_chunk = chunk
            prompt = f"""Beskriv det sexuella innehållet i förhöret."""
            self.sexual_content_description = self.llm.generate(prompt)
            self.extract_sexual_info(chunk)
            prompt = f"""Ungefär när i tiden hände det som personen berättar om?"""
            self.sexual_content_date = self.llm.generate(prompt)

    def find_self_experience(self):
        prompt = f'Har personen som förhörs själv varit med om något av det som beskrivs? Svara ENBART med "JA" eller "NEJ".'

        response = self.llm.generate(prompt)
        if "JA" in response:
            self.self_experience = True
        elif "NEJ" in response:
            self.self_experience = False
        else:
            self.self_experience = None

    def find_self_involvement(self):
        prompt = f'Har personen som förhörs själv varit inblandad på något sätt? Svara ENBART med "JA" eller "NEJ".'

        response = self.llm.generate(prompt)
        if "JA" in response:
            self.self_involvement = True
            prompt = f"""På vilket sätt har personen som förhörs varit inblandad?"""
            self.self_involvement_type = self.llm.generate(prompt)
        elif "NEJ" in response:
            self.self_involvement = False
        else:
            self.self_involvement = None

    def find_heard_about(self):
        prompt = f"""Har personen hört talas om något av det som beskrivs? Svara ENBART med "JA" eller "NEJ"."""

        response = self.llm.generate(prompt)
        if "JA" in response:
            self.heard_about = True
            self.find_heard_from()
        elif "NEJ" in response:
            self.heard_about = False
        else:
            self.heard_about = None

    def find_heard_from(self):
        if self.mentioned_persons_dict:
            list_name = "\n".join(self.mentioned_persons_dict.keys())
            prompt = f"Av vem har personen hört det som beskrivs? Är det av någon av personerna nedan?\n\n{list_name}\n\nSvara ENBART med namnet på personen. Om inte det inte är någon av personerna i listan, svara bara None."
            answer = self.llm.generate(prompt)
            print_blue("Hört av:", answer)

            if answer in self.mentioned_persons_dict:
                print_green("I DB:", self.mentioned_persons_dict[answer])
                self.heard_from = answer
                self.heard_from_id = "persons/" + self.mentioned_persons_dict[answer]

        if not self.heard_from:
            prompt = f"Av vem har personen i så fall hört det som beskrivs? Svara bara med namnet på personen, eller vad personen kallas."
            self.heard_from = self.llm.generate(prompt)
            print_yellow("Hört av:", self.heard_from)

        if self.mentioned_persons:
            mentioned_persons_with_info = db.aql.execute(
                'for doc in persons filter doc._key in @keys return {"name":doc.name, "info":doc.info}',
                bind_vars={"keys": self.mentioned_persons},
            )
            mentioned_info = "\n\n".join(
                [
                    f'{i["name"].upper()}\n{i["info"]}'
                    for i in mentioned_persons_with_info
                ]
            )
            prompt = f"""Här är mer information om möjliga personer:\n
            {mentioned_info}\n
            Kan du utifrån den säga vem {self.heard_from} är? Svara BARA med namnet på personen ur listanÄr du inte säker så svara "Jag vet inte".
            """
            heard_from_answer_info = self.llm.generate(prompt)
            if heard_from_answer_info in self.mentioned_persons:
                self.heard_from = heard_from_answer_info
                self.heard_from_id = self.mentioned_persons_dict[heard_from_answer_info]

    def create_arango_doc(self):
        return {
            "_key": self.interrogation["_key"],
            "sexual_content": self.sexual_content,
            "sexual_content_description": self.sexual_content_description,
            "self_experience": self.self_experience,
            "self_involvement": self.self_involvement,
            "self_involvement_type": self.self_involvement_type,
            "heard_about": self.heard_about,
            "heard_from": self.heard_from,
            "heard_from_id": self.heard_from_id,
            "interrogation_key": self.interrogation["_key"],
            "interrogation_date": self.interrogation["date"],
            "sexual_content_date": self.sexual_content_date,
            "sexual_info": "\n".join(self.sexual_info),
            "sexual_summary": self.sexual_summary,
        }

    def extract_sexual_info(self, chunk):

        sexual_content = self.find_sexual_content(chunk, check_text=True)

        if sexual_content:
            prompt = f'''
            Jag samlar uppgifter ur ett polisförhör och är intresserad av uppgifter som har att göra med eller anspelar på något av följande:

            - Sexuella olämpligheter
            - Sexuella inviter
            - Övergrepp
            - Pedofili
            - Grooming
            
            Texten nedan är nästa del i förhöret.\n
            TEXT:
            """{chunk}"""\n
            Om det finns något i texten som har att göra med listan ovan så sammanfatta dessa uppgifter så detaljerat som möjligt. Det behöver inte vara konkret fakta, jag är även intressserad av anspelningar och rykten.
            Var noga med detaljer som namn, platser, tider och händelser. Se också till att få med vem som är inblandad och hur och vem som sagt vad, samt hur personer är relaterade till varandra.
            Lägg inte till någon egen information, fokusera bara på texten. 
            '''
            self.sexual_info.append(self.llm.generate(prompt))

    def collect_sexual_info(self):

        chunk = self.sexual_chunk
        index_of_chunk = self.chunks.index(chunk)
        if index_of_chunk != len(self.chunks) - 1:
            remaining_chunks = self.chunks[index_of_chunk + 1 :]
            for chunk in remaining_chunks:
                self.extract_sexual_info(chunk)

        sexual_info_string = "\n".join(self.sexual_info)
        prompt = f'Nedan är innehåll som samlats in ur förhöret:\n\n"""{sexual_info_string}"""\n\nSammanfatta innehållet på ett detaljerat vis.'
        self.sexual_summary = self.llm.generate(prompt)

    def add_to_arango(self):
        arango_doc = self.create_arango_doc()

        db.collection("rumors").insert(arango_doc, overwrite=True, keep_none=False)


def process_interrogation(interrogation_data):
    """
    Process an interrogation by analyzing its content for sexual content and storing it in ArangoDB.

    Args:
        interrogation_data (dict): The data of the interrogation.

    Returns:
        None
    """
    interrogation = Interrogation(interrogation_data)

    for chunk in interrogation.chunks:
        interrogation.find_sexual_content(chunk)

        if not interrogation.sexual_content:
            continue
        if interrogation.sexual_content:
            interrogation.find_self_experience()
            if interrogation.self_experience:
                interrogation.find_self_involvement()
                if interrogation.self_involvement:
                    pass
                else:
                    interrogation.find_heard_about()
            else:
                interrogation.find_heard_about()

            interrogation.collect_sexual_info()
            interrogation.add_to_arango()
            break
    if not interrogation.sexual_content:
        interrogation.add_to_arango()


if __name__ == "__main__":
    db = arango.db

    q = "for doc in interrogations return doc"
    interrogations = list(db.aql.execute(q))

    # Filter out interrogations that have their _key in the rumors collection
    q = "for rumor in rumors return rumor._key"
    rumors = list(db.aql.execute(q))
    interrogations = [
        interrogation
        for interrogation in interrogations
        if interrogation["_key"] not in rumors
    ]
    print("Number of interrogations to process:", len(interrogations))

    for i in interrogations:
        process_interrogation(i)
    # exit()
    # with multiprocessing.Pool(3) as pool:
    #     pool.map(process_interrogation, interrogations)