chore: import upstream snapshot with attribution

kt-kernel/python/cli/commands/chat.py

@@ -0,0 +1,572 @@
+"""
+Chat command for kt-cli.
+
+Provides interactive chat interface with running model server.
+"""
+
+import json
+import os
+import sys
+from datetime import datetime
+from pathlib import Path
+from typing import Optional
+
+import typer
+from rich.console import Console
+from rich.markdown import Markdown
+from rich.panel import Panel
+from rich.prompt import Prompt, Confirm
+
+from kt_kernel.cli.config.settings import get_settings
+from kt_kernel.cli.i18n import t
+from kt_kernel.cli.utils.console import (
+    console,
+    print_error,
+    print_info,
+    print_success,
+    print_warning,
+)
+
+# Try to import OpenAI SDK
+try:
+    from openai import OpenAI
+
+    HAS_OPENAI = True
+except ImportError:
+    HAS_OPENAI = False
+
+
+def chat(
+    host: Optional[str] = typer.Option(
+        None,
+        "--host",
+        "-H",
+        help="Server host address",
+    ),
+    port: Optional[int] = typer.Option(
+        None,
+        "--port",
+        "-p",
+        help="Server port",
+    ),
+    model: Optional[str] = typer.Option(
+        None,
+        "--model",
+        "-m",
+        help="Model name (if server hosts multiple models)",
+    ),
+    temperature: float = typer.Option(
+        0.7,
+        "--temperature",
+        "-t",
+        help="Sampling temperature (0.0 to 2.0)",
+    ),
+    max_tokens: int = typer.Option(
+        2048,
+        "--max-tokens",
+        help="Maximum tokens to generate",
+    ),
+    system_prompt: Optional[str] = typer.Option(
+        None,
+        "--system",
+        "-s",
+        help="System prompt",
+    ),
+    save_history: bool = typer.Option(
+        True,
+        "--save-history/--no-save-history",
+        help="Save conversation history",
+    ),
+    history_file: Optional[Path] = typer.Option(
+        None,
+        "--history-file",
+        help="Path to save conversation history",
+    ),
+    stream: bool = typer.Option(
+        True,
+        "--stream/--no-stream",
+        help="Enable streaming output",
+    ),
+) -> None:
+    """Start interactive chat with a running model server.
+
+    Examples:
+        kt chat                          # Connect to default server
+        kt chat --host 127.0.0.1 -p 8080 # Connect to specific server
+        kt chat -t 0.9 --max-tokens 4096 # Adjust generation parameters
+    """
+    if not HAS_OPENAI:
+        print_error(t("chat_openai_required"))
+        console.print()
+        console.print(t("chat_install_hint"))
+        console.print("  pip install openai")
+        raise typer.Exit(1)
+
+    settings = get_settings()
+
+    # Resolve server connection
+    final_host = host or settings.get("server.host", "127.0.0.1")
+    final_port = port or settings.get("server.port", 30000)
+
+    # Construct base URL for OpenAI-compatible API
+    base_url = f"http://{final_host}:{final_port}/v1"
+
+    console.print()
+    console.print(
+        Panel.fit(
+            f"[bold cyan]{t('chat_title')}[/bold cyan]\n\n"
+            f"{t('chat_server')}: [yellow]{final_host}:{final_port}[/yellow]\n"
+            f"{t('chat_temperature')}: [cyan]{temperature}[/cyan] | {t('chat_max_tokens')}: [cyan]{max_tokens}[/cyan]\n\n"
+            f"[dim]{t('chat_help_hint')}[/dim]",
+            border_style="cyan",
+        )
+    )
+    console.print()
+
+    # Check for proxy environment variables
+    proxy_vars = ["HTTP_PROXY", "HTTPS_PROXY", "http_proxy", "https_proxy", "ALL_PROXY", "all_proxy"]
+    detected_proxies = {var: os.environ.get(var) for var in proxy_vars if os.environ.get(var)}
+
+    if detected_proxies:
+        proxy_info = ", ".join(f"{k}={v}" for k, v in detected_proxies.items())
+        console.print()
+        print_warning(t("chat_proxy_detected"))
+        console.print(f"  [dim]{proxy_info}[/dim]")
+        console.print()
+
+        use_proxy = Confirm.ask(t("chat_proxy_confirm"), default=False)
+
+        if not use_proxy:
+            # Temporarily disable proxy for this connection
+            for var in proxy_vars:
+                if var in os.environ:
+                    del os.environ[var]
+            print_info(t("chat_proxy_disabled"))
+            console.print()
+
+    # Initialize OpenAI client
+    try:
+        client = OpenAI(
+            base_url=base_url,
+            api_key="EMPTY",  # SGLang doesn't require API key
+        )
+
+        # Test connection
+        print_info(t("chat_connecting"))
+        models = client.models.list()
+        available_models = [m.id for m in models.data]
+
+        if not available_models:
+            print_error(t("chat_no_models"))
+            raise typer.Exit(1)
+
+        # Select model
+        if model:
+            if model not in available_models:
+                print_warning(t("chat_model_not_found", model=model, available=", ".join(available_models)))
+                selected_model = available_models[0]
+            else:
+                selected_model = model
+        else:
+            selected_model = available_models[0]
+
+        print_success(t("chat_connected", model=selected_model))
+        console.print()
+
+        # Load tokenizer for accurate token counting
+        tokenizer = None
+        try:
+            from transformers import AutoTokenizer
+
+            # selected_model is the model path
+            tokenizer = AutoTokenizer.from_pretrained(selected_model, trust_remote_code=True)
+            console.print(f"[dim]Loaded tokenizer from {selected_model}[/dim]")
+            console.print()
+        except Exception as e:
+            console.print(f"[dim yellow]Warning: Could not load tokenizer, token counts will be estimated[/dim]")
+            console.print()
+
+    except Exception as e:
+        print_error(t("chat_connect_failed", error=str(e)))
+        console.print()
+        console.print(t("chat_server_not_running"))
+        console.print("  kt run <model>")
+        raise typer.Exit(1)
+
+    # Initialize conversation history
+    messages = []
+
+    # Add system prompt if provided
+    if system_prompt:
+        messages.append({"role": "system", "content": system_prompt})
+
+    # Setup history file
+    if save_history:
+        if history_file is None:
+            history_dir = settings.config_dir / "chat_history"
+            history_dir.mkdir(parents=True, exist_ok=True)
+            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+            history_file = history_dir / f"chat_{timestamp}.json"
+        else:
+            history_file = Path(history_file)
+            history_file.parent.mkdir(parents=True, exist_ok=True)
+
+    # Main chat loop
+    try:
+        while True:
+            # Get user input - use console.input() for better CJK character support
+            try:
+                user_input = console.input(f"[bold green]{t('chat_user_prompt')}[/bold green]: ")
+            except (EOFError, KeyboardInterrupt):
+                console.print()
+                print_info(t("chat_goodbye"))
+                break
+
+            if not user_input.strip():
+                continue
+
+            # Handle special commands
+            if user_input.startswith("/"):
+                if _handle_command(user_input, messages, temperature, max_tokens):
+                    continue
+                else:
+                    break  # Exit command
+
+            # Add user message to history
+            messages.append({"role": "user", "content": user_input})
+
+            # Generate response
+            console.print()
+            console.print(f"[bold cyan]{t('chat_assistant_prompt')}[/bold cyan]")
+
+            try:
+                if stream:
+                    # Streaming response
+                    response_content = _stream_response(
+                        client, selected_model, messages, temperature, max_tokens, tokenizer
+                    )
+                else:
+                    # Non-streaming response
+                    response_content = _generate_response(
+                        client, selected_model, messages, temperature, max_tokens, tokenizer
+                    )
+
+                # Add assistant response to history
+                messages.append({"role": "assistant", "content": response_content})
+
+                console.print()
+
+            except Exception as e:
+                print_error(t("chat_generation_error", error=str(e)))
+                # Remove the user message that caused the error
+                messages.pop()
+                continue
+
+            # Save history if enabled
+            if save_history:
+                _save_history(history_file, messages, selected_model)
+
+    except KeyboardInterrupt:
+        console.print()
+        console.print()
+        print_info(t("chat_interrupted"))
+
+    # Final history save
+    if save_history and messages:
+        _save_history(history_file, messages, selected_model)
+        console.print(f"[dim]{t('chat_history_saved', path=str(history_file))}[/dim]")
+        console.print()
+
+
+def _stream_response(
+    client: "OpenAI",
+    model: str,
+    messages: list,
+    temperature: float,
+    max_tokens: int,
+    tokenizer=None,
+) -> str:
+    """Generate streaming response and display in real-time."""
+    import time
+
+    response_content = ""
+    reasoning_content = ""
+
+    # Performance tracking
+    first_token_time = None
+    chunk_count = 0
+
+    try:
+        # Start timing before sending request
+        start_time = time.time()
+
+        stream = client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stream=True,
+        )
+
+        for chunk in stream:
+            delta = chunk.choices[0].delta if chunk.choices else None
+            if delta:
+                reasoning_delta = getattr(delta, "reasoning_content", None)
+                if reasoning_delta:
+                    if first_token_time is None:
+                        first_token_time = time.time()
+                    reasoning_content += reasoning_delta
+                    console.print(reasoning_delta, end="", style="dim")
+                    chunk_count += 1
+
+                if delta.content:
+                    if first_token_time is None:
+                        first_token_time = time.time()
+                    content = delta.content
+                    response_content += content
+                    console.print(content, end="")
+                    chunk_count += 1
+
+        console.print()  # Newline after streaming
+
+        # Display performance metrics
+        end_time = time.time()
+        if first_token_time and chunk_count > 0:
+            ttft = first_token_time - start_time
+            total_time = end_time - start_time
+
+            # Calculate TPOT based on chunks
+            if chunk_count > 1:
+                generation_time = total_time - ttft
+                tpot = generation_time / (chunk_count - 1)
+            else:
+                tpot = 0
+
+            # Calculate accurate token counts using tokenizer
+            if tokenizer:
+                input_tokens = _count_tokens_with_tokenizer(messages, tokenizer)
+                output_tokens = _count_tokens_with_tokenizer(
+                    [{"role": "assistant", "content": response_content}], tokenizer
+                )
+                token_prefix = ""
+            else:
+                # Fallback to estimation
+                input_tokens = _estimate_tokens(messages)
+                output_tokens = _estimate_tokens([{"role": "assistant", "content": response_content}])
+                token_prefix = "~"
+
+            # Build metrics display
+            metrics = f"[dim]Total: {total_time*1000:.0f}ms | TTFT: {ttft*1000:.0f}ms"
+            if tpot > 0:
+                metrics += f" | TPOT: {tpot*1000:.1f}ms"
+            metrics += f" | In: {token_prefix}{input_tokens} | Out: {token_prefix}{output_tokens}"
+            metrics += "[/dim]"
+
+            console.print(metrics)
+
+    except Exception as e:
+        raise Exception(f"Streaming error: {e}")
+
+    return response_content
+
+
+def _count_tokens_with_tokenizer(messages: list, tokenizer) -> int:
+    """Count tokens accurately using the model's tokenizer."""
+    try:
+        # Concatenate all message content
+        text = ""
+        for msg in messages:
+            role = msg.get("role", "")
+            content = msg.get("content", "")
+            # Simple format: role + content
+            text += f"{role}: {content}\n"
+
+        # Encode and count tokens - suppress any debug output from custom tokenizers
+        import os
+        import sys
+        from contextlib import redirect_stdout, redirect_stderr
+
+        with open(os.devnull, "w") as devnull:
+            with redirect_stdout(devnull), redirect_stderr(devnull):
+                tokens = tokenizer.encode(text, add_special_tokens=True)
+        return len(tokens)
+    except Exception:
+        # Fallback to estimation if tokenizer fails
+        return _estimate_tokens(messages)
+
+
+def _estimate_tokens(messages: list) -> int:
+    """Estimate token count for messages (rough approximation)."""
+    total_chars = 0
+    for msg in messages:
+        content = msg.get("content", "")
+        total_chars += len(content)
+
+    # Rough estimation:
+    # - English: ~4 chars per token
+    # - Chinese: ~1.5 chars per token
+    # Use 2.5 as average
+    return max(1, int(total_chars / 2.5))
+
+
+def _generate_response(
+    client: "OpenAI",
+    model: str,
+    messages: list,
+    temperature: float,
+    max_tokens: int,
+    tokenizer=None,
+) -> str:
+    """Generate non-streaming response."""
+    import time
+
+    try:
+        start_time = time.time()
+
+        response = client.chat.completions.create(
+            model=model,
+            messages=messages,
+            temperature=temperature,
+            max_tokens=max_tokens,
+            stream=False,
+        )
+
+        end_time = time.time()
+        total_time = end_time - start_time
+
+        content = response.choices[0].message.content
+
+        # Display as markdown
+        md = Markdown(content)
+        console.print(md)
+
+        # Calculate accurate token counts using tokenizer
+        if tokenizer:
+            input_tokens = _count_tokens_with_tokenizer(messages, tokenizer)
+            output_tokens = _count_tokens_with_tokenizer([{"role": "assistant", "content": content}], tokenizer)
+            token_prefix = ""
+        else:
+            # Fallback to API usage or estimation
+            input_tokens = response.usage.prompt_tokens if response.usage else _estimate_tokens(messages)
+            output_tokens = (
+                response.usage.completion_tokens
+                if response.usage
+                else _estimate_tokens([{"role": "assistant", "content": content}])
+            )
+            token_prefix = "" if response.usage else "~"
+
+        # Display performance metrics
+        console.print(
+            f"[dim]Time: {total_time*1000:.0f}ms | "
+            f"In: {token_prefix}{input_tokens} | Out: {token_prefix}{output_tokens}[/dim]"
+        )
+
+        return content
+
+    except Exception as e:
+        raise Exception(f"Generation error: {e}")
+
+
+def _handle_command(command: str, messages: list, temperature: float, max_tokens: int) -> bool:
+    """Handle special commands. Returns True to continue chat, False to exit."""
+    cmd = command.lower().strip()
+
+    if cmd in ["/quit", "/exit", "/q"]:
+        console.print()
+        print_info(t("chat_goodbye"))
+        return False
+
+    elif cmd in ["/help", "/h"]:
+        console.print()
+        console.print(
+            Panel(
+                f"[bold]{t('chat_help_title')}[/bold]\n\n{t('chat_help_content')}",
+                title="Help",
+                border_style="cyan",
+            )
+        )
+        console.print()
+        return True
+
+    elif cmd in ["/clear", "/c"]:
+        messages.clear()
+        console.print()
+        print_success(t("chat_history_cleared"))
+        console.print()
+        return True
+
+    elif cmd in ["/history", "/hist"]:
+        console.print()
+        if not messages:
+            print_info(t("chat_no_history"))
+        else:
+            console.print(
+                Panel(
+                    _format_history(messages),
+                    title=t("chat_history_title", count=len(messages)),
+                    border_style="cyan",
+                )
+            )
+        console.print()
+        return True
+
+    elif cmd in ["/info", "/i"]:
+        console.print()
+        console.print(
+            Panel(
+                f"[bold]{t('chat_info_title')}[/bold]\n\n{t('chat_info_content', temperature=temperature, max_tokens=max_tokens, messages=len(messages))}",
+                title="Info",
+                border_style="cyan",
+            )
+        )
+        console.print()
+        return True
+
+    elif cmd in ["/retry", "/r"]:
+        if len(messages) >= 2 and messages[-1]["role"] == "assistant":
+            # Remove last assistant response
+            messages.pop()
+            print_info(t("chat_retrying"))
+            console.print()
+        else:
+            print_warning(t("chat_no_retry"))
+            console.print()
+        return True
+
+    else:
+        print_warning(t("chat_unknown_command", command=command))
+        console.print(f"[dim]{t('chat_unknown_hint')}[/dim]")
+        console.print()
+        return True
+
+
+def _format_history(messages: list) -> str:
+    """Format conversation history for display."""
+    lines = []
+    for i, msg in enumerate(messages, 1):
+        role = msg["role"].capitalize()
+        content = msg["content"]
+
+        # Truncate long messages
+        if len(content) > 200:
+            content = content[:200] + "..."
+
+        lines.append(f"[bold]{i}. {role}:[/bold] {content}")
+
+    return "\n\n".join(lines)
+
+
+def _save_history(file_path: Path, messages: list, model: str) -> None:
+    """Save conversation history to file."""
+    try:
+        history_data = {
+            "model": model,
+            "timestamp": datetime.now().isoformat(),
+            "messages": messages,
+        }
+
+        with open(file_path, "w", encoding="utf-8") as f:
+            json.dump(history_data, f, indent=2, ensure_ascii=False)
+
+    except Exception as e:
+        print_warning(f"Failed to save history: {e}")