salim0986/graph-bug-ai
0
1import tiktoken2from typing import List, Dict, Any, Optional3from tree_sitter import Node4import os5 6class SemanticChunker:7 """8 Splits source code into semantic chunks respecting AST boundaries.9 Uses tiktoken for fast token estimation.10 """11 def __init__(self, max_tokens: int = 512, overlap_tokens: int = 64):12 self.max_tokens = max_tokens13 self.overlap_tokens = overlap_tokens14 # Use cl100k_base as a fast generic tokenizer approximation15 self.encoder = tiktoken.get_encoding("cl100k_base")16 17 # AST node types that represent structural boundaries18 self.function_types = {19 "function_definition", "function_declaration", "function_item",20 "method_definition", "method_declaration",21 "arrow_function", "generator_function_declaration"22 }23 self.class_types = {24 "class_definition", "class_declaration",25 "interface_declaration", "impl_item"26 }27 28 def _get_tokens(self, text: str) -> List[int]:29 return self.encoder.encode(text, disallowed_special=())30 31 def _decode_tokens(self, tokens: List[int]) -> str:32 return self.encoder.decode(tokens)33 34 def _split_large_text(self, text: str, start_line: int, parent_func: Optional[str], parent_class: Optional[str]) -> List[Dict[str, Any]]:35 """Split a large text block into overlapping chunks."""36 tokens = self._get_tokens(text)37 chunks = []38 39 if len(tokens) <= self.max_tokens:40 return [{41 "text": text,42 "start_line": start_line,43 "end_line": start_line + text.count("\n"),44 "parent_function": parent_func,45 "parent_class": parent_class46 }]47 48 step = self.max_tokens - self.overlap_tokens49 if step <= 0:50 step = self.max_tokens # Fallback: no overlap51 52 # Pre-compute cumulative newline counts at each token boundary.53 # This turns line-number estimation from O(N²) into O(N) total.54 # We walk the text once, mapping token offset → newline count.55 token_newlines: List[int] = [0] * (len(tokens) + 1)56 decoded_so_far = ""57 for idx in range(len(tokens)):58 decoded_so_far = self._decode_tokens(tokens[:idx + 1])59 token_newlines[idx + 1] = decoded_so_far.count("\n")60 61 for i in range(0, len(tokens), step):62 chunk_tokens = tokens[i:i + self.max_tokens]63 chunk_text = self._decode_tokens(chunk_tokens)64 chunk_start_line = start_line + token_newlines[i]65 chunk_end_line = start_line + token_newlines[min(i + self.max_tokens, len(tokens))]66 67 chunks.append({68 "text": chunk_text,69 "start_line": chunk_start_line,70 "end_line": chunk_end_line,71 "parent_function": parent_func,72 "parent_class": parent_class73 })74 75 return chunks76 77 def _extract_name(self, node: Node, code_bytes: bytes) -> Optional[str]:78 """Extract the name of a function or class node."""79 name_node = node.child_by_field_name("name")80 if name_node:81 return code_bytes[name_node.start_byte:name_node.end_byte].decode("utf8", errors="ignore")82 83 # Fallback: find the first identifier child84 for child in node.children:85 if child.type == "identifier" or child.type == "type_identifier":86 return code_bytes[child.start_byte:child.end_byte].decode("utf8", errors="ignore")87 return None88 89 def _process_node(90 self, 91 node: Node, 92 code_bytes: bytes, 93 current_class: Optional[str] = None94 ) -> List[Dict[str, Any]]:95 """Recursively process nodes, extracting semantic chunks."""96 chunks = []97 node_type = node.type98 99 if node_type in self.class_types:100 # We found a class, update current_class and process children101 class_name = self._extract_name(node, code_bytes) or current_class102 node_text = code_bytes[node.start_byte:node.end_byte].decode("utf8", errors="ignore")103 tokens = self._get_tokens(node_text)104 if len(tokens) <= self.max_tokens:105 chunks.append({106 "text": node_text,107 "start_line": node.start_point[0],108 "end_line": node.end_point[0],109 "parent_function": None,110 "parent_class": class_name111 })112 else:113 for child in node.children:114 chunks.extend(self._process_node(child, code_bytes, class_name))115 return chunks116 117 elif node_type in self.function_types:118 # We found a function/method. Treat it as a single unit (split if too large)119 func_name = self._extract_name(node, code_bytes)120 node_text = code_bytes[node.start_byte:node.end_byte].decode("utf8", errors="ignore")121 chunks.extend(self._split_large_text(122 text=node_text,123 start_line=node.start_point[0],124 parent_func=func_name,125 parent_class=current_class126 ))127 return chunks128 129 # If it's a root node, process its children130 if node.parent is None:131 for child in node.children:132 chunks.extend(self._process_node(child, code_bytes, current_class))133 return chunks134 135 # Top-level statements (e.g. imports) or class-level statements (e.g. fields)136 # We can identify them if their parent is the root or if their parent is a class/block137 # To avoid leaf nodes, we only chunk statements (nodes with children or specific types)138 # A simpler way: if we are here and we didn't recurse yet, we should recurse to find functions, 139 # but also capture non-function statements.140 # Actually, let's just recurse. For simplicity, we might lose some non-function statements if they are 141 # nested and we don't capture them. Let's capture them if they are direct children of root or class body.142 143 is_top_level = node.parent and node.parent.parent is None144 is_class_level = current_class and node.parent and node.parent.type in ("block", "declaration_list")145 146 if is_top_level or is_class_level:147 node_text = code_bytes[node.start_byte:node.end_byte].decode("utf8", errors="ignore")148 if node_text.strip() and not any(child.type in self.function_types for child in node.children):149 chunks.extend(self._split_large_text(150 text=node_text,151 start_line=node.start_point[0],152 parent_func=None,153 parent_class=current_class154 ))155 return chunks156 157 # Otherwise, keep recursing to find nested functions/classes158 for child in node.children:159 chunks.extend(self._process_node(child, code_bytes, current_class))160 161 return chunks162 163 def chunk_ast(self, tree: Any, code_bytes: bytes, file_path: str, language: str) -> List[Dict[str, Any]]:164 """165 Takes a tree-sitter AST and code, returns list of chunks.166 Each chunk: {file, start_line, end_line, parent_function, parent_class, language, raw_code}167 """168 raw_chunks = self._process_node(tree.root_node, code_bytes)169 170 # Post-process to add file and language metadata, and group tiny adjacent chunks?171 # Grouping tiny adjacent chunks (like sequential imports) makes retrieval better.172 grouped_chunks = []173 current_group = None174 current_tokens = 0175 176 for c in raw_chunks:177 c_tokens = len(self._get_tokens(c["text"]))178 179 # If the chunk has a function or class context, keep it separate from global blocks180 has_context = c["parent_function"] is not None or c["parent_class"] is not None181 182 if not has_context and current_group is not None and (current_tokens + c_tokens) <= self.max_tokens:183 # Append to current group184 current_group["text"] += "\n" + c["text"]185 current_group["end_line"] = c["end_line"]186 current_tokens += c_tokens187 else:188 if current_group is not None:189 grouped_chunks.append(current_group)190 191 if has_context or c_tokens > 20: # Only start new group if it's substantial or has context192 current_group = c.copy()193 current_tokens = c_tokens194 elif not has_context:195 # Start a new group anyway for small global statements196 current_group = c.copy()197 current_tokens = c_tokens198 199 if current_group is not None:200 grouped_chunks.append(current_group)201 202 # Finalize format203 final_chunks = []204 for c in grouped_chunks:205 final_chunks.append({206 "file": file_path,207 "language": language,208 "start_line": c["start_line"],209 "end_line": c["end_line"],210 "parent_function": c.get("parent_function"),211 "parent_class": c.get("parent_class"),212 "raw_code": c["text"].strip()213 })214 215 return final_chunks216