Spaces:
Running on Zero
Running on Zero
| """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>" | |
| 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": []}) | |
| def handle_undo(history: List[Dict[str, Any]]) -> List[Dict[str, Any]]: | |
| """Remove last interaction turn.""" | |
| return history[:-2] if len(history) >= 2 else [] | |
| def handle_clear() -> List[Any]: | |
| """Reset chat conversation.""" | |
| return [] | |