policy_analyser/extraction.py

"""
    Extraction
    @author : Sakshi Tantak
"""

# Imports
import os
import re
import json

from policy_analyser import EXTRACTION_PROMPT
from policy_analyser.llm import call_openai


def extract(text):
    raw_response = ''
    try:
        raw_response = call_openai(EXTRACTION_PROMPT + '\nPolicy Document : ' + text)
        print(raw_response)
        response = re.sub(r'`|json', '', raw_response)
        try:
            response = json.loads(response)
            response = [{'entityName' : k, 'entityValue' : v} for k, v in response.items()]
        except:
            response = []
    except Exception as e:
        print(e)
        response = []
    return raw_response, response