chat_ui.py

#!/usr/bin/env python3
"""
SFGPT Chat Interface

A Gradio-based chat UI for interacting with vLLM or the full RAG system.
All conversations are logged to JSON files for easy debugging.

Usage:
    python chat_ui.py                    # Default: vLLM on localhost:8000
    python chat_ui.py --mode rag         # Use full RAG system
    python chat_ui.py --port 7860        # Custom Gradio port
    python chat_ui.py --share            # Create public URL
"""

import argparse
import json
import os
import uuid
from datetime import datetime
from pathlib import Path
from typing import Generator

import gradio as gr
from openai import OpenAI

# Configuration
VLLM_BASE_URL = os.getenv("VLLM_BASE_URL", "http://localhost:8000/v1")
RAG_BASE_URL = os.getenv("RAG_BASE_URL", "http://localhost:8080")
LOG_DIR = Path(__file__).parent / "chat_logs"
LOG_DIR.mkdir(exist_ok=True)

# Default system prompts
SYSTEM_PROMPTS = {
    "default": "You are a helpful AI assistant powered by NVIDIA Nemotron. Be concise and accurate.",
    "sf_data": """You are SFGPT, an AI assistant specializing in San Francisco government data.
You help users explore and understand SF Open Data datasets covering topics like:
- Public safety and crime statistics
- Transportation and traffic
- Housing and building permits
- City infrastructure and services
- Health and social services

Be helpful, accurate, and cite specific datasets when relevant.""",
    "technical": "You are a technical assistant. Provide precise, detailed answers with code examples when appropriate.",
}


class ConversationLogger:
    """Logs all conversations to JSON files for debugging."""

    def __init__(self, log_dir: Path):
        self.log_dir = log_dir
        self.log_dir.mkdir(exist_ok=True)
        self.session_id = self.create_session()

    def create_session(self) -> str:
        """Create a new session ID."""
        return f"{datetime.now().strftime('%Y%m%d_%H%M%S')}_{uuid.uuid4().hex[:8]}"

    def new_session(self) -> str:
        """Start a new session."""
        self.session_id = self.create_session()
        return self.session_id

    def log_message(self, role: str, content: str, metadata: dict = None):
        """Append a message to the session log."""
        log_file = self.log_dir / f"{self.session_id}.json"

        # Load existing log or create new
        if log_file.exists():
            with open(log_file) as f:
                log_data = json.load(f)
        else:
            log_data = {
                "session_id": self.session_id,
                "created_at": datetime.now().isoformat(),
                "messages": []
            }

        # Add message
        message = {
            "timestamp": datetime.now().isoformat(),
            "role": role,
            "content": content,
        }
        if metadata:
            message["metadata"] = metadata

        log_data["messages"].append(message)
        log_data["updated_at"] = datetime.now().isoformat()

        # Save
        with open(log_file, "w") as f:
            json.dump(log_data, f, indent=2)

    def get_session_log(self) -> dict:
        """Retrieve current session log."""
        log_file = self.log_dir / f"{self.session_id}.json"
        if log_file.exists():
            with open(log_file) as f:
                return json.load(f)
        return None

    def export_logs(self) -> str:
        """Export current session logs as JSON string."""
        log_data = self.get_session_log()
        if log_data:
            return json.dumps(log_data, indent=2)
        return "No logs available for current session."


class ChatBackend:
    """Backend for chat completions via vLLM or RAG."""

    def __init__(self, mode: str = "vllm", vllm_url: str = VLLM_BASE_URL):
        self.mode = mode
        self.vllm_url = vllm_url
        self.logger = ConversationLogger(LOG_DIR)

        # Initialize OpenAI client for vLLM
        self.client = OpenAI(
            base_url=vllm_url,
            api_key="not-needed"
        )

        # Get available model
        try:
            models = self.client.models.list()
            self.model_id = models.data[0].id if models.data else "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"
        except Exception as e:
            print(f"Warning: Could not fetch models: {e}")
            self.model_id = "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16"

        print(f"Initialized ChatBackend in {mode} mode")
        print(f"Model: {self.model_id}")
        print(f"Session ID: {self.logger.session_id}")
        print(f"Logs: {LOG_DIR}")

    def chat_stream(
        self,
        message: str,
        history: list,
        system_prompt: str,
        temperature: float,
        max_tokens: int
    ) -> Generator[str, None, None]:
        """Stream chat response from vLLM."""
        # Build messages for API
        messages = [{"role": "system", "content": system_prompt}]

        # Convert Gradio history format to OpenAI format
        for msg in history:
            if isinstance(msg, dict):
                messages.append({"role": msg.get("role", "user"), "content": msg.get("content", "")})
            elif isinstance(msg, (list, tuple)) and len(msg) >= 2:
                # Legacy tuple format
                messages.append({"role": "user", "content": str(msg[0])})
                if msg[1]:
                    messages.append({"role": "assistant", "content": str(msg[1])})

        messages.append({"role": "user", "content": message})

        # Log user message
        self.logger.log_message("user", message, {"temperature": temperature, "max_tokens": max_tokens})

        if self.mode == "rag":
            yield from self._chat_rag(message, temperature, max_tokens)
        else:
            yield from self._chat_vllm(messages, temperature, max_tokens)

    def _chat_vllm(self, messages: list, temperature: float, max_tokens: int) -> Generator[str, None, None]:
        """Stream from vLLM."""
        try:
            stream = self.client.chat.completions.create(
                model=self.model_id,
                messages=messages,
                temperature=temperature,
                max_tokens=max_tokens,
                stream=True
            )

            full_response = ""
            for chunk in stream:
                if chunk.choices[0].delta.content:
                    content = chunk.choices[0].delta.content
                    full_response += content
                    yield full_response

            # Log assistant response
            self.logger.log_message("assistant", full_response, {"model": self.model_id, "mode": self.mode})

        except Exception as e:
            error_msg = f"Error: {str(e)}"
            self.logger.log_message("error", error_msg)
            yield error_msg

    def _chat_rag(self, message: str, temperature: float, max_tokens: int) -> Generator[str, None, None]:
        """Query RAG system."""
        try:
            import requests
            response = requests.post(
                f"{RAG_BASE_URL}/api/query",
                json={"question": message, "temperature": temperature, "max_tokens": max_tokens},
                timeout=60
            )

            if response.ok:
                result = response.json()
                answer = result.get("answer", "No answer returned")
                sources = result.get("sources", [])

                full_response = answer
                if sources:
                    full_response += "\n\n**Sources:**\n"
                    for src in sources[:3]:
                        full_response += f"- {src.get('name', 'Unknown')}\n"

                self.logger.log_message("assistant", full_response, {"sources": sources, "mode": "rag"})
                yield full_response
            else:
                yield f"RAG API error: {response.status_code}"

        except Exception as e:
            # Fall back to vLLM
            yield f"RAG unavailable, using vLLM: {str(e)[:50]}..."
            messages = [{"role": "user", "content": message}]
            yield from self._chat_vllm(messages, temperature, max_tokens)


def create_app(backend: ChatBackend) -> gr.Blocks:
    """Create the Gradio application."""

    current_system_prompt = [SYSTEM_PROMPTS["default"]]
    current_temperature = [0.7]
    current_max_tokens = [1024]

    def respond(message: str, history: list) -> Generator[str, None, None]:
        """Handle chat response with streaming."""
        if not message.strip():
            yield ""
            return

        for response in backend.chat_stream(
            message,
            history,
            current_system_prompt[0],
            current_temperature[0],
            current_max_tokens[0]
        ):
            yield response

    def update_system_prompt(preset: str, custom: str):
        if preset == "custom":
            current_system_prompt[0] = custom
        else:
            current_system_prompt[0] = SYSTEM_PROMPTS.get(preset, SYSTEM_PROMPTS["default"])
        return current_system_prompt[0]

    def update_temperature(val: float):
        current_temperature[0] = val

    def update_max_tokens(val: int):
        current_max_tokens[0] = val

    def new_session():
        session_id = backend.logger.new_session()
        return None, f"New session: {session_id}"

    def export_logs():
        return backend.logger.export_logs()

    # Build UI
    with gr.Blocks(title="SFGPT Chat") as app:
        gr.Markdown(f"""
# SFGPT Chat Interface
**Mode:** {backend.mode.upper()} | **Model:** `{backend.model_id}`

Chat logs saved to `{LOG_DIR}/` for debugging.
        """)

        with gr.Row():
            with gr.Column(scale=3):
                chat = gr.ChatInterface(
                    fn=respond,
                    title="",
                    examples=[
                        "What can you help me with?",
                        "Tell me about San Francisco's open data",
                        "How many datasets are available?",
                    ],
                )

            with gr.Column(scale=1):
                gr.Markdown("### Settings")

                preset = gr.Dropdown(
                    choices=list(SYSTEM_PROMPTS.keys()) + ["custom"],
                    value="default",
                    label="System Prompt Preset",
                    elem_id="sfgpt-preset"
                )

                system_prompt_box = gr.Textbox(
                    value=SYSTEM_PROMPTS["default"],
                    label="System Prompt",
                    lines=4,
                    elem_id="sfgpt-system-prompt"
                )

                temp_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    label="Temperature",
                    elem_id="sfgpt-temperature"
                )

                tokens_slider = gr.Slider(
                    minimum=64,
                    maximum=4096,
                    value=1024,
                    step=64,
                    label="Max Tokens",
                    elem_id="sfgpt-max-tokens"
                )

                gr.Markdown("### Session")

                session_info = gr.Textbox(
                    value=f"Session: {backend.logger.session_id}",
                    label="Current Session",
                    interactive=False,
                    elem_id="sfgpt-session"
                )

                with gr.Row():
                    new_btn = gr.Button("New Session", size="sm", elem_id="sfgpt-new-btn")
                    export_btn = gr.Button("Export Logs", size="sm", elem_id="sfgpt-export-btn")

                log_output = gr.Code(
                    label="Exported Logs",
                    language="json",
                    elem_id="sfgpt-logs"
                )

        # Event bindings
        preset.change(
            update_system_prompt,
            inputs=[preset, system_prompt_box],
            outputs=[system_prompt_box]
        )

        system_prompt_box.change(
            lambda x: update_system_prompt("custom", x),
            inputs=[system_prompt_box],
            outputs=[]
        )

        temp_slider.change(update_temperature, inputs=[temp_slider])
        tokens_slider.change(update_max_tokens, inputs=[tokens_slider])

        new_btn.click(new_session, outputs=[chat.chatbot, session_info])
        export_btn.click(export_logs, outputs=[log_output])

    return app


def main():
    parser = argparse.ArgumentParser(description="SFGPT Chat Interface")
    parser.add_argument(
        "--mode",
        choices=["vllm", "rag"],
        default="vllm",
        help="Backend mode: vllm (direct) or rag (full system)"
    )
    parser.add_argument(
        "--vllm-url",
        default=VLLM_BASE_URL,
        help="vLLM API base URL"
    )
    parser.add_argument(
        "--port",
        type=int,
        default=7860,
        help="Gradio server port"
    )
    parser.add_argument(
        "--share",
        action="store_true",
        help="Create a public URL"
    )
    parser.add_argument(
        "--host",
        default="0.0.0.0",
        help="Host to bind to"
    )

    args = parser.parse_args()

    # Create backend and app
    backend = ChatBackend(mode=args.mode, vllm_url=args.vllm_url)
    app = create_app(backend)

    print(f"\nStarting SFGPT Chat on http://{args.host}:{args.port}")
    print(f"Chat logs directory: {LOG_DIR}")

    app.launch(
        server_name=args.host,
        server_port=args.port,
        share=args.share
    )


if __name__ == "__main__":
    main()