ai-coding-assistant 0.5.0__py3-none-any.whl

coding_assistant/context/__init__.py

@@ -0,0 +1,19 @@
+"""Context management module."""
+
+from coding_assistant.context.retriever import SemanticRetriever
+from coding_assistant.context.chunker import SmartChunker, CodeChunk
+from coding_assistant.context.hybrid_search import HybridSearch
+from coding_assistant.context.ranker import ContextRanker
+from coding_assistant.context.enhanced_retriever import EnhancedSemanticRetriever
+from coding_assistant.context.window import TokenWindowManager, TokenBudget
+
+__all__ = [
+    'SemanticRetriever',
+    'SmartChunker',
+    'CodeChunk',
+    'HybridSearch',
+    'ContextRanker',
+    'EnhancedSemanticRetriever',
+    'TokenWindowManager',
+    'TokenBudget',
+]

coding_assistant/context/chunker.py

@@ -0,0 +1,443 @@
+"""Smart code chunking for Python and JavaScript/TypeScript."""
+
+import ast
+from typing import List, Dict, Optional
+from dataclasses import dataclass
+
+
+@dataclass
+class CodeChunk:
+    """Represents a chunk of code."""
+    type: str  # 'function', 'class', 'method', 'file'
+    name: str
+    content: str
+    file_path: str
+    language: str
+    start_line: int
+    end_line: int
+    docstring: Optional[str] = None
+    imports: Optional[str] = None
+    metadata: Optional[Dict] = None
+
+    def __repr__(self):
+        return f"CodeChunk({self.type}:{self.name} in {self.file_path}:{self.start_line}-{self.end_line})"
+
+
+class SmartChunker:
+    """
+    Smart code chunking that preserves semantic units.
+
+    Chunks code at function/class level while preserving imports and context.
+    """
+
+    def __init__(self, max_chunk_tokens: int = 500):
+        """
+        Initialize the smart chunker.
+
+        Args:
+            max_chunk_tokens: Maximum tokens per chunk (approximate)
+        """
+        self.max_chunk_tokens = max_chunk_tokens
+
+    def chunk_code(self, code: str, file_path: str, language: str) -> List[CodeChunk]:
+        """
+        Chunk code based on language.
+
+        Args:
+            code: The source code to chunk
+            file_path: Path to the file
+            language: Programming language
+
+        Returns:
+            List of CodeChunk objects
+        """
+        language = language.lower()
+
+        if language == 'python':
+            return self._chunk_python(code, file_path)
+        elif language in ('javascript', 'typescript', 'jsx', 'tsx'):
+            return self._chunk_javascript(code, file_path, language)
+        else:
+            # Fallback: create a single file-level chunk
+            return [CodeChunk(
+                type='file',
+                name=file_path.split('/')[-1],
+                content=code[:self.max_chunk_tokens * 4],  # Rough char estimate
+                file_path=file_path,
+                language=language,
+                start_line=1,
+                end_line=len(code.split('\n'))
+            )]
+
+    def _chunk_python(self, code: str, file_path: str) -> List[CodeChunk]:
+        """
+        Chunk Python code using AST parsing.
+
+        Extracts functions, methods, and classes as separate chunks,
+        preserving imports with each chunk.
+        """
+        chunks = []
+
+        try:
+            tree = ast.parse(code)
+        except SyntaxError:
+            # If can't parse, return file-level chunk
+            return [CodeChunk(
+                type='file',
+                name=file_path.split('/')[-1],
+                content=code,
+                file_path=file_path,
+                language='python',
+                start_line=1,
+                end_line=len(code.split('\n'))
+            )]
+
+        # Extract imports
+        imports = self._extract_python_imports(tree, code)
+
+        # Extract top-level functions and classes
+        for node in ast.iter_child_nodes(tree):
+            if isinstance(node, ast.FunctionDef):
+                chunk = self._python_function_to_chunk(node, code, file_path, imports)
+                if chunk:
+                    chunks.append(chunk)
+
+            elif isinstance(node, ast.ClassDef):
+                # For classes, create one chunk for the whole class
+                class_chunk = self._python_class_to_chunk(node, code, file_path, imports)
+                if class_chunk:
+                    chunks.append(class_chunk)
+
+                # Also create individual chunks for each method
+                for class_node in node.body:
+                    if isinstance(class_node, ast.FunctionDef):
+                        method_chunk = self._python_function_to_chunk(
+                            class_node, code, file_path, imports,
+                            class_name=node.name
+                        )
+                        if method_chunk:
+                            chunks.append(method_chunk)
+
+        # If no chunks were created, return file-level chunk
+        if not chunks:
+            chunks.append(CodeChunk(
+                type='file',
+                name=file_path.split('/')[-1],
+                content=code,
+                file_path=file_path,
+                language='python',
+                start_line=1,
+                end_line=len(code.split('\n')),
+                imports=imports
+            ))
+
+        return chunks
+
+    def _extract_python_imports(self, tree: ast.AST, code: str) -> str:
+        """Extract import statements from Python code."""
+        imports_lines = []
+        lines = code.split('\n')
+
+        for node in ast.walk(tree):
+            if isinstance(node, (ast.Import, ast.ImportFrom)):
+                if hasattr(node, 'lineno'):
+                    line_idx = node.lineno - 1
+                    if line_idx < len(lines):
+                        imports_lines.append(lines[line_idx])
+
+        return '\n'.join(imports_lines)
+
+    def _python_function_to_chunk(self, node: ast.FunctionDef, code: str,
+                                  file_path: str, imports: str,
+                                  class_name: Optional[str] = None) -> Optional[CodeChunk]:
+        """Convert Python function AST node to CodeChunk."""
+        lines = code.split('\n')
+
+        # Get function source
+        start_line = node.lineno
+        end_line = node.end_lineno if hasattr(node, 'end_lineno') else start_line
+
+        if start_line > len(lines) or end_line > len(lines):
+            return None
+
+        function_code = '\n'.join(lines[start_line - 1:end_line])
+
+        # Get docstring
+        docstring = ast.get_docstring(node)
+
+        # Combine imports with function code
+        full_content = imports + '\n\n' + function_code if imports else function_code
+
+        # Determine chunk type and name
+        chunk_type = 'method' if class_name else 'function'
+        name = f"{class_name}.{node.name}" if class_name else node.name
+
+        return CodeChunk(
+            type=chunk_type,
+            name=name,
+            content=full_content,
+            file_path=file_path,
+            language='python',
+            start_line=start_line,
+            end_line=end_line,
+            docstring=docstring,
+            imports=imports,
+            metadata={
+                'class': class_name,
+                'decorators': [d.id if isinstance(d, ast.Name) else str(d) for d in node.decorator_list],
+                'args': [arg.arg for arg in node.args.args]
+            }
+        )
+
+    def _python_class_to_chunk(self, node: ast.ClassDef, code: str,
+                               file_path: str, imports: str) -> Optional[CodeChunk]:
+        """Convert Python class AST node to CodeChunk."""
+        lines = code.split('\n')
+
+        start_line = node.lineno
+        end_line = node.end_lineno if hasattr(node, 'end_lineno') else start_line
+
+        if start_line > len(lines) or end_line > len(lines):
+            return None
+
+        class_code = '\n'.join(lines[start_line - 1:end_line])
+
+        # Get docstring
+        docstring = ast.get_docstring(node)
+
+        # Get method names
+        methods = [n.name for n in node.body if isinstance(n, ast.FunctionDef)]
+
+        # Combine imports with class code
+        full_content = imports + '\n\n' + class_code if imports else class_code
+
+        return CodeChunk(
+            type='class',
+            name=node.name,
+            content=full_content,
+            file_path=file_path,
+            language='python',
+            start_line=start_line,
+            end_line=end_line,
+            docstring=docstring,
+            imports=imports,
+            metadata={
+                'methods': methods,
+                'bases': [base.id if isinstance(base, ast.Name) else str(base) for base in node.bases]
+            }
+        )
+
+    def _chunk_javascript(self, code: str, file_path: str, language: str) -> List[CodeChunk]:
+        """
+        Chunk JavaScript/TypeScript code.
+
+        Uses regex patterns to identify functions and classes.
+        For better accuracy, tree-sitter could be used.
+        """
+        chunks = []
+        lines = code.split('\n')
+
+        # Extract imports (ES6 and CommonJS)
+        imports = self._extract_js_imports(code)
+
+        # Try to find functions and classes using regex
+        # This is a simplified approach; tree-sitter would be more accurate
+
+        import re
+
+        # Pattern for function declarations
+        func_pattern = r'^(?:export\s+)?(?:async\s+)?function\s+(\w+)\s*\('
+        # Pattern for arrow functions assigned to const/let/var
+        arrow_pattern = r'^(?:export\s+)?(?:const|let|var)\s+(\w+)\s*=\s*(?:async\s*)?\([^)]*\)\s*=>'
+        # Pattern for class declarations
+        class_pattern = r'^(?:export\s+)?class\s+(\w+)'
+
+        current_line = 0
+        while current_line < len(lines):
+            line = lines[current_line].strip()
+
+            # Check for function
+            func_match = re.match(func_pattern, line)
+            arrow_match = re.match(arrow_pattern, line)
+            class_match = re.match(class_pattern, line)
+
+            if func_match or arrow_match:
+                # Found a function
+                func_name = func_match.group(1) if func_match else arrow_match.group(1)
+                start_line = current_line + 1
+
+                # Find end of function (look for closing brace)
+                end_line = self._find_js_block_end(lines, current_line)
+
+                if end_line > start_line:
+                    function_code = '\n'.join(lines[current_line:end_line])
+                    full_content = imports + '\n\n' + function_code if imports else function_code
+
+                    chunks.append(CodeChunk(
+                        type='function',
+                        name=func_name,
+                        content=full_content,
+                        file_path=file_path,
+                        language=language,
+                        start_line=start_line,
+                        end_line=end_line,
+                        imports=imports
+                    ))
+
+                    current_line = end_line
+                    continue
+
+            elif class_match:
+                # Found a class
+                class_name = class_match.group(1)
+                start_line = current_line + 1
+
+                # Find end of class
+                end_line = self._find_js_block_end(lines, current_line)
+
+                if end_line > start_line:
+                    class_code = '\n'.join(lines[current_line:end_line])
+                    full_content = imports + '\n\n' + class_code if imports else class_code
+
+                    chunks.append(CodeChunk(
+                        type='class',
+                        name=class_name,
+                        content=full_content,
+                        file_path=file_path,
+                        language=language,
+                        start_line=start_line,
+                        end_line=end_line,
+                        imports=imports
+                    ))
+
+                    current_line = end_line
+                    continue
+
+            current_line += 1
+
+        # If no chunks, return file-level chunk
+        if not chunks:
+            chunks.append(CodeChunk(
+                type='file',
+                name=file_path.split('/')[-1],
+                content=code,
+                file_path=file_path,
+                language=language,
+                start_line=1,
+                end_line=len(lines),
+                imports=imports
+            ))
+
+        return chunks
+
+    def _extract_js_imports(self, code: str) -> str:
+        """Extract import statements from JavaScript/TypeScript code."""
+        import re
+
+        lines = code.split('\n')
+        import_lines = []
+
+        # Patterns for imports
+        import_pattern = r'^import\s+'
+        require_pattern = r'^(?:const|let|var)\s+.*=\s*require\('
+
+        for line in lines:
+            stripped = line.strip()
+            if re.match(import_pattern, stripped) or re.match(require_pattern, stripped):
+                import_lines.append(line)
+
+        return '\n'.join(import_lines)
+
+    def _find_js_block_end(self, lines: List[str], start_idx: int) -> int:
+        """
+        Find the end of a JavaScript code block (function/class).
+
+        Uses brace matching to find the closing brace.
+        """
+        brace_count = 0
+        in_block = False
+
+        for i in range(start_idx, len(lines)):
+            line = lines[i]
+
+            # Count braces
+            for char in line:
+                if char == '{':
+                    brace_count += 1
+                    in_block = True
+                elif char == '}':
+                    brace_count -= 1
+
+                if in_block and brace_count == 0:
+                    return i + 1
+
+        # If no matching brace found, return end of file
+        return len(lines)
+
+    def chunk_with_token_limit(self, chunks: List[CodeChunk],
+                               max_tokens: int = 500) -> List[CodeChunk]:
+        """
+        Further split chunks that exceed token limit.
+
+        Args:
+            chunks: List of CodeChunk objects
+            max_tokens: Maximum tokens per chunk
+
+        Returns:
+            List of chunks, with large chunks split
+        """
+        result = []
+
+        for chunk in chunks:
+            # Rough estimate: 1 token ≈ 4 characters
+            estimated_tokens = len(chunk.content) / 4
+
+            if estimated_tokens <= max_tokens:
+                result.append(chunk)
+            else:
+                # Split into smaller chunks using sliding window
+                sub_chunks = self._sliding_window_split(chunk, max_tokens)
+                result.extend(sub_chunks)
+
+        return result
+
+    def _sliding_window_split(self, chunk: CodeChunk, max_tokens: int) -> List[CodeChunk]:
+        """Split a large chunk using sliding window approach."""
+        max_chars = max_tokens * 4  # Rough estimate
+        overlap = max_chars // 4  # 25% overlap
+
+        content_lines = chunk.content.split('\n')
+        sub_chunks = []
+
+        start = 0
+        chunk_num = 0
+
+        while start < len(content_lines):
+            # Take window of lines
+            end = start
+            current_chars = 0
+
+            while end < len(content_lines) and current_chars < max_chars:
+                current_chars += len(content_lines[end])
+                end += 1
+
+            # Create sub-chunk
+            sub_content = '\n'.join(content_lines[start:end])
+            sub_chunks.append(CodeChunk(
+                type=f"{chunk.type}_part",
+                name=f"{chunk.name}_part{chunk_num}",
+                content=sub_content,
+                file_path=chunk.file_path,
+                language=chunk.language,
+                start_line=chunk.start_line + start,
+                end_line=chunk.start_line + end,
+                imports=chunk.imports,
+                metadata={**chunk.metadata, 'is_partial': True} if chunk.metadata else {'is_partial': True}
+            ))
+
+            chunk_num += 1
+
+            # Move window with overlap
+            start = max(start + 1, end - (overlap // 50))  # Overlap in lines
+
+        return sub_chunks