cognify-code 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

ai_code_assistant/agent/code_agent.py

@@ -9,6 +9,7 @@ from ai_code_assistant.agent.intent_classifier import IntentClassifier, Intent,
 from ai_code_assistant.agent.code_generator import CodeGenerator, CodeGenerationRequest, GeneratedCode
 from ai_code_assistant.agent.diff_engine import DiffEngine, ChangeSet, FileDiff
 from ai_code_assistant.agent.code_reviewer import CodeReviewer, ReviewResult
+from ai_code_assistant.context import ContextSelector, ContextConfig
 
 
 @dataclass
@@ -30,7 +31,7 @@ class AgentResponse:
 class CodeAgent:
     """Main agent that orchestrates code operations based on user intent."""
     
-    def __init__(self, llm_manager, root_path: Optional[Path] = None):
+    def __init__(self, llm_manager, root_path: Optional[Path] = None, auto_context: bool = True):
         self.llm = llm_manager
         self.file_manager = FileContextManager(root_path)
         self.intent_classifier = IntentClassifier(llm_manager)
@@ -38,6 +39,13 @@ class CodeAgent:
         self.diff_engine = DiffEngine(self.file_manager)
         self.code_reviewer = CodeReviewer(llm_manager, self.file_manager)
         
+        # Context selector for smart context gathering
+        self.auto_context = auto_context
+        self.context_selector = ContextSelector(
+            root_path=root_path,
+            config=ContextConfig(max_tokens=8000, max_files=10)
+        )
+        
         # Pending changes awaiting confirmation
         self._pending_changeset: Optional[ChangeSet] = None
     
@@ -84,6 +92,28 @@ class CodeAgent:
         self._pending_changeset = None
         return "Changes discarded."
     
+    def _get_relevant_context(self, query: str, file_path: Optional[Path] = None) -> str:
+        """Get relevant context for a query using the context selector."""
+        if not self.auto_context:
+            return ""
+        
+        try:
+            # Get context based on query and optional file path
+            target_file = str(file_path) if file_path else None
+            context_result = self.context_selector.select_for_query(
+                query=query,
+                target_file=target_file
+            )
+            
+            if not context_result.files:
+                return ""
+            
+            # Use the built-in formatting
+            return context_result.format_for_prompt(include_summary=False)
+        except Exception:
+            # If context gathering fails, continue without it
+            return ""
+    
     def _handle_generate(self, message: str, intent: Intent) -> AgentResponse:
         """Handle code generation requests."""
         request = CodeGenerationRequest(
@@ -778,339 +808,6 @@ Provide a helpful, concise response.
         response = AgentResponse(message=full_response, intent=intent)
         yield ("", response)
 
-
-
-    def process_stream(self, message: str, use_llm_classification: bool = True) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """
-        Process a user message with streaming output.
-        
-        Yields tuples of (chunk, final_response).
-        During streaming, final_response is None.
-        The last yield will have the complete AgentResponse.
-        """
-        # Classify intent (non-streaming, it's fast)
-        if use_llm_classification:
-            intent = self.intent_classifier.classify_with_llm(message)
-        else:
-            intent = self.intent_classifier.classify(message)
-        
-        # Route to appropriate streaming handler
-        streaming_handlers = {
-            IntentType.CODE_GENERATE: self._handle_generate_stream,
-            IntentType.CODE_EDIT: self._handle_edit_stream,
-            IntentType.CODE_REVIEW: self._handle_review_stream,
-            IntentType.CODE_EXPLAIN: self._handle_explain_stream,
-            IntentType.CODE_REFACTOR: self._handle_refactor_stream,
-            IntentType.GENERAL_CHAT: self._handle_general_chat_stream,
-        }
-        
-        handler = streaming_handlers.get(intent.type)
-        
-        if handler:
-            yield from handler(message, intent)
-        else:
-            # Fall back to non-streaming for other intents
-            response = self.process(message, use_llm_classification)
-            yield (response.message, response)
-    
-    def _handle_explain_stream(self, message: str, intent: Intent) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """Handle code explanation with streaming."""
-        if not intent.file_paths:
-            response = AgentResponse(
-                message="Please specify which file or code you want me to explain.",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        file_path = intent.file_paths[0]
-        content = self.file_manager.read_file(file_path)
-        
-        if not content:
-            response = AgentResponse(
-                message=f"Cannot find file: {file_path}",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        prompt = f"""Explain the following code in a clear, educational way.
-
-## Code ({file_path})
-```
-{content[:5000]}
-```
-
-## Instructions
-1. Start with a high-level overview
-2. Explain the main components/functions
-3. Describe the flow of execution
-4. Note any important patterns or techniques used
-5. Keep the explanation concise but thorough
-"""
-        
-        # Stream the explanation
-        full_response = f"📖 **Explanation of {file_path}**\n\n"
-        yield (f"📖 **Explanation of {file_path}**\n\n", None)
-        
-        for chunk in self.llm.stream(prompt):
-            full_response += chunk
-            yield (chunk, None)
-        
-        # Final response
-        response = AgentResponse(
-            message=full_response,
-            intent=intent,
-        )
-        yield ("", response)
-    
-    def _handle_review_stream(self, message: str, intent: Intent) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """Handle code review with streaming."""
-        if not intent.file_paths:
-            context = self.file_manager.get_project_context()
-            py_files = [f.relative_path for f in context.files if f.extension == ".py"][:5]
-            
-            if py_files:
-                msg = f"Which file would you like me to review? Found these Python files:\n" + \
-                      "\n".join(f"  • {f}" for f in py_files)
-            else:
-                msg = "Please specify which file you want me to review."
-            
-            response = AgentResponse(message=msg, intent=intent)
-            yield (msg, response)
-            return
-        
-        file_path = intent.file_paths[0]
-        content = self.file_manager.read_file(file_path)
-        
-        if not content:
-            response = AgentResponse(
-                message=f"Cannot find file: {file_path}",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        yield (f"🔍 **Reviewing {file_path}...**\n\n", None)
-        
-        # Use streaming for the review
-        prompt = f"""Review the following code for issues, bugs, and improvements.
-
-## Code ({file_path})
-```
-{content[:5000]}
-```
-
-## Review Format
-Provide a structured review with:
-1. **Summary** - Brief overview
-2. **Issues** - List any bugs, security issues, or problems
-3. **Suggestions** - Improvements and best practices
-4. **Score** - Rate the code quality (1-10)
-
-Be specific and actionable.
-"""
-        
-        full_response = f"🔍 **Reviewing {file_path}...**\n\n"
-        
-        for chunk in self.llm.stream(prompt):
-            full_response += chunk
-            yield (chunk, None)
-        
-        response = AgentResponse(
-            message=full_response,
-            intent=intent,
-        )
-        yield ("", response)
-    
-    def _handle_generate_stream(self, message: str, intent: Intent) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """Handle code generation with streaming."""
-        yield ("🔨 **Generating code...**\n\n", None)
-        
-        request = CodeGenerationRequest(
-            description=message,
-            language=intent.language,
-            file_path=intent.file_paths[0] if intent.file_paths else None,
-        )
-        
-        # Generate code (this part streams)
-        full_code = ""
-        prompt = self.code_generator._build_prompt(request)
-        
-        for chunk in self.llm.stream(prompt):
-            full_code += chunk
-            yield (chunk, None)
-        
-        # Extract and create changeset
-        code = self._extract_code(full_code)
-        file_path = request.file_path or f"generated.{request.language or 'py'}"
-        
-        generated = GeneratedCode(
-            code=code,
-            language=request.language or "python",
-            file_path=file_path,
-            description=request.description,
-        )
-        
-        changeset = ChangeSet(description=f"Generate: {message[:50]}...")
-        diff = self.diff_engine.create_file_diff(generated.file_path, generated.code)
-        changeset.diffs.append(diff)
-        
-        self._pending_changeset = changeset
-        
-        response = AgentResponse(
-            message=f"\n\n✅ Code generated for {file_path}",
-            intent=intent,
-            generated_code=generated,
-            changeset=changeset,
-            requires_confirmation=True,
-        )
-        yield ("\n\n✅ Code generated. Apply changes? (yes/no)", response)
-    
-    def _handle_edit_stream(self, message: str, intent: Intent) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """Handle code editing with streaming."""
-        if not intent.file_paths:
-            response = AgentResponse(
-                message="Please specify which file you want to edit.",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        file_path = intent.file_paths[0]
-        original = self.file_manager.read_file(file_path)
-        
-        if not original:
-            response = AgentResponse(
-                message=f"Cannot find file: {file_path}",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        yield (f"✏️ **Editing {file_path}...**\n\n", None)
-        
-        prompt = f"""Edit the following code according to the user's request.
-
-## Original Code ({file_path})
-```
-{original[:5000]}
-```
-
-## User Request
-{message}
-
-## Instructions
-Return the COMPLETE modified file.
-
-```
-"""
-        
-        full_response = ""
-        for chunk in self.llm.stream(prompt):
-            full_response += chunk
-            yield (chunk, None)
-        
-        new_code = self._extract_code(full_response)
-        
-        changeset = ChangeSet(description=f"Edit: {message[:50]}...")
-        diff = self.diff_engine.create_diff(original, new_code, file_path)
-        changeset.diffs.append(diff)
-        
-        self._pending_changeset = changeset
-        
-        response = AgentResponse(
-            message=f"\n\n✅ Edit complete for {file_path}",
-            intent=intent,
-            changeset=changeset,
-            requires_confirmation=True,
-        )
-        yield ("\n\n✅ Edit complete. Apply changes? (yes/no)", response)
-    
-    def _handle_refactor_stream(self, message: str, intent: Intent) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """Handle code refactoring with streaming."""
-        if not intent.file_paths:
-            response = AgentResponse(
-                message="Please specify which file you want to refactor.",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        file_path = intent.file_paths[0]
-        original = self.file_manager.read_file(file_path)
-        
-        if not original:
-            response = AgentResponse(
-                message=f"Cannot find file: {file_path}",
-                intent=intent,
-            )
-            yield (response.message, response)
-            return
-        
-        yield (f"🔄 **Refactoring {file_path}...**\n\n", None)
-        
-        prompt = f"""Refactor the following code to improve its quality.
-
-## Original Code ({file_path})
-```
-{original[:5000]}
-```
-
-## User Request
-{message}
-
-Return the COMPLETE refactored file.
-
-```
-"""
-        
-        full_response = ""
-        for chunk in self.llm.stream(prompt):
-            full_response += chunk
-            yield (chunk, None)
-        
-        new_code = self._extract_code(full_response)
-        
-        changeset = ChangeSet(description=f"Refactor: {file_path}")
-        diff = self.diff_engine.create_diff(original, new_code, file_path)
-        changeset.diffs.append(diff)
-        
-        self._pending_changeset = changeset
-        
-        response = AgentResponse(
-            message=f"\n\n✅ Refactoring complete for {file_path}",
-            intent=intent,
-            changeset=changeset,
-            requires_confirmation=True,
-        )
-        yield ("\n\n✅ Refactoring complete. Apply changes? (yes/no)", response)
-    
-    def _handle_general_chat_stream(self, message: str, intent: Intent) -> Iterator[Tuple[str, Optional[AgentResponse]]]:
-        """Handle general chat with streaming."""
-        context = self.file_manager.get_project_context()
-        
-        prompt = f"""You are a helpful coding assistant. Answer the user's question.
-
-Project context:
-- Root: {context.root_path.name}
-- Languages: {', '.join(context.languages)}
-- Files: {context.total_code_files} code files
-
-User: {message}
-
-Provide a helpful, concise response.
-"""
-        
-        full_response = ""
-        for chunk in self.llm.stream(prompt):
-            full_response += chunk
-            yield (chunk, None)
-        
-        response = AgentResponse(message=full_response, intent=intent)
-        yield ("", response)
-
-
     def _extract_code(self, response: str) -> str:
         """Extract code from LLM response."""
         import re

ai_code_assistant/cli.py

@@ -17,6 +17,7 @@ from ai_code_assistant.generator import CodeGenerator
 from ai_code_assistant.chat import ChatSession
 from ai_code_assistant.editor import FileEditor
 from ai_code_assistant.utils import FileHandler, get_formatter
+from ai_code_assistant.context import ContextSelector, ContextConfig
 
 console = Console()
 
@@ -65,6 +66,59 @@ def get_components(config_path: Optional[Path] = None):
     return config, llm
 
 
+def validate_llm_connection(llm: LLMManager, exit_on_failure: bool = True) -> bool:
+    """
+    Validate LLM connection before operations.
+
+    Args:
+        llm: The LLM manager instance
+        exit_on_failure: If True, exit with error code on failure
+
+    Returns:
+        True if connection is valid, False otherwise
+    """
+    try:
+        info = llm.get_model_info()
+        provider = info.get("provider", "ollama")
+
+        # Quick validation - don't do full connection check for cloud providers with API keys
+        if provider != "ollama" and llm.config.llm.api_key:
+            return True
+
+        # For Ollama, check connection
+        if provider == "ollama":
+            # Try a lightweight check first
+            import socket
+            base_url = info.get("base_url", "http://localhost:11434")
+            host = base_url.replace("http://", "").replace("https://", "").split(":")[0]
+            port = int(base_url.split(":")[-1].split("/")[0]) if ":" in base_url else 11434
+
+            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            sock.settimeout(2)
+            result = sock.connect_ex((host, port))
+            sock.close()
+
+            if result != 0:
+                console.print(f"\n[red]Error: Cannot connect to Ollama at {base_url}[/red]")
+                console.print("\n[yellow]Quick fix:[/yellow]")
+                console.print("  1. Make sure Ollama is installed: [cyan]https://ollama.ai[/cyan]")
+                console.print(f"  2. Pull the model: [cyan]ollama pull {info.get('model', 'deepseek-coder:6.7b')}[/cyan]")
+                console.print("  3. Start Ollama: [cyan]ollama serve[/cyan]")
+                console.print("\n[dim]Or use a cloud provider: cognify status --help[/dim]")
+
+                if exit_on_failure:
+                    sys.exit(1)
+                return False
+
+        return True
+
+    except Exception as e:
+        console.print(f"\n[red]Error validating LLM connection: {e}[/red]")
+        if exit_on_failure:
+            sys.exit(1)
+        return False
+
+
 @click.group(invoke_without_command=True)
 @click.version_option(version=__version__, prog_name="cognify")
 @click.option("--config", "-c", type=click.Path(exists=True, path_type=Path), help="Config file path")
@@ -89,15 +143,21 @@ def main(ctx, config: Optional[Path], verbose: bool):
               type=click.Choice(["console", "markdown", "json"]), help="Output format")
 @click.option("--output", "-o", type=click.Path(path_type=Path), help="Output file path")
 @click.option("--recursive", "-r", is_flag=True, help="Recursively review directories")
+@click.option("--context", multiple=True, type=click.Path(exists=True, path_type=Path),
+              help="Additional context files to include")
+@click.option("--auto-context", is_flag=True, help="Automatically include related files as context")
+@click.option("--max-context-tokens", type=int, default=8000, help="Max tokens for context")
 @click.pass_context
 def review(ctx, files: Tuple[Path, ...], review_type: str, output_format: str,
-           output: Optional[Path], recursive: bool):
+           output: Optional[Path], recursive: bool, context: Tuple[Path, ...],
+           auto_context: bool, max_context_tokens: int):
     """Review code files for issues and improvements."""
     if not files:
         console.print("[red]Error:[/red] No files specified")
         sys.exit(1)
 
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     analyzer = CodeAnalyzer(config, llm)
     file_handler = FileHandler(config)
     formatter = get_formatter(output_format, config.output.use_colors)
@@ -157,16 +217,22 @@ def review(ctx, files: Tuple[Path, ...], review_type: str, output_format: str,
 @click.option("--source", "-s", type=click.Path(exists=True, path_type=Path),
               help="Source file (for test mode)")
 @click.option("--stream/--no-stream", default=True, help="Stream output in real-time")
+@click.option("--context", multiple=True, type=click.Path(exists=True, path_type=Path),
+              help="Context files to include for better generation")
+@click.option("--auto-context", is_flag=True, help="Automatically find relevant context files")
+@click.option("--max-context-tokens", type=int, default=8000, help="Max tokens for context")
 @click.pass_context
 def generate(ctx, description: str, mode: str, language: str, name: Optional[str],
              params: Optional[str], output: Optional[Path], output_format: str,
-             source: Optional[Path], stream: bool):
+             source: Optional[Path], stream: bool, context: Tuple[Path, ...],
+             auto_context: bool, max_context_tokens: int):
     """Generate code from natural language description."""
     from rich.live import Live
     from rich.markdown import Markdown
     from rich.panel import Panel
     
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     generator = CodeGenerator(config, llm)
     formatter = get_formatter(output_format, config.output.use_colors)
 
@@ -256,6 +322,7 @@ def generate(ctx, description: str, mode: str, language: str, name: Optional[str
 def chat(ctx, context: Tuple[Path, ...], stream: bool):
     """Start an interactive chat session about code."""
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     session = ChatSession(config, llm)
 
     # Load context files
@@ -528,10 +595,15 @@ def search(ctx, query: str, top_k: int, file_filter: Optional[str],
               type=click.Choice(["console", "json"]), help="Output format")
 @click.option("--start-line", "-s", type=int, help="Start line for targeted edit")
 @click.option("--end-line", "-e", type=int, help="End line for targeted edit")
+@click.option("--context", multiple=True, type=click.Path(exists=True, path_type=Path),
+              help="Additional context files to include")
+@click.option("--auto-context", is_flag=True, help="Automatically include related files as context")
+@click.option("--max-context-tokens", type=int, default=8000, help="Max tokens for context")
 @click.pass_context
 def edit(ctx, file: Path, instruction: str, mode: str, preview: bool,
          no_backup: bool, output_format: str, start_line: Optional[int],
-         end_line: Optional[int]):
+         end_line: Optional[int], context: Tuple[Path, ...], auto_context: bool,
+         max_context_tokens: int):
     """Edit a file using AI based on natural language instructions.
 
     Examples:
@@ -542,6 +614,7 @@ def edit(ctx, file: Path, instruction: str, mode: str, preview: bool,
         ai-assist edit config.py "Update the timeout value" -s 10 -e 20
     """
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     editor = FileEditor(config, llm)
 
     # Determine edit mode
@@ -631,6 +704,7 @@ def refactor(ctx, instruction: str, files: Tuple[Path, ...], pattern: Optional[s
     from ai_code_assistant.utils import FileHandler
 
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     editor = MultiFileEditor(config, llm)
     file_handler = FileHandler(config)
 
@@ -772,6 +846,7 @@ def rename(ctx, old_name: str, new_name: str, symbol_type: str, files: Tuple[Pat
     from ai_code_assistant.utils import FileHandler
 
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     editor = MultiFileEditor(config, llm)
     file_handler = FileHandler(config)
 
@@ -1166,6 +1241,7 @@ def git_commit(ctx, message: Optional[str], stage_all: bool, push_after: bool, n
     from ai_code_assistant.git import GitManager, CommitMessageGenerator
 
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
 
     try:
         git_mgr = GitManager()
@@ -1299,6 +1375,7 @@ def git_sync(ctx, message: Optional[str], no_confirm: bool):
     from ai_code_assistant.git import GitManager, CommitMessageGenerator
 
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
 
     try:
         git_mgr = GitManager()
@@ -1435,9 +1512,10 @@ def agent(ctx, path: Path):
     from ai_code_assistant.chat import AgentChatSession
     from rich.markdown import Markdown
     from rich.prompt import Prompt
-    
+
     config, llm = get_components(ctx.obj.get("config_path"))
-    
+    validate_llm_connection(llm)
+
     # Initialize agent session
     session = AgentChatSession(config, llm, path.resolve())
     
@@ -1575,10 +1653,11 @@ def agent_review(ctx, file: Path, path: Path, stream: bool):
         ai-assist agent-review main.py --no-stream
     """
     from ai_code_assistant.agent import CodeAgent
-    
+
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     agent = CodeAgent(llm, path.resolve())
-    
+
     console.print(f"\n[bold]Reviewing {file}...[/bold]\n")
     
     if stream:
@@ -1612,10 +1691,11 @@ def agent_generate(ctx, description: str, file: Optional[Path], language: Option
         ai-assist agent-generate "hello world" --no-stream
     """
     from ai_code_assistant.agent import CodeAgent
-    
+
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     agent = CodeAgent(llm, path.resolve())
-    
+
     # Build the request
     request = description
     if file:
@@ -1665,10 +1745,11 @@ def agent_explain(ctx, file: Path, path: Path, stream: bool):
         ai-assist agent-explain main.py --no-stream
     """
     from ai_code_assistant.agent import CodeAgent
-    
+
     config, llm = get_components(ctx.obj.get("config_path"))
+    validate_llm_connection(llm)
     agent = CodeAgent(llm, path.resolve())
-    
+
     console.print(f"\n[bold]Explaining {file}...[/bold]\n")
     
     if stream:

ai_code_assistant/context/__init__.py

@@ -0,0 +1,12 @@
+"""Context-aware codebase understanding module."""
+
+from .analyzer import ContextAnalyzer, FileContext
+from .selector import ContextSelector, ContextConfig, ContextResult
+
+__all__ = [
+    "ContextAnalyzer",
+    "ContextSelector", 
+    "ContextConfig",
+    "ContextResult",
+    "FileContext",
+]