""" Code Analyzer Module Uses LlamaIndex + AI models for intelligent code analysis. """ import logging from dataclasses import dataclass from typing import Optional, List, Dict from ..config import get_config logger = logging.getLogger("codeatlas.analyzer") # LlamaIndex imports try: from llama_index.core.llms import ChatMessage LLAMAINDEX_AVAILABLE = True except ImportError: LLAMAINDEX_AVAILABLE = False logger.warning("LlamaIndex core not available") try: from llama_index.llms.gemini import Gemini GEMINI_AVAILABLE = True except ImportError: GEMINI_AVAILABLE = False logger.warning("LlamaIndex Gemini not available") try: from llama_index.llms.openai import OpenAI OPENAI_AVAILABLE = True except ImportError: OPENAI_AVAILABLE = False logger.warning("LlamaIndex OpenAI not available") # Fallback to google-genai try: from google import genai from google.genai import types GENAI_AVAILABLE = True except ImportError: GENAI_AVAILABLE = False # System prompts for different analysis tasks ARCHITECT_PROMPT = """You are CodeAtlas, an expert software architect. Generate a Graphviz DOT diagram showing code architecture with RELATIONSHIPS. CRITICAL CONSTRAINTS: - Maximum 15-20 nodes (focus on KEY architectural components) - Maximum 25-30 edges (most important relationships) - Group related components into subgraphs - Omit trivial files (tests, configs, utilities) WHAT TO SHOW: - Main entry points and core modules - Key classes/services with clear responsibilities - Important data flow and dependencies - Layer boundaries (API, Business Logic, Data) RULES: 1. Start with: digraph CodeArchitecture { 2. Every diagram MUST have arrows showing component connections 3. Use actual class/file names from the code 4. Group related items in clusters with descriptive labels 5. Use colors to distinguish layers/types EXAMPLE: ```dot digraph CodeArchitecture { rankdir=TB; node [shape=box, style="rounded,filled", fontname="Helvetica"]; subgraph cluster_api { label="API Layer"; style="rounded,filled"; fillcolor="#e8f5e9"; Routes; Handlers; } subgraph cluster_services { label="Business Logic"; style="rounded,filled"; fillcolor="#e3f2fd"; UserService; DataProcessor; } Routes -> Handlers; Handlers -> UserService; Handlers -> DataProcessor; } ``` Generate ONLY valid DOT code. Focus on architectural clarity.""" SUMMARY_PROMPT = """You are CodeAtlas. Analyze the codebase and provide a concise summary. Include: 1. **Project Overview**: What does this codebase do? 2. **Technology Stack**: Languages, frameworks, key dependencies 3. **Architecture Pattern**: MVC, microservices, monolith, etc. 4. **Key Components**: Main modules and their responsibilities 5. **Entry Points**: Where does execution start? Keep it concise (200-300 words). Be specific about actual file/class names.""" CHAT_PROMPT = """You are CodeAtlas, an expert software architect assistant. You're analyzing a codebase and helping answer questions about its architecture. Use the provided code context to give accurate, specific answers. Reference actual file names, class names, and code patterns when relevant. Be helpful, concise, and technical. If you're unsure about something, say so.""" @dataclass class AnalysisResult: """Result of code analysis.""" content: str success: bool = True error: Optional[str] = None class CodeAnalyzer: """Analyzes code using LlamaIndex and AI models.""" def __init__(self, api_key: Optional[str] = None, model_name: Optional[str] = None): self.config = get_config() self.api_key = api_key or self.config.gemini_api_key self.model_name = model_name or self.config.models.get_model_id(self.config.current_model) self._llm = None @property def llm(self): """Get or create the LLM instance.""" if self._llm is None: self._llm = self._create_llm() return self._llm def _create_llm(self): """Create appropriate LLM based on model name.""" is_openai = self.config.models.is_openai_model(self.model_name) if is_openai: if not OPENAI_AVAILABLE: raise ValueError("OpenAI support not available. Install llama-index-llms-openai") api_key = self.config.openai_api_key or self.api_key return OpenAI(api_key=api_key, model=self.model_name, temperature=0.7, max_tokens=4096) else: if GEMINI_AVAILABLE: return Gemini( api_key=self.api_key, model=f"models/{self.model_name}", temperature=0.7, max_tokens=4096, ) elif GENAI_AVAILABLE: return None # Will use fallback else: raise ValueError("No AI backend available") def _generate_with_llamaindex(self, system_prompt: str, user_prompt: str) -> str: """Generate content using LlamaIndex.""" if self.llm is None: return self._generate_with_genai(system_prompt, user_prompt) messages = [ ChatMessage(role="system", content=system_prompt), ChatMessage(role="user", content=user_prompt), ] response = self.llm.chat(messages) return response.message.content def _generate_with_genai(self, system_prompt: str, user_prompt: str) -> str: """Generate content using google-genai directly (fallback).""" if not GENAI_AVAILABLE: raise ValueError("No AI backend available") client = genai.Client(api_key=self.api_key) response = client.models.generate_content( model=self.model_name, contents=[user_prompt], config=types.GenerateContentConfig( system_instruction=system_prompt, temperature=0.7, max_output_tokens=4096, ) ) return response.text or "" def generate_diagram(self, code_context: str) -> AnalysisResult: """Generate an architecture diagram from code context. Args: code_context: Formatted code content Returns: AnalysisResult with DOT diagram or error """ user_prompt = f"""Analyze this codebase and generate an architecture diagram: {code_context} Generate a Graphviz DOT diagram showing the main components and their relationships.""" try: logger.info(f"Generating diagram with {self.model_name}") content = self._generate_with_llamaindex(ARCHITECT_PROMPT, user_prompt) if not content.strip(): return AnalysisResult(content="", success=False, error="Empty response from AI") # Extract DOT content import re match = re.search(r"```(?:dot|graphviz)?\s*(.*?)\s*```", content, re.DOTALL) dot_content = match.group(1).strip() if match else content.strip() # Validate DOT code if "digraph" not in dot_content and "graph" not in dot_content: return AnalysisResult( content="", success=False, error=f"Invalid DOT code: {dot_content[:200]}" ) return AnalysisResult(content=dot_content) except Exception as e: logger.exception("Diagram generation failed") error_str = str(e) if "429" in error_str or "RESOURCE_EXHAUSTED" in error_str: return AnalysisResult(content="", success=False, error="Rate limited. Please wait and try again.") elif "401" in error_str or "403" in error_str: return AnalysisResult(content="", success=False, error="Invalid API key.") return AnalysisResult(content="", success=False, error=str(e)) def generate_summary(self, code_context: str) -> AnalysisResult: """Generate a summary of the codebase. Args: code_context: Formatted code content Returns: AnalysisResult with summary or error """ user_prompt = f"""Analyze this codebase: {code_context} Provide a concise summary.""" try: logger.info(f"Generating summary with {self.model_name}") content = self._generate_with_llamaindex(SUMMARY_PROMPT, user_prompt) return AnalysisResult(content=content.strip()) except Exception as e: logger.exception("Summary generation failed") return AnalysisResult(content="", success=False, error=str(e)) def chat(self, message: str, code_context: str, history: Optional[List[Dict]] = None) -> AnalysisResult: """Chat about the codebase. Args: message: User's question code_context: Formatted code content (or DOT diagram) history: Previous chat messages Returns: AnalysisResult with response or error """ # Build context from history history_text = "" if history: for msg in history[-6:]: # Last 3 exchanges if isinstance(msg, dict): role = "User" if msg.get("role") == "user" else "Assistant" content = msg.get("content", "") if content: history_text += f"{role}: {content}\n" user_prompt = f"""Code context: {code_context} {history_text} Current question: {message}""" try: content = self._generate_with_llamaindex(CHAT_PROMPT, user_prompt) return AnalysisResult(content=content.strip()) except Exception as e: logger.exception("Chat failed") return AnalysisResult(content="", success=False, error=str(e))