terminal_chatbot.py

#!/usr/bin/env python3
"""
Terminal-based AI Coding Assistant
A command-line interface for the 5B parameter coding model via API
"""

import requests
import json
import time
import sys
import os
from typing import List, Dict, Any
from rich.console import Console
from rich.panel import Panel
from rich.syntax import Syntax
from rich.prompt import Prompt, Confirm
from rich.markdown import Markdown
from rich.table import Table
from rich import print as rprint

console = Console()

class TerminalChatbot:
    """Terminal-based chatbot client for the AI coding assistant."""
    
    def __init__(self, server_url: str = "http://localhost:8000"):
        self.server_url = server_url.rstrip('/')
        self.api_url = f"{server_url}/api/chat"
        self.history: List[Dict[str, str]] = []
        self.current_language = "python"
        self.temperature = 0.7
        
    def check_server_connection(self) -> bool:
        """Check if the model server is running."""
        try:
            response = requests.get(f"{self.server_url}/health", timeout=5)
            return response.status_code == 200
        except requests.exceptions.RequestException:
            return False
    
    def send_message(self, message: str) -> Dict[str, Any]:
        """Send a message to the model server and get response."""
        try:
            payload = {
                "message": message,
                "history": self.history,
                "language": self.current_language,
                "temperature": self.temperature
            }
            
            response = requests.post(
                self.api_url,
                json=payload,
                headers={"Content-Type": "application/json"},
                timeout=60
            )
            
            if response.status_code == 200:
                return response.json()
            else:
                return {
                    "choices": [{"message": {"content": f"Server error: {response.status_code}"}}],
                    "history": self.history
                }
                
        except requests.exceptions.RequestException as e:
            return {
                "choices": [{"message": {"content": f"Connection error: {str(e)}"}}],
                "history": self.history
            }
    
    def format_response(self, response: str) -> None:
        """Format and display the model's response."""
        if not response:
            return
        
        # Split response into code blocks and text
        parts = response.split('```')
        
        for i, part in enumerate(parts):
            if i % 2 == 0:  # Text parts
                if part.strip():
                    try:
                        markdown = Markdown(part.strip())
                        console.print(markdown)
                    except:
                        console.print(part.strip())
            else:  # Code parts
                lines = part.split('\n', 1)
                if len(lines) >= 2:
                    language = lines[0].strip() if lines[0].strip() else 'text'
                    code = lines[1]
                else:
                    language = 'text'
                    code = part
                
                if code.strip():
                    syntax = Syntax(code.strip(), language, theme="monokai", line_numbers=True)
                    console.print(syntax)
    
    def show_welcome(self) -> None:
        """Display welcome message and help."""
        welcome_text = """
# 🤖 AI Coding Assistant - Terminal Version

Welcome to your AI-powered coding companion! I can help you with:

• **Code Generation** - Write functions, classes, and complete programs
• **Debugging** - Find and fix errors in your code
• **Algorithm Implementation** - From simple to complex algorithms
• **Best Practices** - Clean, efficient, and readable code
• **Concept Explanation** - Understand programming concepts

## Quick Start Commands:
• `/help` - Show this help
• `/lang <language>` - Change programming language
• `/temp <value>` - Set creativity (0.1-1.0)
• `/clear` - Clear chat history
• `/quit` or `/exit` - Exit the program

## Example Prompts:
• "Write a Python function to reverse a linked list"
• "Create a React component for user authentication"
• "Explain Big O notation with code examples"
• "Debug this JavaScript code: [paste your code]"

**Ready to code? Just ask me anything!**
        """
        
        panel = Panel(
            Markdown(welcome_text),
            title="AI Coder Terminal",
            border_style="blue",
            padding=(1, 2)
        )
        console.print(panel)
    
    def show_settings(self) -> None:
        """Display current settings."""
        table = Table(title="Current Settings")
        table.add_column("Setting", style="cyan")
        table.add_column("Value", style="green")
        table.add_column("Description", style="yellow")
        
        table.add_row("Language", self.current_language, "Target programming language")
        table.add_row("Temperature", str(self.temperature), "Creativity level (0.1-1.0)")
        table.add_row("Server", self.server_url, "Model server URL")
        table.add_row("History", str(len(self.history)), "Messages in conversation")
        
        console.print(table)
    
    def handle_command(self, command: str) -> bool:
        """Handle special commands. Returns True if command was processed."""
        cmd = command.lower().strip()
        
        if cmd in ['/help', '/h']:
            self.show_help()
            return True
        elif cmd.startswith('/lang '):
            language = command.split(' ', 1)[1].strip()
            self.current_language = language
            console.print(f"[green]✓[/green] Language set to: {language}")
            return True
        elif cmd.startswith('/temp '):
            try:
                temp = float(command.split(' ', 1)[1].strip())
                if 0.1 <= temp <= 1.0:
                    self.temperature = temp
                    console.print(f"[green]✓[/green] Temperature set to: {temp}")
                else:
                    console.print("[red]Temperature must be between 0.1 and 1.0[/red]")
            except ValueError:
                console.print("[red]Invalid temperature value[/red]")
            return True
        elif cmd in ['/settings', '/config']:
            self.show_settings()
            return True
        elif cmd in ['/clear', '/reset']:
            self.history = []
            console.print("[green]✓[/green] Chat history cleared")
            return True
        elif cmd in ['/quit', '/exit', '/q']:
            console.print("[yellow]Goodbye! 👋[/yellow]")
            sys.exit(0)
        else:
            console.print(f"[red]Unknown command: {command}[/red]")
            return True
    
    def show_help(self) -> None:
        """Display detailed help information."""
        help_text = """
# Available Commands

## Chat Commands
- **Regular text**: Ask questions or request code
- **/help** or **/h**: Show this help message
- **/settings**: Display current settings
- **/clear**: Clear chat history

## Configuration Commands
- **/lang <language>**: Change programming language
  - Example: `/lang javascript`
- **/temp <value>**: Set creativity level (0.1-1.0)
  - Example: `/temp 0.3` (more precise)
  - Example: `/temp 0.9` (more creative)

## Exit Commands
- **/quit** or **/exit**: Exit the program

## Programming Languages Supported
python, javascript, java, cpp, c, go, rust, typescript, 
php, ruby, swift, kotlin, sql, html, css, bash, powershell
        """
        
        console.print(Panel(Markdown(help_text), title="Help", border_style="green"))
    
    def run(self) -> None:
        """Main chatbot loop."""
        self.show_welcome()
        
        # Check server connection
        console.print("\n[yellow]Checking server connection...[/yellow]")
        if not self.check_server_connection():
            console.print(f"[red]❌ Cannot connect to server at {self.server_url}[/red]")
            console.print("[yellow]💡 Make sure the model server is running with:[/yellow]")
            console.print("[cyan]python model_server.py[/cyan]")
            return
        
        console.print("[green]✓[/green] Connected to model server!")
        
        # Main interaction loop
        while True:
            try:
                # Get user input
                user_input = Prompt.ask(
                    f"[bold blue]You[/bold blue] ({self.current_language})"
                ).strip()
                
                if not user_input:
                    continue
                
                # Handle commands
                if user_input.startswith('/'):
                    self.handle_command(user_input)
                    continue
                
                # Show typing indicator
                with console.status("[bold green]AI is thinking...[/bold green]"):
                    start_time = time.time()
                    
                    # Send to server and get response
                    response_data = self.send_message(user_input)
                    
                    end_time = time.time()
                    response_time = end_time - start_time
                
                # Display response
                if response_data and "choices" in response_data:
                    response = response_data["choices"][0]["message"]["content"]
                    
                    # Display response with timing
                    console.print(f"\n[dim]Response time: {response_time:.2f}s[/dim]")
                    console.print(f"[bold green]AI:[/bold green]")
                    
                    # Format and display response
                    self.format_response(response)
                    
                    # Update history
                    self.history = response_data.get("history", self.history)
                    
                    console.print()  # Add spacing
                else:
                    console.print("[red]❌ Invalid response from server[/red]")
                    
            except KeyboardInterrupt:
                console.print("\n[yellow]Interrupted by user[/yellow]")
                if Confirm.ask("Exit the program?"):
                    break
            except Exception as e:
                console.print(f"[red]❌ Error: {str(e)}[/red]")

def main():
    """Main entry point."""
    # Check for custom server URL
    server_url = "http://localhost:8000"
    if len(sys.argv) > 1:
        server_url = sys.argv[1]
    
    console.print(f"[cyan]AI Coding Assistant Terminal[/cyan]")
    console.print(f"[dim]Server: {server_url}[/dim]\n")
    
    # Create and run chatbot
    chatbot = TerminalChatbot(server_url)
    chatbot.run()

if __name__ == "__main__":
    main()