src/ui/controller.py

"""Controller layer coordinating UI events with domain and application services."""

import logging
from typing import Generator, List, Dict, Any, Optional, Tuple
import gradio as gr

from src.config import GenerationConfig, HardwareSpecs
from src.registry.base import BaseModelRegistry
from src.ingestion.service import KnowledgeBaseService
from src.services.inference_service import InferenceService

logger = logging.getLogger(__name__)


class UIController:
    """Manages event callbacks for model loading, inference streaming, and document ingestion."""

    def __init__(
        self,
        registry: BaseModelRegistry,
        inference: InferenceService,
        knowledge_base: KnowledgeBaseService,
        hw_specs: HardwareSpecs,
        spaces: Optional[Any] = None,
    ) -> None:
        self.registry = registry
        self.inference = inference
        self.knowledge_base = knowledge_base
        self.hw_specs = hw_specs
        self.spaces = spaces

        # Wrap handle_inference with ZeroGPU decorator if spaces accelerator is provided
        if self.spaces is not None:
            gpu_decorator = self.spaces.GPU(duration=self.hw_specs.extended_execution_window_sec)
            self.stream_response = gpu_decorator(self.handle_inference)
        else:
            self.stream_response = self.handle_inference

    def handle_load_model(
        self,
        model_id: str,
        quantization: str
    ) -> Generator[Tuple[Any, Any, Any], None, None]:
        """Manage two-stage loading with real-time feedback and button lock."""
        logger.info("Controller: Loading model '%s' (%s)", model_id, quantization)

        loading_html = (
            f"<div style='padding: 8px 14px; border-radius: 6px; background: rgba(99, 102, 241, 0.12); "
            f"border: 1px solid #6366f1; color: #a5b4fc; font-weight: 500; font-size: 0.88rem;'>"
            f"⏳ Loading <strong>{model_id}</strong> ({quantization})... Please wait.</div>"
        )
        yield (
            gr.update(value=loading_html),
            gr.update(interactive=False, value="⏳ Loading..."),
            gr.update(visible=False),
        )

        try:
            success, msg = self.registry.load_model_and_processor(
                model_id=model_id,
                quantization=quantization,
            )
        except Exception as exc:
            logger.exception("Error during model load:")
            success = False
            msg = f"Unexpected failure: {str(exc)}"

        if success:
            ready_html = (
                f"<div style='padding: 8px 14px; border-radius: 6px; background: rgba(16, 185, 129, 0.12); "
                f"border: 1px solid #10b981; color: #10b981; font-weight: 500; font-size: 0.88rem;'>"
                f"✅ {msg}</div>"
            )
            yield (
                gr.update(value=ready_html),
                gr.update(interactive=True, value="🔄 Reload Model"),
                gr.update(visible=True),
            )
        else:
            preview_html = (
                f"<div style='padding: 8px 14px; border-radius: 6px; background: rgba(245, 158, 11, 0.12); "
                f"border: 1px solid #f59e0b; color: #f59e0b; font-weight: 500; font-size: 0.88rem;'>"
                f"ℹ️ {msg} (Preview mode active)</div>"
            )
            yield (
                gr.update(value=preview_html),
                gr.update(interactive=True, value="🔄 Reload Model"),
                gr.update(visible=True),
            )

    def handle_inference(
        self,
        history: List[Dict[str, Any]],
        downsample_mode: str,
        max_slice_nums: int,
        temperature: float,
        top_p: float,
        top_k: int,
        repetition_penalty: float,
        max_new_tokens: int,
        enable_thinking: bool,
        enable_rag: bool,
        rag_max_context_chars: int,
        rag_top_k: int,
        system_prompt: str,
    ) -> Generator[List[Dict[str, Any]], None, None]:
        """Execute stream generation with optional context grounding."""
        if not history:
            yield history
            return

        if history[-1]["role"] != "assistant":
            history.append({"role": "assistant", "content": ""})

        gen_config = GenerationConfig(
            downsample_mode=downsample_mode,
            max_slice_nums=int(max_slice_nums),
            temperature=float(temperature),
            top_p=float(top_p),
            top_k=int(top_k),
            repetition_penalty=float(repetition_penalty),
            max_new_tokens=int(max_new_tokens),
            enable_thinking=bool(enable_thinking),
            enable_rag=bool(enable_rag),
            rag_max_context_chars=int(rag_max_context_chars),
            rag_top_k=int(rag_top_k),
            system_prompt=system_prompt,
        )

        rag_context = ""
        if enable_rag and self.knowledge_base.chunks:
            latest_query = self.inference.extract_latest_user_query(history[:-1])
            if latest_query:
                rag_context = self.knowledge_base.retrieve_context(
                    query=latest_query,
                    top_k=int(rag_top_k),
                    max_context_chars=int(rag_max_context_chars),
                )

        formatted_messages = self.inference.parse_gradio_history(
            raw_history=history[:-1],
            system_prompt=system_prompt,
            rag_context=rag_context,
        )

        try:
            for chunk in self.inference.stream_inference(formatted_messages, gen_config):
                history[-1]["content"] = chunk
                yield history
        except Exception as exc:
            logger.exception("Inference error occurred:")
            history[-1]["content"] = f"⚠️ Generation error: {str(exc)}"
            yield history

    def handle_document_upload(self, files: Optional[List[str]]) -> Tuple[Any, Any, str]:
        """Process and index uploaded documents."""
        if not files:
            names = self.knowledge_base.get_document_names()
            return (
                self.knowledge_base.get_document_table_data(),
                gr.update(choices=names, value=names[0] if names else None),
                "<span style='color: #94a3b8; font-size: 0.85rem;'>No files selected.</span>",
            )

        res = self.knowledge_base.add_documents(files)
        table_data = self.knowledge_base.get_document_table_data()
        names = self.knowledge_base.get_document_names()
        msg = (
            f"<span style='color: #10b981; font-weight: 500; font-size: 0.85rem;'>"
            f"✅ Indexed {res['new_docs']} new file(s). "
            f"Total: {res['count']} document(s) & {res['chunks']} chunks active.</span>"
        )
        return table_data, gr.update(choices=names, value=names[0] if names else None), msg

    def handle_document_delete(self, doc_name: Optional[str]) -> Tuple[Any, Any, str]:
        """Remove a selected document from the knowledge base."""
        if not doc_name:
            names = self.knowledge_base.get_document_names()
            return (
                self.knowledge_base.get_document_table_data(),
                gr.update(choices=names, value=names[0] if names else None),
                "<span style='color: #ef4444; font-size: 0.85rem;'>Select a document to delete.</span>",
            )

        removed = self.knowledge_base.remove_document(doc_name)
        table_data = self.knowledge_base.get_document_table_data()
        names = self.knowledge_base.get_document_names()
        if removed:
            msg = f"<span style='color: #f59e0b; font-size: 0.85rem;'>🗑️ Removed '{doc_name}'. Remaining: {len(names)} document(s).</span>"
        else:
            msg = f"<span style='color: #ef4444; font-size: 0.85rem;'>Document '{doc_name}' not found.</span>"

        return table_data, gr.update(choices=names, value=names[0] if names else None), msg

    def handle_document_clear(self) -> Tuple[List[Any], Any, str]:
        """Clear all indexed documents."""
        self.knowledge_base.clear()
        return [], gr.update(choices=[], value=None), "<span style='color: #94a3b8; font-size: 0.85rem;'>Knowledge base cleared.</span>"

    @staticmethod
    def handle_user_submit(message: Dict[str, Any], history: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], Any]:
        """Parse user multimodal submission into history."""
        text = (message.get("text") or "").strip()
        files = message.get("files") or []

        if not text and not files:
            return history, gr.MultimodalTextbox(value={"text": "", "files": []})

        content_elements = [{"path": f} for f in files]
        if text:
            content_elements.append(text)

        user_payload = content_elements if len(content_elements) > 1 else content_elements[0]
        history.append({"role": "user", "content": user_payload})

        return history, gr.MultimodalTextbox(value={"text": "", "files": []})

    @staticmethod
    def handle_undo(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
        """Remove last interaction turn."""
        return history[:-2] if len(history) >= 2 else []

    @staticmethod
    def handle_clear() -> List[Any]:
        """Reset chat conversation."""
        return []