cognify-code 0.2.3__py3-none-any.whl → 0.2.5__py3-none-any.whl

ai_code_assistant/agent/code_agent.py

@@ -2,13 +2,14 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, List, Optional, Tuple
+from typing import Callable, Iterator, List, Optional, Tuple
 
 from ai_code_assistant.agent.file_manager import FileContextManager
 from ai_code_assistant.agent.intent_classifier import IntentClassifier, Intent, IntentType
 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(
@@ -447,6 +477,670 @@ Provide a helpful, concise response. If the question is about code, you can sugg
         
         return AgentResponse(message=response, intent=intent)
     
+
+    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 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