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

coding_assistant/documentation/agents/coordinator.py

@@ -0,0 +1,510 @@
+"""Coordinate multiple LLM agents for parallel documentation generation.
+
+This module implements the multi-agent coordination system that enables
+scalable, parallel documentation generation across repository partitions.
+"""
+
+from typing import List, Dict, Optional, Set
+from dataclasses import dataclass, field
+from pathlib import Path
+import asyncio
+from concurrent.futures import ThreadPoolExecutor
+import networkx as nx
+
+from coding_assistant.documentation.decomposition.partitioner import Partition
+from coding_assistant.documentation.decomposition.context_preserver import ModuleContext
+from coding_assistant.llm.client import LLMClientFactory
+from coding_assistant.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+@dataclass
+class DocumentationTask:
+    """
+    Represents a documentation generation task for an LLM agent.
+
+    Tasks are executed by individual agents and may have dependencies
+    on other tasks completing first.
+    """
+    task_id: str
+    partition: Partition
+    context: ModuleContext
+    priority: int = 0
+    dependencies: List[str] = field(default_factory=list)  # Task IDs that must complete first
+    status: str = 'pending'  # pending, running, completed, failed
+    result: Optional[str] = None
+    error: Optional[str] = None
+    metadata: Dict = field(default_factory=dict)
+
+    def __repr__(self):
+        return f"DocumentationTask({self.task_id}, status={self.status}, deps={len(self.dependencies)})"
+
+
+class MultiAgentCoordinator:
+    """
+    Coordinate multiple LLM agents for parallel documentation generation.
+
+    Inspired by CodeWiki's recursive multi-agent system, this coordinator:
+    - Manages task dependencies and scheduling
+    - Executes tasks in parallel waves
+    - Handles failures and retries
+    - Synthesizes results into coherent documentation
+    """
+
+    def __init__(self,
+                 max_concurrent_agents: int = 3,
+                 llm_client=None,
+                 use_async: bool = True):
+        """
+        Initialize the multi-agent coordinator.
+
+        Args:
+            max_concurrent_agents: Maximum number of parallel agents
+            llm_client: LLM client instance (creates default if None)
+            use_async: Whether to use async execution (True) or threaded (False)
+        """
+        self.max_concurrent = max_concurrent_agents
+        self.llm_client = llm_client or LLMClientFactory.create_client()
+        self.use_async = use_async
+        self.executor = ThreadPoolExecutor(max_workers=max_concurrent_agents)
+
+        # Task tracking
+        self.tasks: Dict[str, DocumentationTask] = {}
+        self.completed_tasks: Set[str] = set()
+
+        logger.info(f"MultiAgentCoordinator initialized with {max_concurrent_agents} concurrent agents")
+
+    async def generate_documentation(self,
+                                       partitions: List[Partition],
+                                       contexts: Dict[str, ModuleContext],
+                                       parsed_files: Optional[Dict[str, Dict]] = None) -> Dict[str, str]:
+        """
+        Generate documentation for all partitions in parallel.
+
+        Args:
+            partitions: List of repository partitions
+            contexts: Dictionary mapping partition names to their contexts
+            parsed_files: Optional dictionary of parsed file information
+
+        Returns:
+            Dictionary mapping partition names to generated documentation (markdown)
+        """
+        logger.info(f"Starting parallel documentation generation for {len(partitions)} partitions")
+
+        # Step 1: Create tasks
+        tasks = self._create_tasks(partitions, contexts)
+        self.tasks = {task.task_id: task for task in tasks}
+
+        logger.info(f"Created {len(tasks)} documentation tasks")
+
+        # Step 2: Topologically sort tasks based on dependencies
+        sorted_tasks = self._topological_sort(tasks)
+
+        logger.info(f"Tasks sorted into {len(sorted_tasks)} execution waves")
+
+        # Step 3: Execute tasks in waves (respecting dependencies)
+        results = await self._execute_tasks_parallel(sorted_tasks, parsed_files)
+
+        logger.info(f"Documentation generation complete: {len(results)} partitions documented")
+
+        return results
+
+    def _create_tasks(self,
+                       partitions: List[Partition],
+                       contexts: Dict[str, ModuleContext]) -> List[DocumentationTask]:
+        """
+        Create documentation tasks with priorities and dependencies.
+
+        Args:
+            partitions: List of partitions to document
+            contexts: Context information for partitions
+
+        Returns:
+            List of DocumentationTask objects
+        """
+        tasks = []
+
+        for partition in partitions:
+            # Compute priority based on dependencies
+            # Partitions with fewer dependencies get higher priority
+            priority = len(partition.dependencies)
+
+            task = DocumentationTask(
+                task_id=partition.name,
+                partition=partition,
+                context=contexts.get(partition.name, ModuleContext(partition.name)),
+                priority=priority,
+                dependencies=partition.dependencies.copy()  # Copy to avoid mutation
+            )
+
+            tasks.append(task)
+
+        return tasks
+
+    def _topological_sort(self, tasks: List[DocumentationTask]) -> List[List[DocumentationTask]]:
+        """
+        Topologically sort tasks into execution waves.
+
+        Tasks in the same wave have no dependencies on each other and
+        can be executed in parallel.
+
+        Args:
+            tasks: List of tasks to sort
+
+        Returns:
+            List of task waves (each wave is a list of tasks)
+        """
+        # Build dependency graph
+        graph = nx.DiGraph()
+
+        for task in tasks:
+            graph.add_node(task.task_id)
+
+            for dep in task.dependencies:
+                if dep in [t.task_id for t in tasks]:
+                    graph.add_edge(dep, task.task_id)
+
+        # Group into waves by dependency level
+        waves = []
+        remaining = set(graph.nodes())
+        processed = set()
+
+        while remaining:
+            # Find nodes with no unprocessed dependencies
+            current_wave_ids = []
+
+            for node in remaining:
+                predecessors = set(graph.predecessors(node))
+                if not predecessors or predecessors.issubset(processed):
+                    current_wave_ids.append(node)
+
+            if not current_wave_ids:
+                # Circular dependency or disconnected graph
+                # Add all remaining nodes to final wave
+                logger.warning("Circular dependencies detected, adding remaining tasks to final wave")
+                current_wave_ids = list(remaining)
+
+            # Get task objects for this wave
+            current_wave = [task for task in tasks if task.task_id in current_wave_ids]
+
+            # Sort wave by priority (lower priority value = higher priority)
+            current_wave.sort(key=lambda t: t.priority)
+
+            waves.append(current_wave)
+            remaining -= set(current_wave_ids)
+            processed.update(current_wave_ids)
+
+        return waves
+
+    async def _execute_tasks_parallel(self,
+                                        task_waves: List[List[DocumentationTask]],
+                                        parsed_files: Optional[Dict[str, Dict]]) -> Dict[str, str]:
+        """
+        Execute tasks in parallel waves.
+
+        Args:
+            task_waves: List of task waves to execute
+            parsed_files: Optional parsed file information
+
+        Returns:
+            Dictionary mapping task IDs to generated documentation
+        """
+        results = {}
+
+        for wave_num, wave_tasks in enumerate(task_waves):
+            logger.info(f"Executing wave {wave_num + 1}/{len(task_waves)} with {len(wave_tasks)} tasks")
+
+            # Execute all tasks in this wave concurrently
+            if self.use_async:
+                wave_results = await self._execute_wave_async(wave_tasks, results, parsed_files)
+            else:
+                wave_results = await self._execute_wave_threaded(wave_tasks, results, parsed_files)
+
+            # Add results
+            for task_id, doc in wave_results.items():
+                results[task_id] = doc
+                self.completed_tasks.add(task_id)
+
+        return results
+
+    async def _execute_wave_async(self,
+                                    tasks: List[DocumentationTask],
+                                    dependency_results: Dict[str, str],
+                                    parsed_files: Optional[Dict[str, Dict]]) -> Dict[str, str]:
+        """Execute a wave of tasks asynchronously."""
+        # Create coroutines for all tasks in wave
+        coroutines = [
+            self._generate_partition_docs(task, dependency_results, parsed_files)
+            for task in tasks
+        ]
+
+        # Execute in parallel with semaphore to limit concurrency
+        semaphore = asyncio.Semaphore(self.max_concurrent)
+
+        async def bounded_task(coro, task):
+            async with semaphore:
+                try:
+                    result = await coro
+                    task.status = 'completed'
+                    task.result = result
+                    return task.task_id, result
+                except Exception as e:
+                    logger.error(f"Task {task.task_id} failed: {e}")
+                    task.status = 'failed'
+                    task.error = str(e)
+                    return task.task_id, f"# Error\n\nDocumentation generation failed: {e}"
+
+        bounded_coroutines = [
+            bounded_task(coro, task)
+            for coro, task in zip(coroutines, tasks)
+        ]
+
+        results = await asyncio.gather(*bounded_coroutines)
+
+        return dict(results)
+
+    async def _execute_wave_threaded(self,
+                                       tasks: List[DocumentationTask],
+                                       dependency_results: Dict[str, str],
+                                       parsed_files: Optional[Dict[str, Dict]]) -> Dict[str, str]:
+        """Execute a wave of tasks using thread pool (fallback for non-async LLMs)."""
+        loop = asyncio.get_event_loop()
+
+        def sync_generate(task):
+            try:
+                # Run async function in sync context
+                result = asyncio.run(self._generate_partition_docs(task, dependency_results, parsed_files))
+                task.status = 'completed'
+                task.result = result
+                return task.task_id, result
+            except Exception as e:
+                logger.error(f"Task {task.task_id} failed: {e}")
+                task.status = 'failed'
+                task.error = str(e)
+                return task.task_id, f"# Error\n\nDocumentation generation failed: {e}"
+
+        # Execute in thread pool
+        futures = [
+            loop.run_in_executor(self.executor, sync_generate, task)
+            for task in tasks
+        ]
+
+        results = await asyncio.gather(*futures)
+
+        return dict(results)
+
+    async def _generate_partition_docs(self,
+                                        task: DocumentationTask,
+                                        dependency_docs: Dict[str, str],
+                                        parsed_files: Optional[Dict[str, Dict]]) -> str:
+        """
+        Generate documentation for a single partition.
+
+        Args:
+            task: Documentation task
+            dependency_docs: Documentation from dependent partitions
+            parsed_files: Optional parsed file information
+
+        Returns:
+            Generated documentation as markdown string
+        """
+        logger.debug(f"Generating documentation for partition: {task.partition.name}")
+
+        task.status = 'running'
+
+        # Build comprehensive prompt
+        prompt = self._build_documentation_prompt(task, dependency_docs, parsed_files)
+
+        # Generate documentation using LLM
+        try:
+            # LLM clients use messages format and return generators
+            messages = [{"role": "user", "content": prompt}]
+
+            # Generate and collect response chunks
+            response_chunks = []
+            for chunk in self.llm_client.generate(messages=messages, stream=True):
+                response_chunks.append(chunk)
+
+            response = ''.join(response_chunks)
+            logger.debug(f"Documentation generated for {task.partition.name} ({len(response)} chars)")
+
+            return response
+
+        except Exception as e:
+            logger.error(f"LLM generation failed for {task.partition.name}: {e}")
+            # Return placeholder documentation
+            return self._generate_fallback_documentation(task)
+
+    def _build_documentation_prompt(self,
+                                     task: DocumentationTask,
+                                     dependency_docs: Dict[str, str],
+                                     parsed_files: Optional[Dict[str, Dict]]) -> str:
+        """
+        Build comprehensive prompt for partition documentation.
+
+        Args:
+            task: Documentation task
+            dependency_docs: Documentation from dependencies
+            parsed_files: Optional parsed files
+
+        Returns:
+            Prompt string for LLM
+        """
+        partition = task.partition
+        context = task.context
+
+        # Build file content summary
+        file_summaries = []
+        total_lines = 0
+
+        for file_path in partition.files[:10]:  # Limit to first 10 files
+            file_name = Path(file_path).name
+
+            if parsed_files and file_path in parsed_files:
+                parsed = parsed_files[file_path]
+                functions = parsed.get('functions', [])
+                classes = parsed.get('classes', [])
+
+                summary = f"**{file_name}**\n"
+                if classes:
+                    summary += f"  - Classes: {', '.join(c['name'] for c in classes[:5])}\n"
+                if functions:
+                    summary += f"  - Functions: {', '.join(f['name'] for f in functions[:5])}\n"
+
+                file_summaries.append(summary)
+            else:
+                file_summaries.append(f"**{file_name}**\n")
+
+            total_lines += 1
+
+        if len(partition.files) > 10:
+            file_summaries.append(f"\n... and {len(partition.files) - 10} more files")
+
+        # Build dependency context
+        dep_context = ""
+        if task.dependencies and dependency_docs:
+            dep_context = "\n### Dependencies\n\nThis module depends on:\n\n"
+            for dep_name in task.dependencies[:3]:
+                if dep_name in dependency_docs:
+                    dep_doc = dependency_docs[dep_name]
+                    # Extract first paragraph
+                    first_para = dep_doc.split('\n\n')[0] if dep_doc else "No documentation"
+                    dep_context += f"**{dep_name}**: {first_para[:200]}...\n\n"
+
+        prompt = f"""Generate comprehensive documentation for the following code module.
+
+## Module: {partition.name}
+
+### Overview
+- **Files**: {len(partition.files)}
+- **Lines of Code**: {partition.size_loc}
+- **Cohesion Score**: {partition.cohesion_score:.2f}
+- **Architectural Role**: {context.architectural_role}
+- **Level**: {partition.level} (0=top-level, 1=module, 2=component)
+
+### Architectural Context
+- **Related Modules**: {', '.join(context.related_partitions[:5]) if context.related_partitions else 'None'}
+- **Design Patterns**: {', '.join(context.common_patterns) if context.common_patterns else 'None detected'}
+- **Public Interfaces**: {len(context.exports)} exports
+
+### Files in This Module
+{chr(10).join(file_summaries)}
+
+{dep_context}
+
+## Task
+Generate detailed, professional documentation in Markdown format that includes:
+
+1. **Module Overview** (2-3 sentences)
+   - What this module does
+   - Its purpose in the overall system
+
+2. **Architecture** (1 paragraph)
+   - How it's organized internally
+   - Key structural decisions
+
+3. **Key Components**
+   - Main classes/functions and their responsibilities
+   - Important data structures
+
+4. **Dependencies & Relationships**
+   - What this module depends on
+   - What depends on this module
+   - How it fits in the larger architecture
+
+5. **Public API** (if applicable)
+   - Exported functions/classes
+   - Usage examples
+   - API contracts
+
+6. **Design Decisions**
+   - Notable patterns or architectural choices
+   - Trade-offs made
+
+Format: Use clear Markdown with headers, bullet points, and code examples where appropriate.
+Tone: Professional, concise, developer-focused.
+"""
+
+        return prompt
+
+    def _generate_fallback_documentation(self, task: DocumentationTask) -> str:
+        """Generate basic fallback documentation when LLM fails."""
+        partition = task.partition
+        context = task.context
+
+        doc = f"""# {partition.name}
+
+## Overview
+
+This module contains {len(partition.files)} files with approximately {partition.size_loc} lines of code.
+
+**Architectural Role**: {context.architectural_role}
+
+## Files
+
+"""
+        for file_path in partition.files[:20]:
+            doc += f"- `{Path(file_path).name}`\n"
+
+        if len(partition.files) > 20:
+            doc += f"\n... and {len(partition.files) - 20} more files\n"
+
+        if context.common_patterns:
+            doc += f"\n## Design Patterns\n\n"
+            for pattern in context.common_patterns:
+                doc += f"- {pattern}\n"
+
+        if context.related_partitions:
+            doc += f"\n## Related Modules\n\n"
+            for related in context.related_partitions[:10]:
+                doc += f"- {related}\n"
+
+        doc += f"\n## Metrics\n\n"
+        doc += f"- **Cohesion**: {partition.cohesion_score:.2f}\n"
+        doc += f"- **Dependencies**: {len(partition.dependencies)}\n"
+
+        doc += "\n---\n\n*Note: Detailed documentation could not be generated. This is a placeholder.*\n"
+
+        return doc
+
+    def get_task_status(self, task_id: str) -> Optional[DocumentationTask]:
+        """Get status of a specific task."""
+        return self.tasks.get(task_id)
+
+    def get_completion_stats(self) -> Dict:
+        """Get statistics about task completion."""
+        total = len(self.tasks)
+        completed = len([t for t in self.tasks.values() if t.status == 'completed'])
+        failed = len([t for t in self.tasks.values() if t.status == 'failed'])
+        running = len([t for t in self.tasks.values() if t.status == 'running'])
+        pending = len([t for t in self.tasks.values() if t.status == 'pending'])
+
+        return {
+            'total': total,
+            'completed': completed,
+            'failed': failed,
+            'running': running,
+            'pending': pending,
+            'completion_rate': completed / total if total > 0 else 0
+        }

coding_assistant/documentation/agents/module_documenter.py

@@ -0,0 +1,111 @@
+"""Per-module documentation generation.
+
+This module provides utilities for generating documentation for individual
+code modules/partitions.
+"""
+
+from typing import Dict, List, Optional
+from pathlib import Path
+
+from coding_assistant.documentation.decomposition.partitioner import Partition
+from coding_assistant.documentation.decomposition.context_preserver import ModuleContext
+from coding_assistant.utils.logger import get_logger
+
+logger = get_logger(__name__)
+
+
+class ModuleDocumenter:
+    """
+    Generate documentation for individual modules.
+
+    This class encapsulates the logic for documenting a single module/partition,
+    which is then used by the coordinator for parallel execution.
+    """
+
+    def __init__(self, llm_client=None):
+        """
+        Initialize module documenter.
+
+        Args:
+            llm_client: LLM client instance
+        """
+        from coding_assistant.llm.client import LLMClientFactory
+        self.llm_client = llm_client or LLMClientFactory.create_client()
+
+    def generate(self,
+                 partition: Partition,
+                 context: ModuleContext,
+                 parsed_files: Optional[Dict[str, Dict]] = None) -> str:
+        """
+        Generate documentation for a single partition.
+
+        Args:
+            partition: Partition to document
+            context: Architectural context
+            parsed_files: Optional parsed file data
+
+        Returns:
+            Generated documentation as markdown
+        """
+        logger.info(f"Generating documentation for module: {partition.name}")
+
+        prompt = self._build_prompt(partition, context, parsed_files)
+
+        try:
+            # LLM clients use messages format and return generators
+            messages = [{"role": "user", "content": prompt}]
+
+            # Generate and collect response chunks
+            response_chunks = []
+            for chunk in self.llm_client.generate(messages=messages, stream=True):
+                response_chunks.append(chunk)
+
+            response = ''.join(response_chunks)
+            return response
+
+        except Exception as e:
+            logger.error(f"Documentation generation failed: {e}")
+            return self._generate_fallback(partition, context)
+
+    def _build_prompt(self,
+                       partition: Partition,
+                       context: ModuleContext,
+                       parsed_files: Optional[Dict[str, Dict]]) -> str:
+        """Build prompt for module documentation."""
+        # Similar to coordinator's prompt builder
+        file_list = "\n".join(f"- {Path(f).name}" for f in partition.files[:15])
+
+        if len(partition.files) > 15:
+            file_list += f"\n- ... and {len(partition.files) - 15} more files"
+
+        return f"""Document the following code module:
+
+**Module**: {partition.name}
+**Role**: {context.architectural_role}
+**Files**: {len(partition.files)}
+**LOC**: {partition.size_loc}
+
+**Files**:
+{file_list}
+
+**Design Patterns**: {', '.join(context.common_patterns) if context.common_patterns else 'None'}
+
+Generate concise, professional documentation covering:
+1. Purpose and overview
+2. Key components
+3. Architecture
+4. Public API (if any)
+
+Format as markdown."""
+
+    def _generate_fallback(self, partition: Partition, context: ModuleContext) -> str:
+        """Generate fallback documentation."""
+        return f"""# {partition.name}
+
+## Overview
+Module with {len(partition.files)} files ({partition.size_loc} LOC).
+
+**Role**: {context.architectural_role}
+
+*Detailed documentation could not be generated.*
+"""