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

ai_code_assistant/agent/code_agent.py

@@ -2,7 +2,7 @@
 
 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
@@ -447,6 +447,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

ai_code_assistant/chat/agent_session.py

@@ -3,7 +3,7 @@
 from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
-from typing import Dict, List, Literal, Optional
+from typing import Dict, Iterator, List, Literal, Optional, Tuple
 
 from ai_code_assistant.config import Config
 from ai_code_assistant.llm import LLMManager
@@ -67,6 +67,96 @@ class AgentChatSession:
         
         return assistant_msg
     
+
+    def send_message_stream(self, user_input: str) -> Iterator[Tuple[str, Optional[AgentMessage]]]:
+        """
+        Process user message with streaming output.
+        
+        Yields tuples of (chunk, final_message).
+        During streaming, final_message is None.
+        The last yield will have the complete AgentMessage.
+        """
+        # Add user message to history
+        user_msg = AgentMessage(role="user", content=user_input)
+        self.history.append(user_msg)
+        
+        # Check for confirmation/rejection of pending changes
+        if self._awaiting_confirmation:
+            msg = self._handle_confirmation(user_input)
+            yield (msg.content, msg)
+            return
+        
+        # Process through agent with streaming
+        full_content = ""
+        final_response = None
+        
+        for chunk, response in self.agent.process_stream(user_input):
+            full_content += chunk
+            yield (chunk, None)
+            if response is not None:
+                final_response = response
+        
+        # Create assistant message
+        assistant_msg = AgentMessage(
+            role="assistant",
+            content=full_content,
+            response=final_response,
+            pending_action=final_response.requires_confirmation if final_response else False,
+        )
+        self.history.append(assistant_msg)
+        
+        # Track if we're awaiting confirmation
+        if final_response:
+            self._awaiting_confirmation = final_response.requires_confirmation
+        
+        yield ("", assistant_msg)
+
+
+
+    def send_message_stream(self, user_input: str) -> Iterator[Tuple[str, Optional[AgentMessage]]]:
+        """
+        Process user message with streaming output.
+        
+        Yields tuples of (chunk, final_message).
+        During streaming, final_message is None.
+        The last yield will have the complete AgentMessage.
+        """
+        # Add user message to history
+        user_msg = AgentMessage(role="user", content=user_input)
+        self.history.append(user_msg)
+        
+        # Check for confirmation/rejection of pending changes
+        if self._awaiting_confirmation:
+            msg = self._handle_confirmation(user_input)
+            yield (msg.content, msg)
+            return
+        
+        # Process through agent with streaming
+        full_content = ""
+        final_response = None
+        
+        for chunk, response in self.agent.process_stream(user_input):
+            full_content += chunk
+            yield (chunk, None)
+            if response is not None:
+                final_response = response
+        
+        # Create assistant message
+        assistant_msg = AgentMessage(
+            role="assistant",
+            content=full_content,
+            response=final_response,
+            pending_action=final_response.requires_confirmation if final_response else False,
+        )
+        self.history.append(assistant_msg)
+        
+        # Track if we're awaiting confirmation
+        if final_response:
+            self._awaiting_confirmation = final_response.requires_confirmation
+        
+        yield ("", assistant_msg)
+
+
     def _handle_confirmation(self, user_input: str) -> AgentMessage:
         """Handle user confirmation or rejection of pending changes."""
         lower_input = user_input.lower().strip()

ai_code_assistant/cli.py

@@ -156,41 +156,91 @@ def review(ctx, files: Tuple[Path, ...], review_type: str, output_format: str,
               type=click.Choice(["console", "markdown", "json"]))
 @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.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]):
+             source: Optional[Path], stream: bool):
     """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"))
     generator = CodeGenerator(config, llm)
     formatter = get_formatter(output_format, config.output.use_colors)
 
     console.print(f"\n[bold]Generating {mode} in {language}...[/bold]\n")
 
-    with console.status("[bold green]Generating code..."):
-        if mode == "function":
-            result = generator.generate_function(
-                description=description, name=name or "generated_function",
-                language=language, parameters=params or "",
-            )
-        elif mode == "class":
-            result = generator.generate_class(
-                description=description, name=name or "GeneratedClass", language=language,
-            )
-        elif mode == "script":
-            result = generator.generate_script(
-                description=description, requirements=[description], language=language,
-            )
-        elif mode == "test":
-            if not source:
-                console.print("[red]Error:[/red] --source required for test mode")
-                sys.exit(1)
-            source_code = source.read_text()
-            result = generator.generate_tests(source_code=source_code, language=language)
-        else:
-            result = generator.generate(description=description, language=language)
+    # Handle test mode source requirement
+    source_code = ""
+    if mode == "test":
+        if not source:
+            console.print("[red]Error:[/red] --source required for test mode")
+            sys.exit(1)
+        source_code = source.read_text()
 
-    formatted = formatter.format_generation(result)
+    if stream:
+        # Streaming mode - show output as it generates
+        full_response = ""
+        final_code = ""
+        
+        console.print("[dim]Streaming response...[/dim]\n")
+        
+        for chunk, is_complete in generator.generate_stream(
+            description=description,
+            mode=mode,
+            language=language,
+            name=name or "",
+            parameters=params or "",
+            source_code=source_code,
+        ):
+            if is_complete:
+                final_code = chunk
+            else:
+                console.print(chunk, end="", highlight=False)
+                full_response += chunk
+        
+        console.print("\n")
+        
+        # Create result for formatting
+        from ai_code_assistant.generator import GenerationResult
+        result = GenerationResult(
+            code=final_code,
+            language=language,
+            mode=mode,
+            description=description,
+            raw_response=full_response,
+        )
+        
+        # Show extracted code in a panel
+        console.print(Panel(
+            final_code,
+            title=f"[bold green]Generated {mode.title()}[/bold green]",
+            border_style="green",
+        ))
+    else:
+        # Non-streaming mode (original behavior)
+        with console.status("[bold green]Generating code..."):
+            if mode == "function":
+                result = generator.generate_function(
+                    description=description, name=name or "generated_function",
+                    language=language, parameters=params or "",
+                )
+            elif mode == "class":
+                result = generator.generate_class(
+                    description=description, name=name or "GeneratedClass", language=language,
+                )
+            elif mode == "script":
+                result = generator.generate_script(
+                    description=description, requirements=[description], language=language,
+                )
+            elif mode == "test":
+                result = generator.generate_tests(source_code=source_code, language=language)
+            else:
+                result = generator.generate(description=description, language=language)
+
+        formatted = formatter.format_generation(result)
 
     if output and result.success:
         output.parent.mkdir(parents=True, exist_ok=True)
@@ -1441,21 +1491,20 @@ def agent(ctx, path: Path):
                     console.print("[green]No pending changes.[/green]")
                 continue
             
-            # Process through agent
-            with console.status("[bold green]Thinking..."):
-                response = session.send_message(user_input)
-            
-            # Display response
+            # Process through agent with streaming
             console.print(f"\n[bold green]Agent[/bold green]")
             
-            # Use markdown rendering for better formatting
-            try:
-                console.print(Markdown(response.content))
-            except Exception:
-                console.print(response.content)
+            final_msg = None
+            for chunk, msg in session.send_message_stream(user_input):
+                if chunk:
+                    console.print(chunk, end="")
+                if msg is not None:
+                    final_msg = msg
+            
+            console.print()  # Newline after streaming
             
             # Show confirmation prompt if needed
-            if response.pending_action:
+            if final_msg and final_msg.pending_action:
                 console.print("\n[yellow]Apply these changes? (yes/no)[/yellow]")
         
         except KeyboardInterrupt:
@@ -1515,13 +1564,15 @@ def _show_agent_help():
 @click.argument("file", type=click.Path(exists=True, path_type=Path))
 @click.option("--path", "-p", type=click.Path(exists=True, path_type=Path),
               default=".", help="Project root path")
+@click.option("--stream/--no-stream", default=True, help="Stream output in real-time")
 @click.pass_context
-def agent_review(ctx, file: Path, path: Path):
+def agent_review(ctx, file: Path, path: Path, stream: bool):
     """Quick code review using the agent.
     
     Examples:
         ai-assist agent-review src/main.py
         ai-assist agent-review utils.py --path ./my-project
+        ai-assist agent-review main.py --no-stream
     """
     from ai_code_assistant.agent import CodeAgent
     
@@ -1530,10 +1581,15 @@ def agent_review(ctx, file: Path, path: Path):
     
     console.print(f"\n[bold]Reviewing {file}...[/bold]\n")
     
-    with console.status("[bold green]Analyzing..."):
-        response = agent.process(f"review {file}")
-    
-    console.print(response.message)
+    if stream:
+        for chunk, response in agent.process_stream(f"review {file}"):
+            if chunk:
+                console.print(chunk, end="")
+        console.print()
+    else:
+        with console.status("[bold green]Analyzing..."):
+            response = agent.process(f"review {file}")
+        console.print(response.message)
 
 
 @main.command("agent-generate")
@@ -1543,15 +1599,17 @@ def agent_review(ctx, file: Path, path: Path):
 @click.option("--path", "-p", type=click.Path(exists=True, path_type=Path),
               default=".", help="Project root path")
 @click.option("--apply", "-a", is_flag=True, help="Apply changes without confirmation")
+@click.option("--stream/--no-stream", default=True, help="Stream output in real-time")
 @click.pass_context
 def agent_generate(ctx, description: str, file: Optional[Path], language: Optional[str],
-                   path: Path, apply: bool):
+                   path: Path, apply: bool, stream: bool):
     """Generate code using the agent.
     
     Examples:
         ai-assist agent-generate "a function to validate email"
         ai-assist agent-generate "REST API for users" -f src/api.py
         ai-assist agent-generate "sorting algorithm" -l python
+        ai-assist agent-generate "hello world" --no-stream
     """
     from ai_code_assistant.agent import CodeAgent
     
@@ -1567,12 +1625,20 @@ def agent_generate(ctx, description: str, file: Optional[Path], language: Option
     
     console.print(f"\n[bold]Generating code...[/bold]\n")
     
-    with console.status("[bold green]Generating..."):
-        response = agent.process(request)
-    
-    console.print(response.message)
+    final_response = None
+    if stream:
+        for chunk, response in agent.process_stream(request):
+            if chunk:
+                console.print(chunk, end="")
+            if response is not None:
+                final_response = response
+        console.print()
+    else:
+        with console.status("[bold green]Generating..."):
+            final_response = agent.process(request)
+        console.print(final_response.message)
     
-    if response.requires_confirmation:
+    if final_response and final_response.requires_confirmation:
         if apply:
             success, msg = agent.confirm_changes()
             console.print(f"\n{msg}")
@@ -1588,29 +1654,36 @@ def agent_generate(ctx, description: str, file: Optional[Path], language: Option
 @click.argument("file", type=click.Path(exists=True, path_type=Path))
 @click.option("--path", "-p", type=click.Path(exists=True, path_type=Path),
               default=".", help="Project root path")
+@click.option("--stream/--no-stream", default=True, help="Stream output in real-time")
 @click.pass_context
-def agent_explain(ctx, file: Path, path: Path):
+def agent_explain(ctx, file: Path, path: Path, stream: bool):
     """Explain code using the agent.
     
     Examples:
         ai-assist agent-explain src/main.py
         ai-assist agent-explain config.py --path ./my-project
+        ai-assist agent-explain main.py --no-stream
     """
     from ai_code_assistant.agent import CodeAgent
-    from rich.markdown import Markdown
     
     config, llm = get_components(ctx.obj.get("config_path"))
     agent = CodeAgent(llm, path.resolve())
     
     console.print(f"\n[bold]Explaining {file}...[/bold]\n")
     
-    with console.status("[bold green]Analyzing..."):
-        response = agent.process(f"explain {file}")
-    
-    try:
-        console.print(Markdown(response.message))
-    except Exception:
-        console.print(response.message)
+    if stream:
+        for chunk, response in agent.process_stream(f"explain {file}"):
+            if chunk:
+                console.print(chunk, end="")
+        console.print()
+    else:
+        from rich.markdown import Markdown
+        with console.status("[bold green]Analyzing..."):
+            response = agent.process(f"explain {file}")
+        try:
+            console.print(Markdown(response.message))
+        except Exception:
+            console.print(response.message)
 
 
 if __name__ == "__main__":

ai_code_assistant/generator/code_gen.py

@@ -3,7 +3,7 @@
 import re
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import List, Literal, Optional
+from typing import Iterator, List, Literal, Optional, Tuple
 
 from ai_code_assistant.config import Config
 from ai_code_assistant.llm import LLMManager
@@ -214,6 +214,160 @@ class CodeGenerator:
                 error=str(e),
             )
 
+
+    def generate_stream(
+        self,
+        description: str,
+        mode: str = "generic",
+        language: str = "python",
+        name: str = "",
+        parameters: str = "",
+        source_code: str = "",
+    ) -> Iterator[Tuple[str, bool]]:
+        """
+        Stream code generation with real-time output.
+        
+        Yields tuples of (chunk, is_complete).
+        The final yield will have is_complete=True and contain the extracted code.
+        """
+        # Build the prompt based on mode
+        if mode == "function":
+            prompt_template = GENERATION_PROMPTS["function"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                name=name or "generated_function",
+                parameters=parameters or "None specified",
+                return_type="appropriate type",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+        elif mode == "class":
+            prompt_template = GENERATION_PROMPTS["class"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                name=name or "GeneratedClass",
+                attributes="None specified",
+                methods="None specified",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+        elif mode == "script":
+            prompt_template = GENERATION_PROMPTS["script"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                requirements=f"- {description}",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+        elif mode == "test":
+            prompt_template = GENERATION_PROMPTS["test"]
+            formatted = prompt_template.format(
+                language=language,
+                source_code=source_code,
+                test_framework="pytest",
+            )
+        else:  # generic
+            prompt_template = GENERATION_PROMPTS["generic"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+
+        # Stream the response
+        full_response = ""
+        try:
+            for chunk in self.llm.stream(formatted):
+                full_response += chunk
+                yield (chunk, False)
+            
+            # Extract code and yield final result
+            code = self._extract_code(full_response, language)
+            yield (code, True)
+        except Exception as e:
+            yield (f"Error: {str(e)}", True)
+
+
+    def generate_stream(
+        self,
+        description: str,
+        mode: str = "generic",
+        language: str = "python",
+        name: str = "",
+        parameters: str = "",
+        source_code: str = "",
+    ) -> Iterator[Tuple[str, bool]]:
+        """
+        Stream code generation with real-time output.
+        
+        Yields tuples of (chunk, is_complete).
+        The final yield will have is_complete=True and contain the extracted code.
+        """
+        # Build the prompt based on mode
+        if mode == "function":
+            prompt_template = GENERATION_PROMPTS["function"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                name=name or "generated_function",
+                parameters=parameters or "None specified",
+                return_type="appropriate type",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+        elif mode == "class":
+            prompt_template = GENERATION_PROMPTS["class"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                name=name or "GeneratedClass",
+                attributes="None specified",
+                methods="None specified",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+        elif mode == "script":
+            prompt_template = GENERATION_PROMPTS["script"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                requirements=f"- {description}",
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+        elif mode == "test":
+            prompt_template = GENERATION_PROMPTS["test"]
+            formatted = prompt_template.format(
+                language=language,
+                source_code=source_code,
+                test_framework="pytest",
+            )
+        else:  # generic
+            prompt_template = GENERATION_PROMPTS["generic"]
+            formatted = prompt_template.format(
+                language=language,
+                description=description,
+                include_type_hints=self.config.generation.include_type_hints,
+                include_docstrings=self.config.generation.include_docstrings,
+            )
+
+        # Stream the response
+        full_response = ""
+        try:
+            for chunk in self.llm.stream(formatted):
+                full_response += chunk
+                yield (chunk, False)
+            
+            # Extract code and yield final result
+            code = self._extract_code(full_response, language)
+            yield (code, True)
+        except Exception as e:
+            yield (f"Error: {str(e)}", True)
+
     def _extract_code(self, response: str, language: str) -> str:
         """Extract code block from LLM response."""
         # Try to find language-specific code block

{cognify_code-0.2.2.dist-info → cognify_code-0.2.4.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: cognify-code
-Version: 0.2.2
+Version: 0.2.4
 Summary: Your local AI-powered code assistant. Review, generate, search, and refactor code with an intelligent AI agent—all running locally with complete privacy.
 Author-email: Ashok Kumar <akkssy@users.noreply.github.com>
 Maintainer-email: Ashok Kumar <akkssy@users.noreply.github.com>

{cognify_code-0.2.2.dist-info → cognify_code-0.2.4.dist-info}/RECORD

@@ -1,24 +1,24 @@
 ai_code_assistant/__init__.py,sha256=XnpG4h-2gW3cXseFvqQT_-XyOmVJtikVMrHUnmy8XKI,409
-ai_code_assistant/cli.py,sha256=47TI6XMkcgojYLSLADnk3BphxUffSgmV0tCTHIu4W7o,61350
+ai_code_assistant/cli.py,sha256=XSpwAQgWXhhLX9c628PwdSfx5czK3p5qS4nNNgESnWw,64153
 ai_code_assistant/config.py,sha256=6sAufexwzfCu2JNWvt9KevS9k_gMcjj1TAnwuaO1ZFw,4727
 ai_code_assistant/llm.py,sha256=DfcWJf6zEAUsPSEZLdEmb9o6BQNf1Ja88nswjpy6cOw,4209
 ai_code_assistant/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 ai_code_assistant/agent/__init__.py,sha256=BcVe4Ebopv_J01ApnRl05oN5yOet5mEefBrQmdPsUj0,1284
-ai_code_assistant/agent/code_agent.py,sha256=0_T3bFzPosMKKczYx49UGIWJtMBEuKy4ey-gT-qdfYQ,16155
+ai_code_assistant/agent/code_agent.py,sha256=y0Osc8wzpIAW1B1NGrlsSw_vIYu-ZlMOae05IhM2XYM,38463
 ai_code_assistant/agent/code_generator.py,sha256=rAaziRU-mJ5NooERjR_Cd6_hwO0kuULw3Sp8Ca9kR48,13138
 ai_code_assistant/agent/code_reviewer.py,sha256=YiM7lRJhoN-vBnQb29jF-5nmE9ppL-OJffvx4ocTHEU,12066
 ai_code_assistant/agent/diff_engine.py,sha256=A5jszowc5VmWbdidpIW_QhswG_Hats3FYuemP8VoYv4,11018
 ai_code_assistant/agent/file_manager.py,sha256=Inyfo-UXT4joms1ADIMA_TKtIHEjEkBz4I4U2iK5_jI,10742
 ai_code_assistant/agent/intent_classifier.py,sha256=MuIcyWQntocrTlCb4CD54mhc3JfSsWGfulsPYGUoz6E,10667
 ai_code_assistant/chat/__init__.py,sha256=KntIXcjbPgznax1E0fvdrA3XtKF-hCz5Fr1tcRbdl7U,279
-ai_code_assistant/chat/agent_session.py,sha256=VcutqSq8GKwldJhhAs_4_kXOuUvtnjRXZML1Mamm9SM,5495
+ai_code_assistant/chat/agent_session.py,sha256=-sW78d0nifRBNO6PRDiqdd8Sqpkv98kedvZbBQzK3lo,8674
 ai_code_assistant/chat/session.py,sha256=5JRd1DuLjxbtckmsMeHzNjoEZnJS9lx9NoX6z03F0xE,5500
 ai_code_assistant/editor/__init__.py,sha256=892BfTIo6kLdfZdhnvl4OFe0QSnxE4EyfkBoyLdA5rc,340
 ai_code_assistant/editor/diff_handler.py,sha256=LeI-00GuH7ASIetsUzT3Y_pDq4K1wmycuu4UFu5ZkGg,8759
 ai_code_assistant/editor/file_editor.py,sha256=csD8MW0jrfXAek5blWNuot_QWlhkgTTmtQtf8rbIdhY,11143
 ai_code_assistant/editor/prompts.py,sha256=wryxwb4dNaeSbl5mHkDuw2uTAZxkdrdyZ89Gfogafow,4356
 ai_code_assistant/generator/__init__.py,sha256=CfCO58CBye-BlZHjOfLLShovp2TVXg_GHKJuXe6ihu0,273
-ai_code_assistant/generator/code_gen.py,sha256=vJ4xeiG4_1hOZ8TW7YZ-gTAVKoiHwgUzzdcWMvhgeRo,8609
+ai_code_assistant/generator/code_gen.py,sha256=Sp_j1IdR0vbU5xRZHhvTBCbxfNxtvNMU1tEg0NLV00k,14790
 ai_code_assistant/generator/prompts.py,sha256=uoEDpcRzpTd-4TLHNW_EbSHJiADMlu9SoGWZvvo1Adk,3384
 ai_code_assistant/git/__init__.py,sha256=YgqmzneAnZyRrbazMqGoFSPIk5Yf5OTm2LXPbkQmecU,232
 ai_code_assistant/git/commit_generator.py,sha256=CzDH5ZPqEaXyPznBg8FgTz8wbV4adALUQD__kl8au6o,4135
@@ -47,9 +47,9 @@ ai_code_assistant/reviewer/prompts.py,sha256=9RrHEBttS5ngxY2BNsUvqGC6-cTxco-kDPb
 ai_code_assistant/utils/__init__.py,sha256=3HO-1Bj4VvUtM7W1C3MKR4DzQ9Xc875QKSHHkHwuqVs,368
 ai_code_assistant/utils/file_handler.py,sha256=jPxvtI5dJxkpPjELgRJ11WXamtyKKmZANQ1fcfMVtiU,5239
 ai_code_assistant/utils/formatters.py,sha256=5El9ew9HS6JLBucBUxxcw4fO5nLpOucgNJrJj2NC3zw,8945
-cognify_code-0.2.2.dist-info/licenses/LICENSE,sha256=5yu_kWq2bK-XKhWo79Eykdg4Qf3O8V2Ys7cpOO7GyyE,1063
-cognify_code-0.2.2.dist-info/METADATA,sha256=HRjZsc18DNcoS-NOJw9mCXB3qvrwuP1GhkgeLHmucOg,11862
-cognify_code-0.2.2.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
-cognify_code-0.2.2.dist-info/entry_points.txt,sha256=MrBnnWPHZVozqqKyTlnJO63YN2kE5yPWKlr2nnRFRks,94
-cognify_code-0.2.2.dist-info/top_level.txt,sha256=dD_r1x-oX0s1uspYY72kig4jfIsjh3oDKwOBCMYXqpo,18
-cognify_code-0.2.2.dist-info/RECORD,,
+cognify_code-0.2.4.dist-info/licenses/LICENSE,sha256=5yu_kWq2bK-XKhWo79Eykdg4Qf3O8V2Ys7cpOO7GyyE,1063
+cognify_code-0.2.4.dist-info/METADATA,sha256=_S3by3PNwTZmPdVbxZDBTLdKk7MZuoRttLOYZiXIYkI,11862
+cognify_code-0.2.4.dist-info/WHEEL,sha256=wUyA8OaulRlbfwMtmQsvNngGrxQHAvkKcvRmdizlJi0,92
+cognify_code-0.2.4.dist-info/entry_points.txt,sha256=MrBnnWPHZVozqqKyTlnJO63YN2kE5yPWKlr2nnRFRks,94
+cognify_code-0.2.4.dist-info/top_level.txt,sha256=dD_r1x-oX0s1uspYY72kig4jfIsjh3oDKwOBCMYXqpo,18
+cognify_code-0.2.4.dist-info/RECORD,,