CintraAI/code-chunker

1from abc import ABC, abstractmethod2from CodeParser import CodeParser3from utils import count_tokens4 5 6class Chunker(ABC):7    def __init__(self, encoding_name="gpt-4"):8        self.encoding_name = encoding_name9 10    @abstractmethod11    def chunk(self, content, token_limit):12        pass13 14    @abstractmethod15    def get_chunk(self, chunked_content, chunk_number):16        pass17 18    @staticmethod19    def print_chunks(chunks):20        for chunk_number, chunk_code in chunks.items():21            print(f"Chunk {chunk_number}:")22            print("=" * 40)23            print(chunk_code)24            print("=" * 40)25 26    @staticmethod27    def consolidate_chunks_into_file(chunks):28        return "\n".join(chunks.values())29 30    @staticmethod31    def count_lines(consolidated_chunks):32        lines = consolidated_chunks.split("\n")33        return len(lines)34 35 36class CodeChunker(Chunker):37    def __init__(self, file_extension, encoding_name="gpt-4"):38        super().__init__(encoding_name)39        self.file_extension = file_extension40 41    def chunk(self, code, token_limit) -> dict:42        code_parser = CodeParser(self.file_extension)43        chunks = {}44        current_chunk = ""45        token_count = 046        lines = code.split("\n")47        i = 048        chunk_number = 149        start_line = 050        breakpoints = sorted(code_parser.get_lines_for_points_of_interest(code, self.file_extension))51        comments = sorted(code_parser.get_lines_for_comments(code, self.file_extension))52        adjusted_breakpoints = []53        for bp in breakpoints:54            current_line = bp - 155            highest_comment_line = None  # Initialize with None to indicate no comment line has been found yet56            while current_line in comments:57                highest_comment_line = current_line  # Update highest comment line found58                current_line -= 1  # Move to the previous line59 60            if highest_comment_line:  # If a highest comment line exists, add it61                adjusted_breakpoints.append(highest_comment_line)62            else:63                adjusted_breakpoints.append(64                    bp)  # If no comments were found before the breakpoint, add the original breakpoint65 66        breakpoints = sorted(set(adjusted_breakpoints))  # Ensure breakpoints are unique and sorted67 68        while i < len(lines):69            line = lines[i]70            new_token_count = count_tokens(line, self.encoding_name)71            if token_count + new_token_count > token_limit:72 73                # Set the stop line to the last breakpoint before the current line74                if i in breakpoints:75                    stop_line = i76                else:77                    stop_line = max(max([x for x in breakpoints if x < i], default=start_line), start_line)78 79                # If the stop line is the same as the start line, it means we haven't reached a breakpoint yet and we need to move to the next line to find one80                if stop_line == start_line and i not in breakpoints:81                    token_count += new_token_count82                    i += 183 84                # If the stop line is the same as the start line and the current line is a breakpoint, it means we can create a chunk with just the current line85                elif stop_line == start_line and i == stop_line:86                    token_count += new_token_count87                    i += 188 89 90                # If the stop line is the same as the start line and the current line is a breakpoint, it means we can create a chunk with just the current line91                elif stop_line == start_line and i in breakpoints:92                    current_chunk = "\n".join(lines[start_line:stop_line])93                    if current_chunk.strip():  # If the current chunk is not just whitespace94                        chunks[chunk_number] = current_chunk  # Using chunk_number as key95                        chunk_number += 196 97                    token_count = 098                    start_line = i99                    i += 1100 101                # If the stop line is different from the start line, it means we're at the end of a block102                else:103                    current_chunk = "\n".join(lines[start_line:stop_line])104                    if current_chunk.strip():105                        chunks[chunk_number] = current_chunk  # Using chunk_number as key106                        chunk_number += 1107 108                    i = stop_line109                    token_count = 0110                    start_line = stop_line111            else:112                # If the token count is still within the limit, add the line to the current chunk113                token_count += new_token_count114                i += 1115 116        # Append remaining code, if any, ensuring it's not empty or whitespace117        current_chunk_code = "\n".join(lines[start_line:])118        if current_chunk_code.strip():  # Checks if the chunk is not just whitespace119            chunks[chunk_number] = current_chunk_code  # Using chunk_number as key120 121        return chunks122 123    def get_chunk(self, chunked_codebase, chunk_number):124        return chunked_codebase[chunk_number]125