# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/01_parse.ipynb.

# %% auto 0
__all__ = ['ChatData']

# %% ../nbs/01_parse.ipynb 2
import json, os
from openai import OpenAI

# %% ../nbs/01_parse.ipynb 4
class ChatData:
    "Process multi-turn chat data from openai."
    def __init__(self, data):
        self.data = data
        self.funcs = self._extract_functions()
    
    def _extract_functions(self):
        "Extract the set of functions"
        self.consistent = True
        functions = []
        for i, row in enumerate(self.data):
            if 'functions' in row:
                for func in row['functions']:
                    if func not in functions:
                        functions.append(func)
                        if i>0: self.consistent = False # this function is unique for this row
        return functions

    @classmethod
    def load_jsonl(cls, file_path):
        "Load openai jsonl data and parse it."
        data = []
        with open(file_path, 'r') as f:
            for line in f:
                data.append(json.loads(line))
        print(f'Loaded {len(data)} rows from {file_path}')
        return cls(data)

    def to_md(self, exclude_funcs=True):
        "Flatten the json messag format to markdown."
        md_data = []
        for row in self.data:
            markdown = ''
            for message in row['messages']:
                if ('weight' not in message or message['weight'] != 0):
                    is_func_call = 'function_call' in message
                    markdown += f"## Role: {'function-call' if is_func_call else message['role']}\n"
                    if message['content']:
                        markdown += f"{message['content']}\n\n"
                    if is_func_call:
                        markdown += f"{message['function_call']}\n\n"
            
            if 'functions' in row and row['functions'] and not exclude_funcs:
                markdown += "## Functions:\n"
                for func in row['functions']:
                    markdown += f"- {func['name']}\n"
                markdown += "\n"
            
            md_data.append(markdown)
        return md_data
