23ws-LLMcoder/LLMcoder-GitHub-Python-Mix-Direct

1import os2import time3import configparser4 5from pathlib import Path6 7 8API_KEYS_LOCATION = os.path.join(os.path.dirname(os.path.realpath(__file__)), 'openaiapirc')9 10class PromptFile:11    context_source_filename = ""12    default_context_filename = "current_context.txt"13    default_file_path = os.path.join(os.path.dirname(__file__), "..", default_context_filename)14    default_config_path = os.path.join(os.path.dirname(__file__), "..", "current_context.config")15 16    def __init__(self, file_name, config):17        self.context_source_filename = "{}-context.txt".format(config['shell']) #  feel free to set your own default context path here18        19        self.file_path = self.default_file_path20        self.config_path = self.default_config_path21 22        # loading in one of the saved contexts23        if file_name != self.default_context_filename:24            self.load_context(file_name, True)25 26    def has_config(self):27        """28        Check if the prompt file has a corresponding config file29        """30        return os.path.isfile(self.config_path)31    32    def read_config(self):33        """34        Read the prompt config and return a dictionary35        """36 37        if self.has_config() == False:38            self.set_config(self.config)39            return self.config40        41        with open(self.config_path, 'r') as f:42            lines = f.readlines()43 44        config = {45            'engine': lines[0].split(':')[1].strip(),46            'temperature': float(lines[1].split(':')[1].strip()),47            'max_tokens': int(lines[2].split(':')[1].strip()),48            'shell': lines[3].split(':')[1].strip(),49            'multi_turn': lines[4].split(':')[1].strip(),50            'token_count': int(lines[5].split(':')[1].strip())51        }52 53        self.config = config54        return self.config 55    56    def set_config(self, config):57        """58        Set the prompt headers with the new config59        """60        self.config = config61        62        with open(self.config_path, 'w') as f:63            f.write('engine: {}\n'.format(self.config['engine']))64            f.write('temperature: {}\n'.format(self.config['temperature']))65            f.write('max_tokens: {}\n'.format(self.config['max_tokens']))66            f.write('shell: {}\n'.format(self.config['shell']))67            f.write('multi_turn: {}\n'.format(self.config['multi_turn']))68            f.write('token_count: {}\n'.format(self.config['token_count']))69    70    def show_config(self):71        print('\n')72        # read the dictionary into a list of # lines73        lines = []74        for key, value in self.config.items():75            lines.append('# {}: {}\n'.format(key, value))76        print(''.join(lines))77    78    def add_input_output_pair(self, user_query, prompt_response):79        """80        Add lines to file_name and update the token_count81        """82 83        with open(self.file_path, 'a') as f:84            f.write(user_query)85            f.write(prompt_response)86        87        if self.config['multi_turn'] == 'on':88            self.config['token_count'] += len(user_query.split()) + len(prompt_response.split())89            self.set_config(self.config)90    91    def read_prompt_file(self, input):92        """93        Get the updated prompt file94        Checks for token overflow and appends the current input95 96        Returns: the prompt file after appending the input97        """98 99        input_tokens_count = len(input.split())100        need_to_refresh = (self.config['token_count'] + input_tokens_count > 2048)101 102        if need_to_refresh:103            # delete first 2 lines of prompt context file104            with open(self.file_path, 'r') as f:105                lines = f.readlines()106                prompt = lines[2:] # drop first 2 lines of prompt107            with open(self.file_path, 'w') as f:108                f.writelines(prompt)109 110        # get input from prompt file111        with open(self.file_path, 'r') as f:112            lines = f.readlines()113 114        return ''.join(lines)115    116    def get_token_count(self):117        """118        Get the actual token count119        """120        token_count = 0121        if self.has_config():122            with open(self.config_path, 'r') as f:123                lines = f.readlines()124                token_count = int(lines[5].split(':')[1].strip())125        126        true_token_count