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

coding_assistant/context/window.py

@@ -0,0 +1,362 @@
+"""Token window management for LLM context limits."""
+
+from typing import List, Dict, Optional, Any
+import tiktoken
+from dataclasses import dataclass
+
+
+@dataclass
+class TokenBudget:
+    """Token budget allocation."""
+    total: int
+    system_prompt: int
+    query: int
+    history: int
+    chunks: int
+    response_reserved: int
+
+    @property
+    def used(self) -> int:
+        """Total tokens used."""
+        return self.system_prompt + self.query + self.history + self.chunks
+
+    @property
+    def available(self) -> int:
+        """Available tokens for response."""
+        return self.total - self.used - self.response_reserved
+
+    @property
+    def utilization(self) -> float:
+        """Context window utilization (0-1)."""
+        return self.used / (self.total - self.response_reserved)
+
+
+class TokenWindowManager:
+    """
+    Manages context within LLM token limits.
+
+    This class handles:
+    - Token counting for text
+    - Context prioritization
+    - Truncation when needed
+    - Budget allocation
+
+    Default allocations:
+    - System prompt: 5-10% (reserved)
+    - Response: 15-20% (reserved)
+    - Query: actual size
+    - History: up to 30%
+    - Code chunks: remaining space
+    """
+
+    def __init__(
+        self,
+        max_tokens: int = 150000,  # Default for Claude Sonnet
+        reserved_response: int = 4000,
+        reserved_system: int = 1000,
+        max_history_ratio: float = 0.3,  # Max 30% for history
+        encoding: str = "cl100k_base"  # OpenAI compatible encoding
+    ):
+        """
+        Initialize token window manager.
+
+        Args:
+            max_tokens: Maximum context window size
+            reserved_response: Tokens reserved for LLM response
+            reserved_system: Tokens reserved for system prompt
+            max_history_ratio: Maximum ratio of tokens for history
+            encoding: Tiktoken encoding to use
+        """
+        self.max_tokens = max_tokens
+        self.reserved_response = reserved_response
+        self.reserved_system = reserved_system
+        self.max_history_ratio = max_history_ratio
+
+        try:
+            self.encoder = tiktoken.get_encoding(encoding)
+        except Exception:
+            # Fallback to basic estimation
+            self.encoder = None
+
+    def count_tokens(self, text: str) -> int:
+        """
+        Count tokens in text.
+
+        Args:
+            text: Text to count tokens for
+
+        Returns:
+            Number of tokens
+        """
+        if not text:
+            return 0
+
+        if self.encoder:
+            try:
+                return len(self.encoder.encode(text))
+            except Exception:
+                pass
+
+        # Fallback: rough estimation (1 token ≈ 4 characters)
+        return len(text) // 4
+
+    def build_context(
+        self,
+        query: str,
+        chunks: List[Dict[str, Any]],
+        history: Optional[List[Dict[str, str]]] = None,
+        system_prompt: Optional[str] = None
+    ) -> Dict[str, Any]:
+        """
+        Build context within token limit.
+
+        Args:
+            query: User query
+            chunks: Code chunks (must have 'content' key)
+            history: Conversation history (list of {role, content})
+            system_prompt: System prompt text
+
+        Returns:
+            Dict with selected context and token budget
+        """
+        history = history or []
+
+        # Count tokens for fixed components
+        query_tokens = self.count_tokens(query)
+        system_tokens = self.count_tokens(system_prompt) if system_prompt else self.reserved_system
+
+        # Available tokens for history and chunks
+        available = self.max_tokens - self.reserved_response - system_tokens - query_tokens
+
+        if available <= 0:
+            # Query too long!
+            return {
+                'query': query,
+                'history': [],
+                'chunks': [],
+                'budget': self._create_budget(
+                    system_tokens, query_tokens, 0, 0
+                ),
+                'truncated': True,
+                'error': 'Query exceeds token limit'
+            }
+
+        # Allocate history (up to max_history_ratio)
+        max_history_tokens = int(available * self.max_history_ratio)
+        included_history, history_tokens = self._select_history(
+            history, max_history_tokens
+        )
+
+        # Remaining tokens for chunks
+        available_for_chunks = available - history_tokens
+
+        # Select chunks that fit
+        included_chunks, chunk_tokens = self._select_chunks(
+            chunks, available_for_chunks
+        )
+
+        budget = self._create_budget(
+            system_tokens, query_tokens, history_tokens, chunk_tokens
+        )
+
+        return {
+            'query': query,
+            'history': included_history,
+            'chunks': included_chunks,
+            'budget': budget,
+            'truncated': len(included_chunks) < len(chunks) or len(included_history) < len(history)
+        }
+
+    def _select_history(
+        self,
+        history: List[Dict[str, str]],
+        max_tokens: int
+    ) -> tuple[List[Dict[str, str]], int]:
+        """
+        Select history messages that fit within token limit.
+
+        Prioritizes recent messages.
+
+        Args:
+            history: Full conversation history
+            max_tokens: Maximum tokens for history
+
+        Returns:
+            Tuple of (selected messages, total tokens)
+        """
+        if not history:
+            return [], 0
+
+        included = []
+        total_tokens = 0
+
+        # Start from most recent and work backwards
+        for msg in reversed(history[-20:]):  # Consider last 20 messages
+            content = msg.get('content', '')
+            msg_tokens = self.count_tokens(content)
+
+            if total_tokens + msg_tokens <= max_tokens:
+                included.insert(0, msg)  # Maintain chronological order
+                total_tokens += msg_tokens
+            else:
+                # Can't fit more messages
+                break
+
+        return included, total_tokens
+
+    def _select_chunks(
+        self,
+        chunks: List[Dict[str, Any]],
+        max_tokens: int
+    ) -> tuple[List[Dict[str, Any]], int]:
+        """
+        Select code chunks that fit within token limit.
+
+        Assumes chunks are already ranked by relevance.
+
+        Args:
+            chunks: Code chunks (ranked)
+            max_tokens: Maximum tokens for chunks
+
+        Returns:
+            Tuple of (selected chunks, total tokens)
+        """
+        if not chunks:
+            return [], 0
+
+        included = []
+        total_tokens = 0
+
+        for chunk in chunks:
+            content = chunk.get('content', '')
+            chunk_tokens = self.count_tokens(content)
+
+            if total_tokens + chunk_tokens <= max_tokens:
+                included.append(chunk)
+                total_tokens += chunk_tokens
+            else:
+                # Try to fit at least one chunk
+                if not included and chunk_tokens <= max_tokens * 1.2:
+                    # Allow slight overflow for first chunk
+                    included.append(chunk)
+                    total_tokens += chunk_tokens
+                break
+
+        return included, total_tokens
+
+    def _create_budget(
+        self,
+        system_tokens: int,
+        query_tokens: int,
+        history_tokens: int,
+        chunk_tokens: int
+    ) -> TokenBudget:
+        """Create token budget summary."""
+        return TokenBudget(
+            total=self.max_tokens,
+            system_prompt=system_tokens,
+            query=query_tokens,
+            history=history_tokens,
+            chunks=chunk_tokens,
+            response_reserved=self.reserved_response
+        )
+
+    def estimate_response_tokens(self, response: str) -> int:
+        """
+        Estimate tokens in a response.
+
+        Args:
+            response: Response text
+
+        Returns:
+            Estimated token count
+        """
+        return self.count_tokens(response)
+
+    def can_fit(self, text: str, available_tokens: int) -> bool:
+        """
+        Check if text fits within token limit.
+
+        Args:
+            text: Text to check
+            available_tokens: Available token budget
+
+        Returns:
+            True if text fits
+        """
+        return self.count_tokens(text) <= available_tokens
+
+    def truncate_text(self, text: str, max_tokens: int) -> str:
+        """
+        Truncate text to fit within token limit.
+
+        Args:
+            text: Text to truncate
+            max_tokens: Maximum tokens allowed
+
+        Returns:
+            Truncated text
+        """
+        tokens = self.count_tokens(text)
+
+        if tokens <= max_tokens:
+            return text
+
+        # Binary search for the right length
+        if self.encoder:
+            encoded = self.encoder.encode(text)
+            truncated = encoded[:max_tokens]
+            return self.encoder.decode(truncated) + "..."
+
+        # Fallback: character-based truncation
+        ratio = max_tokens / tokens
+        char_limit = int(len(text) * ratio)
+        return text[:char_limit] + "..."
+
+    def get_model_limits(self, model_name: str) -> tuple[int, int]:
+        """
+        Get token limits for known models.
+
+        Args:
+            model_name: Model identifier
+
+        Returns:
+            Tuple of (max_tokens, recommended_response_reserve)
+        """
+        limits = {
+            'claude-sonnet': (200000, 4000),
+            'claude-opus': (200000, 4000),
+            'gpt-4': (128000, 4000),
+            'gpt-3.5': (16385, 2000),
+            'ollama-default': (4096, 512),
+        }
+
+        # Match by substring
+        for key, (max_tok, reserve) in limits.items():
+            if key in model_name.lower():
+                return max_tok, reserve
+
+        # Conservative default
+        return 8000, 1000
+
+    def format_budget_summary(self, budget: TokenBudget) -> str:
+        """
+        Format budget for display.
+
+        Args:
+            budget: Token budget
+
+        Returns:
+            Formatted string
+        """
+        return f"""Token Budget:
+  Total Window:    {budget.total:,} tokens
+  System Prompt:   {budget.system_prompt:,} tokens ({budget.system_prompt/budget.total*100:.1f}%)
+  Query:           {budget.query:,} tokens ({budget.query/budget.total*100:.1f}%)
+  History:         {budget.history:,} tokens ({budget.history/budget.total*100:.1f}%)
+  Code Chunks:     {budget.chunks:,} tokens ({budget.chunks/budget.total*100:.1f}%)
+  Reserved (Resp): {budget.response_reserved:,} tokens ({budget.response_reserved/budget.total*100:.1f}%)
+  ─────────────────────────────────────
+  Used:            {budget.used:,} tokens ({budget.utilization*100:.1f}%)
+  Available:       {budget.available:,} tokens
+"""

coding_assistant/documentation/__init__.py

@@ -0,0 +1,23 @@
+"""Documentation generation module for AI Coding Assistant.
+
+This module provides CodeWiki-style automated documentation generation:
+- Repository-level architecture understanding
+- Hierarchical decomposition for scalability
+- Multi-agent documentation generation
+- Visual diagram creation (Mermaid)
+- Incremental updates and continuous synchronization
+"""
+
+from coding_assistant.documentation.graph.dependency_builder import (
+    DependencyGraphBuilder,
+    CodeEntity,
+)
+from coding_assistant.documentation.graph.module_analyzer import ModuleAnalyzer
+
+__all__ = [
+    'DependencyGraphBuilder',
+    'CodeEntity',
+    'ModuleAnalyzer',
+]
+
+__version__ = '0.5.0'

coding_assistant/documentation/agents/__init__.py

@@ -0,0 +1,27 @@
+"""Multi-agent documentation generation system.
+
+This module implements parallel documentation generation using multiple LLM agents,
+inspired by CodeWiki's recursive multi-agent approach.
+"""
+
+from coding_assistant.documentation.agents.coordinator import (
+    MultiAgentCoordinator,
+    DocumentationTask,
+)
+from coding_assistant.documentation.agents.module_documenter import (
+    ModuleDocumenter,
+)
+from coding_assistant.documentation.agents.synthesizer import (
+    DocumentationSynthesizer,
+)
+from coding_assistant.documentation.agents.task_delegator import (
+    TaskDelegator,
+)
+
+__all__ = [
+    'MultiAgentCoordinator',
+    'DocumentationTask',
+    'ModuleDocumenter',
+    'DocumentationSynthesizer',
+    'TaskDelegator',
+]