# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_core.ipynb.

# %% auto 0
__all__ = ['text2dict', 'text2df']

# %% ../nbs/00_core.ipynb 3
from abc import ABC, abstractmethod
import pandas as pd
import openai
import os

# %% ../nbs/00_core.ipynb 5
class Text2StructAbstract(ABC):
    def __init__(self, text, what_to_retrieve, config):
        self.text, self.what_to_retrieve, self.config = text, what_to_retrieve, config
    
    @abstractmethod
    def run(self):
        pass
    
    @abstractmethod
    def get_dict(self):
        pass
    
    @abstractmethod
    def get_df(self):
        pass

# %% ../nbs/00_core.ipynb 6
class OpenAIProvider(Text2StructAbstract):
    
    def __init__(self, text, what_to_retrieve, config=None):
        super().__init__(text, what_to_retrieve, config)
        if config and 'openai_api_key' in config:
            openai.api_key = config['openai_api_key']
        elif os.environ.get('OPENAI_API_KEY'):
            openai.api_key = os.environ.get('OPENAI_API_KEY')
        else:
            raise Error("OpenAI key missing. Either provide OPEN_API_KEY env variable or use config['openai_api_key'] to set it!")
    
    def run(self):
        query_str = f"{self.text}\n\n| {' | '.join(self.what_to_retrieve)} |"
        self.__call(query_str)
        
    
    def __call(self, q):
        self.response = openai.Completion.create(
          model='text-davinci-002', #davinci-002 #curie-001 #babbage-001 #ada-001
          prompt=q,
          temperature=0,
          max_tokens=100,
          top_p=1.0,
          frequency_penalty=0.0,
          presence_penalty=0.0
        )
        
    def get_dict(self):
        elems = self.__parse_response()
        
        return [dict(zip(self.what_to_retrieve, elem)) for elem in elems]
    
    def get_df(self):
        return pd.DataFrame.from_records(self.get_dict())
    
    def __parse_response(self):
        resp = self.response['choices'][0].get('text')
        elems = resp.split('|\n|')[1:]
        elems = [elem.strip('|').strip() for elem in elems]
        elems = [elem.split(" | ") for elem in elems]

        return elems

# %% ../nbs/00_core.ipynb 10
def text2struct(text:str, # text to search for topics in
                topics_to_retrieve:list, # list of topics to retrieve
                config:dict=None): # config element that has things like api keys (like openai_api_key for example)
    o = OpenAIProvider(text, topics_to_retrieve, config)
    o.run()
    return o

# %% ../nbs/00_core.ipynb 11
def text2dict(text:str, # text to search for topics in
                topics_to_retrieve:list, # list of topics to retrieve
                config:dict=None): # config element that has things like api keys (like openai_api_key for example)
    o = text2struct(text, topics_to_retrieve, config)
    return o.get_dict()

def text2df(text:str, # text to search for topics in
            topics_to_retrieve:list, # list of topics to retrieve
            config:dict=None): # config element that has things like api keys (like openai_api_key for example)
    o = text2struct(text, topics_to_retrieve, config)
    return o.get_df()
