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

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

@@ -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()
 
@@ -89,9 +90,14 @@ 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")
@@ -157,10 +163,15 @@ 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
@@ -528,10 +539,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:
@@ -1491,21 +1507,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:
@@ -1565,13 +1580,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
     
@@ -1580,10 +1597,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")
@@ -1593,15 +1615,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
     
@@ -1617,12 +1641,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}")
@@ -1638,29 +1670,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/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",
+]