cognify-code 0.2.0__py3-none-any.whl

ai_code_assistant/chat/agent_session.py

@@ -0,0 +1,156 @@
+"""Agent-enhanced chat session with code generation and review capabilities."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Literal, Optional
+
+from ai_code_assistant.config import Config
+from ai_code_assistant.llm import LLMManager
+from ai_code_assistant.agent import CodeAgent, AgentResponse, IntentType
+
+
+@dataclass
+class AgentMessage:
+    """A message in the agent chat session."""
+    role: Literal["user", "assistant", "system"]
+    content: str
+    timestamp: datetime = field(default_factory=datetime.now)
+    response: Optional[AgentResponse] = None
+    pending_action: bool = False
+
+
+class AgentChatSession:
+    """Chat session enhanced with code agent capabilities."""
+    
+    def __init__(
+        self,
+        config: Config,
+        llm_manager: LLMManager,
+        root_path: Optional[Path] = None,
+    ):
+        self.config = config
+        self.llm = llm_manager
+        self.agent = CodeAgent(llm_manager, root_path or Path.cwd())
+        self.history: List[AgentMessage] = []
+        self._awaiting_confirmation = False
+    
+    @property
+    def has_pending_changes(self) -> bool:
+        """Check if there are pending changes awaiting confirmation."""
+        return self._awaiting_confirmation
+    
+    def send_message(self, user_input: str) -> AgentMessage:
+        """Process user message through the agent."""
+        # 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:
+            return self._handle_confirmation(user_input)
+        
+        # Process through agent
+        response = self.agent.process(user_input)
+        
+        # Create assistant message
+        assistant_msg = AgentMessage(
+            role="assistant",
+            content=response.message,
+            response=response,
+            pending_action=response.requires_confirmation,
+        )
+        self.history.append(assistant_msg)
+        
+        # Track if we're awaiting confirmation
+        self._awaiting_confirmation = response.requires_confirmation
+        
+        return assistant_msg
+    
+    def _handle_confirmation(self, user_input: str) -> AgentMessage:
+        """Handle user confirmation or rejection of pending changes."""
+        lower_input = user_input.lower().strip()
+        
+        # Check for confirmation
+        if lower_input in ("yes", "y", "confirm", "apply", "ok", "sure", "do it"):
+            success, message = self.agent.confirm_changes()
+            self._awaiting_confirmation = False
+            
+            return AgentMessage(
+                role="assistant",
+                content=message,
+            )
+        
+        # Check for rejection
+        elif lower_input in ("no", "n", "cancel", "reject", "discard", "nevermind"):
+            message = self.agent.reject_changes()
+            self._awaiting_confirmation = False
+            
+            return AgentMessage(
+                role="assistant",
+                content=message,
+            )
+        
+        # Unclear response
+        else:
+            return AgentMessage(
+                role="assistant",
+                content="Please confirm with 'yes' to apply changes or 'no' to discard them.",
+                pending_action=True,
+            )
+    
+    def confirm_changes(self) -> str:
+        """Programmatically confirm pending changes."""
+        if not self._awaiting_confirmation:
+            return "No pending changes."
+        
+        success, message = self.agent.confirm_changes()
+        self._awaiting_confirmation = False
+        return message
+    
+    def reject_changes(self) -> str:
+        """Programmatically reject pending changes."""
+        if not self._awaiting_confirmation:
+            return "No pending changes."
+        
+        message = self.agent.reject_changes()
+        self._awaiting_confirmation = False
+        return message
+    
+    def get_project_info(self) -> str:
+        """Get information about the current project."""
+        response = self.agent.process("show project info")
+        return response.message
+    
+    def review_file(self, file_path: str) -> AgentMessage:
+        """Review a specific file."""
+        return self.send_message(f"review {file_path}")
+    
+    def generate_code(self, description: str, file_path: Optional[str] = None) -> AgentMessage:
+        """Generate code based on description."""
+        if file_path:
+            return self.send_message(f"create {file_path}: {description}")
+        return self.send_message(f"generate code: {description}")
+    
+    def explain_file(self, file_path: str) -> AgentMessage:
+        """Explain a file's code."""
+        return self.send_message(f"explain {file_path}")
+    
+    def get_history(self) -> List[AgentMessage]:
+        """Get conversation history."""
+        return self.history.copy()
+    
+    def clear_history(self) -> None:
+        """Clear conversation history."""
+        self.history.clear()
+        self._awaiting_confirmation = False
+    
+    def format_response(self, msg: AgentMessage) -> str:
+        """Format a message for display."""
+        lines = [msg.content]
+        
+        if msg.pending_action:
+            lines.append("")
+            lines.append("[yellow]Apply these changes? (yes/no)[/yellow]")
+        
+        return "\n".join(lines)

ai_code_assistant/chat/session.py

@@ -0,0 +1,165 @@
+"""Interactive chat session for code discussions."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from pathlib import Path
+from typing import Dict, List, Literal, Optional
+
+from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+
+from ai_code_assistant.config import Config
+from ai_code_assistant.llm import LLMManager
+
+
+@dataclass
+class Message:
+    """A single message in the chat history."""
+    role: Literal["user", "assistant", "system"]
+    content: str
+    timestamp: datetime = field(default_factory=datetime.now)
+    code_context: Optional[str] = None
+
+
+CHAT_SYSTEM_PROMPT = """You are an expert programming assistant with deep knowledge of software development.
+
+You can help with:
+- Explaining code and concepts
+- Debugging issues
+- Suggesting improvements
+- Answering programming questions
+- Discussing architecture and design patterns
+
+When discussing code:
+- Be clear and concise
+- Provide code examples when helpful
+- Explain your reasoning
+- Consider edge cases and best practices
+
+If code context is provided, reference it in your responses when relevant."""
+
+
+class ChatSession:
+    """Manages an interactive chat session."""
+
+    def __init__(
+        self,
+        config: Config,
+        llm_manager: LLMManager,
+        system_prompt: Optional[str] = None,
+    ):
+        self.config = config
+        self.llm = llm_manager
+        self.system_prompt = system_prompt or CHAT_SYSTEM_PROMPT
+        self.history: List[Message] = []
+        self._code_context: Dict[str, str] = {}  # filename -> code
+
+    def add_code_context(self, filename: str, code: str) -> None:
+        """Add code file to the conversation context."""
+        self._code_context[filename] = code
+
+    def load_file_context(self, file_path: Path) -> bool:
+        """Load a file into the conversation context."""
+        try:
+            code = file_path.read_text()
+            self.add_code_context(str(file_path), code)
+            return True
+        except Exception:
+            return False
+
+    def clear_context(self) -> None:
+        """Clear all code context."""
+        self._code_context.clear()
+
+    def send_message(self, user_input: str, stream: bool = False):
+        """Send a message and get a response."""
+        # Add user message to history
+        self.history.append(Message(role="user", content=user_input))
+
+        # Build messages for LLM
+        messages = self._build_messages(user_input)
+
+        if stream:
+            return self._stream_response(messages)
+        else:
+            return self._get_response(messages)
+
+    def _build_messages(self, current_input: str) -> List:
+        """Build message list for LLM including context and history."""
+        messages = []
+        
+        # System message with code context
+        system_content = self.system_prompt
+        if self._code_context:
+            context_str = "\n\n".join(
+                f"**File: {name}**\n```\n{code}\n```"
+                for name, code in self._code_context.items()
+            )
+            system_content += f"\n\n**Code Context:**\n{context_str}"
+        
+        messages.append(SystemMessage(content=system_content))
+        
+        # Add conversation history (keep last N messages to manage context window)
+        max_history = 20
+        recent_history = self.history[-(max_history + 1):-1]  # Exclude current message
+        
+        for msg in recent_history:
+            if msg.role == "user":
+                messages.append(HumanMessage(content=msg.content))
+            elif msg.role == "assistant":
+                messages.append(AIMessage(content=msg.content))
+        
+        # Add current message
+        messages.append(HumanMessage(content=current_input))
+        
+        return messages
+
+    def _get_response(self, messages: List) -> str:
+        """Get complete response from LLM."""
+        response = self.llm.llm.invoke(messages)
+        content = str(response.content)
+        
+        # Add assistant response to history
+        self.history.append(Message(role="assistant", content=content))
+        
+        return content
+
+    def _stream_response(self, messages: List):
+        """Stream response from LLM."""
+        full_response = []
+        
+        for chunk in self.llm.llm.stream(messages):
+            chunk_content = str(chunk.content)
+            full_response.append(chunk_content)
+            yield chunk_content
+        
+        # Add complete response to history
+        self.history.append(Message(role="assistant", content="".join(full_response)))
+
+    def get_history(self) -> List[Message]:
+        """Get conversation history."""
+        return self.history.copy()
+
+    def clear_history(self) -> None:
+        """Clear conversation history."""
+        self.history.clear()
+
+    def export_history(self) -> str:
+        """Export conversation history as markdown."""
+        lines = ["# Chat Session\n"]
+        
+        if self._code_context:
+            lines.append("## Code Context\n")
+            for name in self._code_context:
+                lines.append(f"- {name}")
+            lines.append("")
+        
+        lines.append("## Conversation\n")
+        for msg in self.history:
+            timestamp = msg.timestamp.strftime("%H:%M:%S")
+            role = "**You**" if msg.role == "user" else "**Assistant**"
+            lines.append(f"### {role} ({timestamp})\n")
+            lines.append(msg.content)
+            lines.append("")
+        
+        return "\n".join(lines)
+