Update app.py

app.py

@@ -1,66 +1,184 @@
 import os
+import re
 import glob
-from typing import List, Tuple
-import time
+import threading
+import hashlib
+import logging
+from typing import List, Tuple, Optional
 
 import gradio as gr
 import numpy as np
-from sentence_transformers import SentenceTransformer 
+import yaml
+from sentence_transformers import SentenceTransformer
 
 # -----------------------------
-# CONFIG
-# -----------------------------
-KB_DIR = "./kb"  # folder with .txt or .md files
-EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-TOP_K = 3
-CHUNK_SIZE = 500  # characters
-CHUNK_OVERLAP = 100  # characters
-MIN_SIMILARITY_THRESHOLD = 0.3  # Minimum similarity score to include results
+# LOGGING
+# ----------------------------- 
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
+logger = logging.getLogger(__name__)
 
 # -----------------------------
+# CONFIG — loaded from config.yaml, with safe defaults
+# ----------------------------- 
+_CONFIG_PATH = "./config.yaml"
+
+def load_config(path: str = _CONFIG_PATH) -> dict:
+    """Load configuration from YAML file with fallback defaults."""
+    defaults = {
+        "kb": {"directory": "./kb"},
+        "models": {"embedding": "sentence-transformers/all-MiniLM-L6-v2"},
+        "chunking": {"chunk_size": 500, "overlap": 100},
+        "retrieval": {"top_k": 3, "min_similarity": 0.3, "context_chars": 1500},
+        "llm": {"model": "TinyLlama/TinyLlama-1.1B-Chat-v1.0", "max_new_tokens": 256},
+        "client": {"name": "Self-Service KB Assistant"},
+        "messages": {
+            "welcome": "👋 Ask me anything you'd normally search for in a knowledge base.",
+            "no_answer": "❌ I couldn't find anything relevant in the knowledge base for this query.",
+        },
+    }
+    if os.path.exists(path):
+        try:
+            with open(path, "r", encoding="utf-8") as f:
+                user_cfg = yaml.safe_load(f) or {}
+            # Deep merge user config over defaults
+            for key, val in user_cfg.items():
+                if isinstance(val, dict) and key in defaults:
+                    defaults[key].update(val)
+                else:
+                    defaults[key] = val
+            logger.info(f"Config loaded from {path}")
+        except Exception as e:
+            logger.warning(f"Could not load {path}: {e}. Using defaults.")
+    else:
+        logger.warning(f"{path} not found. Using default configuration.")
+    return defaults
+
+CONFIG = load_config()
+
+KB_DIR               = CONFIG["kb"]["directory"]
+EMBEDDING_MODEL_NAME = CONFIG["models"]["embedding"]
+CHUNK_SIZE           = CONFIG["chunking"]["chunk_size"]
+CHUNK_OVERLAP        = CONFIG["chunking"]["overlap"]
+TOP_K                = CONFIG["retrieval"]["top_k"]
+MIN_SIMILARITY       = CONFIG["retrieval"]["min_similarity"]
+CONTEXT_CHARS        = CONFIG["retrieval"]["context_chars"]
+LLM_MODEL_NAME       = CONFIG["llm"]["model"]
+MAX_NEW_TOKENS       = CONFIG["llm"]["max_new_tokens"]
+MAX_QUERY_CHARS      = 1000  # Guard against excessively long user input
+
+# Paths for caching embeddings
+_CACHE_DIR  = "./.cache"
+_CACHE_FILE = os.path.join(_CACHE_DIR, "kb_embeddings.npz")
+
+# ----------------------------- 
 # UTILITIES
-# -----------------------------
+# ----------------------------- 
 
 def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVERLAP) -> List[str]:
-    """Split long text into overlapping chunks so retrieval is more precise."""
+    """
+    Split text into overlapping chunks that break on sentence boundaries
+    where possible, preventing mid-word or mid-sentence splits.
+    """
     if not text:
         return []
 
+    # Split into sentences using punctuation as delimiters
+    sentence_endings = re.compile(r'(?<=[.!?])\s+')
+    sentences = sentence_endings.split(text.strip())
+
     chunks = []
-    start = 0
-    length = len(text)
+    current_chunk = []
+    current_len = 0
+
+    for sentence in sentences:
+        sentence = sentence.strip()
+        if not sentence:
+            continue
+        slen = len(sentence)
+
+        # If a single sentence exceeds chunk_size, force-split it
+        if slen > chunk_size:
+            if current_chunk:
+                chunks.append(" ".join(current_chunk))
+                current_chunk = []
+                current_len = 0
+            # Hard split the long sentence with overlap
+            start = 0
+            while start < slen:
+                end = min(start + chunk_size, slen)
+                chunks.append(sentence[start:end].strip())
+                start += chunk_size - overlap
+            continue
+
+        if current_len + slen + 1 > chunk_size:
+            if current_chunk:
+                chunks.append(" ".join(current_chunk))
+            # Carry over last sentence(s) for overlap
+            overlap_words = []
+            overlap_len = 0
+            for s in reversed(current_chunk):
+                if overlap_len + len(s) + 1 > overlap:
+                    break
+                overlap_words.insert(0, s)
+                overlap_len += len(s) + 1
+            current_chunk = overlap_words
+            current_len = overlap_len
+
+        current_chunk.append(sentence)
+        current_len += slen + 1
+
+    if current_chunk:
+        chunks.append(" ".join(current_chunk))
+
+    return [c for c in chunks if c.strip()]
+
+
+def clean_context(text: str) -> str:
+    """
+    Clean text for LLM context: remove only Markdown heading markers (leading #),
+    collapse whitespace, but preserve legitimate # characters mid-text.
+    """
+    # Only strip leading '#' at the start of lines (Markdown headings)
+    text = re.sub(r'(?m)^#+\s*', '', text)
+    # Collapse multiple whitespace/newlines
+    text = re.sub(r'\s+', ' ', text)
+    return text.strip()
 
-    while start < length:
-        end = min(start + chunk_size, length)
-        chunk = text[start:end].strip()
-        if chunk:
-            chunks.append(chunk)
-        start += chunk_size - overlap
 
-    return chunks
+def kb_fingerprint(kb_dir: str) -> str:
+    """Return an MD5 fingerprint of all KB filenames + modification times for cache invalidation."""
+    paths = sorted(
+        glob.glob(os.path.join(kb_dir, "*.txt")) +
+        glob.glob(os.path.join(kb_dir, "*.md"))
+    )
+    sig = "|".join(f"{p}:{os.path.getmtime(p)}" for p in paths)
+    return hashlib.md5(sig.encode()).hexdigest()
 
 
 def load_kb_texts(kb_dir: str = KB_DIR) -> List[Tuple[str, str]]:
     """
     Load all .txt and .md files from the KB directory.
     Returns a list of (source_name, content).
+    Falls back to built-in demo content if directory is empty or missing.
     """
     texts = []
 
     if os.path.isdir(kb_dir):
-        paths = glob.glob(os.path.join(kb_dir, "*.txt")) + glob.glob(os.path.join(kb_dir, "*.md"))
-        for path in paths:
+        paths = (
+            glob.glob(os.path.join(kb_dir, "*.txt")) +
+            glob.glob(os.path.join(kb_dir, "*.md"))
+        )
+        for path in sorted(paths):
             try:
                 with open(path, "r", encoding="utf-8") as f:
                     content = f.read()
                 if content.strip():
                     texts.append((os.path.basename(path), content))
             except Exception as e:
-                print(f"Could not read {path}: {e}")
+                logger.warning(f"Could not read {path}: {e}")
 
-    # If no files found, fall back to built-in demo content
     if not texts:
-        print("No KB files found. Using built-in demo content.")
+        logger.warning("No KB files found. Using built-in demo content.")
         demo_text = """
         Welcome to the Self-Service KB Assistant.
 
@@ -76,7 +194,7 @@ def load_kb_texts(kb_dir: str = KB_DIR) -> List[Tuple[str, str]]:
 
         Example use cases for a KB assistant:
         - Agents quickly searching for internal procedures.
-        - Customers asking "how do I…" style questions.
+        - Customers asking "how do I..." style questions.
         - Managers analyzing gaps in documentation based on repeated queries.
         """
         texts.append(("demo_content.txt", demo_text))
@@ -84,23 +202,50 @@ def load_kb_texts(kb_dir: str = KB_DIR) -> List[Tuple[str, str]]:
     return texts
 
 
-# -----------------------------
+# ----------------------------- 
 # KB INDEX
-# -----------------------------
+# ----------------------------- 
 
 class KBIndex:
-    def __init__(self, model_name: str = EMBEDDING_MODEL_NAME):
-        print("Loading embedding model...")
+    def __init__(self, kb_dir: str = KB_DIR, model_name: str = EMBEDDING_MODEL_NAME):
+        self.kb_dir = kb_dir
+        self.model_name = model_name
+        self._lock = threading.Lock()  # Protect concurrent generate calls
+
+        logger.info("Loading embedding model...")
         self.model = SentenceTransformer(model_name)
-        print("Embedding model loaded.")
+        logger.info("Embedding model loaded.")
+
         self.chunks: List[str] = []
         self.chunk_sources: List[str] = []
-        self.embeddings = None
+        self.embeddings: Optional[np.ndarray] = None
+        self._fingerprint: Optional[str] = None
+
         self.build_index()
 
-    def build_index(self):
-        """Load KB texts, split into chunks, and build an embedding index."""
-        texts = load_kb_texts(KB_DIR)
+    def build_index(self, force: bool = False):
+        """
+        Load KB texts, split into chunks, and build an embedding index.
+        Uses disk cache keyed on KB fingerprint to skip re-encoding on cold start.
+        Pass force=True to rebuild regardless of cache.
+        """
+        texts = load_kb_texts(self.kb_dir)
+        fingerprint = kb_fingerprint(self.kb_dir)
+
+        # Try loading from cache if fingerprint matches
+        if not force and os.path.exists(_CACHE_FILE):
+            try:
+                cached = np.load(_CACHE_FILE, allow_pickle=True)
+                if str(cached.get("fingerprint")) == fingerprint:
+                    self.chunks = list(cached["chunks"])
+                    self.chunk_sources = list(cached["chunk_sources"])
+                    self.embeddings = cached["embeddings"]
+                    self._fingerprint = fingerprint
+                    logger.info(f"✅ Loaded {len(self.chunks)} chunks from embedding cache.")
+                    return
+            except Exception as e:
+                logger.warning(f"Cache load failed: {e}. Rebuilding index.")
+
         all_chunks = []
         all_sources = []
 
@@ -110,278 +255,300 @@ class KBIndex:
                 all_sources.append(source_name)
 
         if not all_chunks:
-            print("⚠️ No chunks found for KB index.")
+            logger.warning("⚠️ No chunks found for KB index.")
             self.chunks = []
             self.chunk_sources = []
             self.embeddings = None
             return
 
-        print(f"Creating embeddings for {len(all_chunks)} chunks...")
-        embeddings = self.model.encode(all_chunks, show_progress_bar=False, convert_to_numpy=True)
+        logger.info(f"Creating embeddings for {len(all_chunks)} chunks...")
+        embeddings = self.model.encode(
+            all_chunks,
+            show_progress_bar=False,
+            convert_to_numpy=True,
+            batch_size=64,
+        )
         self.chunks = all_chunks
         self.chunk_sources = all_sources
         self.embeddings = embeddings
-        print("KB index ready.")
+        self._fingerprint = fingerprint
+
+        # Save cache
+        try:
+            os.makedirs(_CACHE_DIR, exist_ok=True)
+            np.savez_compressed(
+                _CACHE_FILE,
+                chunks=np.array(all_chunks, dtype=object),
+                chunk_sources=np.array(all_sources, dtype=object),
+                embeddings=embeddings,
+                fingerprint=np.array(fingerprint),
+            )
+            logger.info("Embedding cache saved.")
+        except Exception as e:
+            logger.warning(f"Could not save embedding cache: {e}")
+
+        logger.info("KB index ready.")
+
+    def reload(self):
+        """Rebuild the index from disk, bypassing cache."""
+        logger.info("Reloading KB index...")
+        self.build_index(force=True)
+        logger.info("KB index reloaded.")
 
     def search(self, query: str, top_k: int = TOP_K) -> List[Tuple[str, str, float]]:
         """Return top-k (chunk, source_name, score) for a given query."""
         if not query.strip():
             return []
-
-        if self.embeddings is None or not len(self.chunks):
+        if self.embeddings is None or not self.chunks:
             return []
 
-        query_vec = self.model.encode([query], show_progress_bar=False, convert_to_numpy=True)[0]
+        query_vec = self.model.encode(
+            [query], show_progress_bar=False, convert_to_numpy=True
+        )[0]
 
-        # Cosine similarity
-        dot_scores = np.dot(self.embeddings, query_vec)
-        norm_docs = np.linalg.norm(self.embeddings, axis=1)
-        norm_query = np.linalg.norm(query_vec) + 1e-10
-        scores = dot_scores / (norm_docs * norm_query + 1e-10)
+        # Correct cosine similarity: single epsilon, applied once
+        norm_docs  = np.linalg.norm(self.embeddings, axis=1)
+        norm_query = np.linalg.norm(query_vec)
+        denom = norm_docs * norm_query
+        denom = np.where(denom == 0, 1e-10, denom)  # avoid divide-by-zero cleanly
+        scores = np.dot(self.embeddings, query_vec) / denom
 
         top_idx = np.argsort(scores)[::-1][:top_k]
-        results = []
-        for idx in top_idx:
-            results.append((self.chunks[idx], self.chunk_sources[idx], float(scores[idx])))
+        return [
+            (self.chunks[i], self.chunk_sources[i], float(scores[i]))
+            for i in top_idx
+        ]
+
 
-        return results
+# ----------------------------- 
+# LLM SETUP
+# ----------------------------- 
 
+logger.info("Initializing KB index...")
+kb_index = KBIndex(kb_dir=KB_DIR, model_name=EMBEDDING_MODEL_NAME)
 
-# Initialize KB index
-print("Initializing KB index...")
-kb_index = KBIndex()
+logger.info("Loading LLM for answer generation...")
+llm_available = False
+llm_tokenizer = None
+llm_model = None
+_llm_lock = threading.Lock()  # Serialize LLM generate calls
 
-# Initialize LLM for answer generation
-print("Loading LLM for answer generation...")
 try:
     from transformers import AutoTokenizer, AutoModelForCausalLM
     import torch
-    
-    # Use a small but capable model for faster responses
-    LLM_MODEL_NAME = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"  # Fast and good quality
-    
-    print(f"Loading {LLM_MODEL_NAME}...")
+
+    logger.info(f"Loading {LLM_MODEL_NAME}...")
     llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
     llm_model = AutoModelForCausalLM.from_pretrained(
         LLM_MODEL_NAME,
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
         device_map="auto" if torch.cuda.is_available() else None,
     )
-    
+
     if not torch.cuda.is_available():
         llm_model = llm_model.to("cpu")
-    
+
     llm_model.eval()
-    print(f"✅ LLM loaded successfully on {'GPU' if torch.cuda.is_available() else 'CPU'}")
     llm_available = True
-    
+    logger.info(f"✅ LLM loaded on {'GPU' if torch.cuda.is_available() else 'CPU'}")
+
 except Exception as e:
-    print(f"⚠️ Could not load LLM: {e}")
-    print("⚠️ Will use fallback mode (direct retrieval)")
-    llm_available = False
-    llm_tokenizer = None
-    llm_model = None
+    logger.warning(f"⚠️ Could not load LLM: {e}. Will use fallback retrieval mode.")
 
-print("✅ KB Assistant ready!")
 
-# -----------------------------
-# CHAT LOGIC (With LLM Answer Generation)
-# -----------------------------
+# ----------------------------- 
+# CHAT LOGIC
+# ----------------------------- 
 
-def clean_context(text: str) -> str:
-    """Clean up text for context, removing markdown and excess whitespace."""
-    # Remove markdown headers
-    text = text.replace('#', '')
-    # Remove multiple spaces
-    text = ' '.join(text.split())
-    return text.strip()
-
-
-def generate_answer_with_llm(query: str, context: str, sources: List[str]) -> str:
+def build_llm_prompt(query: str, context: str, history: list) -> str:
     """
-    Generate a natural, conversational answer using LLM based on retrieved context.
+    Build a TinyLlama chat prompt that includes recent conversation history
+    so the model can resolve follow-up questions.
     """
+    system = (
+        "You are a helpful knowledge base assistant. "
+        "Answer the user's question based ONLY on the provided context. "
+        "Be conversational, clear, and concise. "
+        "If the context doesn't contain enough information, say so honestly."
+    )
+
+    # Include up to last 3 history turns to stay within context budget
+    history_text = ""
+    for turn in history[-3:]:
+        user_msg = turn.get("role") == "user" and turn.get("content", "")
+        asst_msg = turn.get("role") == "assistant" and turn.get("content", "")
+        if user_msg:
+            history_text += f"<|user|>\n{turn['content']}\n</s>\n"
+        elif asst_msg:
+            history_text += f"<|assistant|>\n{turn['content']}\n</s>\n"
+
+    prompt = (
+        f"<|system|>\n{system}\n</s>\n"
+        f"{history_text}"
+        f"<|user|>\n"
+        f"Context from knowledge base:\n{context}\n\n"
+        f"Question: {query}\n"
+        f"</s>\n"
+        f"<|assistant|>\n"
+    )
+    return prompt
+
+
+def generate_answer_with_llm(
+    query: str,
+    context: str,
+    sources: List[str],
+    history: list,
+) -> Optional[str]:
+    """Generate a natural answer using the local LLM. Returns None on failure."""
     if not llm_available:
         return None
-    
-    # Create a focused prompt
-    prompt = f"""<|system|>
-You are a helpful knowledge base assistant. Answer the user's question based ONLY on the provided context. Be conversational, clear, and concise. If the context doesn't contain enough information, say so.
-</s>
-<|user|>
-Context from knowledge base:
-{context}
-
-Question: {query}
-</s>
-<|assistant|>
-"""
-    
+
+    import torch
+
+    prompt = build_llm_prompt(query, context, history)
+
     try:
-        # Tokenize
-        inputs = llm_tokenizer(
-            prompt,
-            return_tensors="pt",
-            truncation=True,
-            max_length=1024
-        )
-        
-        if torch.cuda.is_available():
-            inputs = {k: v.to("cuda") for k, v in inputs.items()}
-        
-        # Generate
-        with torch.no_grad():
-            outputs = llm_model.generate(
-                **inputs,
-                max_new_tokens=256,
-                temperature=0.7,
-                top_p=0.9,
-                do_sample=True,
-                pad_token_id=llm_tokenizer.eos_token_id,
+        with _llm_lock:
+            inputs = llm_tokenizer(
+                prompt,
+                return_tensors="pt",
+                truncation=True,
+                max_length=1024,
             )
-        
-        # Decode
+
+            if torch.cuda.is_available():
+                inputs = {k: v.to("cuda") for k, v in inputs.items()}
+
+            with torch.no_grad():
+                outputs = llm_model.generate(
+                    **inputs,
+                    max_new_tokens=MAX_NEW_TOKENS,
+                    temperature=0.7,
+                    top_p=0.9,
+                    do_sample=True,
+                    pad_token_id=llm_tokenizer.eos_token_id,
+                )
+
         full_response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
-        
-        # Extract only the assistant's response
+
+        # Extract only the final assistant turn
         if "<|assistant|>" in full_response:
             answer = full_response.split("<|assistant|>")[-1].strip()
         else:
             answer = full_response.strip()
-        
-        # Clean up the answer
+
         answer = answer.replace("</s>", "").strip()
-        
-        # Add source attribution
+
         sources_text = ", ".join(sources)
-        final_answer = f"{answer}\n\n---\n📚 **Sources:** {sources_text}"
-        
-        return final_answer
-        
+        return f"{answer}\n\n---\n📚 **Sources:** {sources_text}"
+
     except Exception as e:
-        print(f"Error in LLM generation: {e}")
+        logger.error(f"LLM generation error: {e}")
         return None
 
 
 def format_fallback_answer(results: List[Tuple[str, str, float]]) -> str:
-    """
-    Fallback formatting when LLM is not available or fails.
-    """
+    """Format retrieved chunks into a readable answer when LLM is unavailable."""
     if not results:
-        return (
-            "I couldn't find any relevant information in the knowledge base.\n\n"
-            "**Try:**\n"
-            "- Rephrasing your question\n"
-            "- Using different keywords\n"
-            "- Breaking down complex questions"
+        return CONFIG["messages"]["no_answer"]
+
+    parts = []
+    seen_sources = []
+
+    for chunk, source, score in results:
+        cleaned = clean_context(chunk)
+        label = "🥇" if not parts else "💡"
+        confidence = (
+            "High" if score >= 0.6 else
+            "Medium" if score >= 0.4 else
+            "Low"
         )
-    
-    # Get best result
-    best_chunk, best_source, best_score = results[0]
-    
-    # Clean markdown
-    cleaned = clean_context(best_chunk)
-    
-    # Format nicely
-    answer = f"**From {best_source}:**\n\n{cleaned}"
-    
-    # Add other sources if available
-    if len(results) > 1:
-        other_sources = list(set([src for _, src, _ in results[1:]]))
-        if other_sources:
-            answer += f"\n\n💡 **Also see:** {', '.join(other_sources)}"
-    
-    return answer
-
-
-def build_answer(query: str) -> str:
+        parts.append(
+            f"{label} **From {source}** *(relevance: {confidence} — {score:.2f})*\n\n{cleaned}"
+        )
+        if source not in seen_sources:
+            seen_sources.append(source)
+
+    return "\n\n---\n\n".join(parts)
+
+
+def build_answer(query: str, history: list) -> str:
     """
-    Main answer generation function using LLM for natural responses.
-    
-    Process:
-    1. Retrieve relevant chunks from KB
-    2. Build context from top results
-    3. Use LLM to generate natural answer
-    4. Cite sources
+    Main RAG pipeline:
+    1. Retrieve relevant chunks
+    2. Filter by similarity threshold
+    3. Build context from ALL top results (up to TOP_K)
+    4. Generate answer via LLM (or fallback)
     """
-    # Step 1: Search the knowledge base
     results = kb_index.search(query, top_k=TOP_K)
-    
+
     if not results:
         return (
-            "I couldn't find any relevant information in the knowledge base to answer your question.\n\n"
+            "I couldn't find any relevant information in the knowledge base.\n\n"
             "**Suggestions:**\n"
             "- Try rephrasing with different words\n"
-            "- Check if the topic is covered in the KB\n"
-            "- Be more specific about what you're looking for"
+            "- Use more specific keywords\n"
+            "- Break complex questions into smaller ones"
         )
-    
-    # Step 2: Filter by similarity threshold
-    filtered_results = [
-        (chunk, src, score) 
-        for chunk, src, score in results 
-        if score >= MIN_SIMILARITY_THRESHOLD
+
+    filtered = [
+        (chunk, src, score)
+        for chunk, src, score in results
+        if score >= MIN_SIMILARITY
     ]
-    
-    if not filtered_results:
+
+    if not filtered:
         return (
             "I found some content, but it doesn't seem relevant enough to your question.\n\n"
             "Please try being more specific or using different keywords."
         )
-    
-    # Step 3: Build context from top results
+
+    # Build context from ALL filtered results (not just top 2), capped by CONTEXT_CHARS
     context_parts = []
     sources = []
-    
-    for chunk, source, score in filtered_results[:2]:  # Top 2 most relevant
+    for chunk, source, _ in filtered:  # uses all TOP_K results
         cleaned = clean_context(chunk)
         context_parts.append(cleaned)
         if source not in sources:
             sources.append(source)
-    
-    # Combine context (limit to 1000 chars for speed)
-    context = " ".join(context_parts)[:1000]
-    
-    # Step 4: Generate answer with LLM
+
+    context = " ".join(context_parts)[:CONTEXT_CHARS]
+
     if llm_available:
-        llm_answer = generate_answer_with_llm(query, context, sources)
+        llm_answer = generate_answer_with_llm(query, context, sources, history)
         if llm_answer:
             return llm_answer
-    
-    # Step 5: Fallback if LLM fails or unavailable
-    return format_fallback_answer(filtered_results)
 
+    return format_fallback_answer(filtered)
 
-def chat_respond(message: str, history):
-    """
-    Gradio ChatInterface callback.
-    
-    Args:
-        message: Latest user message (str)
-        history: List of previous messages (handled by Gradio)
-    
-    Returns:
-        Assistant's reply as a string
-    """
+
+def chat_respond(message: str, history: list) -> str:
+    """Gradio ChatInterface callback."""
     if not message or not message.strip():
         return "Please ask me a question about the knowledge base."
-    
+
+    # Guard against excessively long queries
+    query = message.strip()
+    if len(query) > MAX_QUERY_CHARS:
+        query = query[:MAX_QUERY_CHARS]
+        logger.warning(f"Query truncated to {MAX_QUERY_CHARS} characters.")
+
     try:
-        answer = build_answer(message.strip())
-        return answer
+        return build_answer(query, history or [])
     except Exception as e:
-        print(f"Error generating answer: {e}")
-        return f"Sorry, I encountered an error processing your question: {str(e)}"
+        logger.error(f"Error generating answer: {e}")
+        return "Sorry, something went wrong while processing your question. Please try again."
 
 
-# -----------------------------
+# ----------------------------- 
 # GRADIO UI
-# -----------------------------
+# ----------------------------- 
 
-description = """
-🚀 **Fast Knowledge Base Search Assistant**
+description = f"""
+🚀 **{CONFIG['client']['name']}**
 
-Ask questions and get instant answers from the knowledge base. This assistant uses semantic search to find the most relevant information quickly.
+Ask questions and get instant answers from the knowledge base using semantic search.
 
 **Tips for better results:**
 - Be specific in your questions
@@ -389,31 +556,37 @@ Ask questions and get instant answers from the knowledge base. This assistant us
 - Ask one question at a time
 """
 
-# Create ChatInterface (without 'type' parameter for compatibility)
+examples = [
+    q["query"]
+    for q in CONFIG.get("quick_actions", [])
+    if "query" in q
+] or [
+    "What makes a good knowledge base article?",
+    "How could a KB assistant help agents?",
+    "Why is self-service important for customer support?",
+]
+
 chat_interface = gr.ChatInterface(
     fn=chat_respond,
-    title="🤖 Self-Service KB Assistant",
+    title=f"🤖 {CONFIG['client']['name']}",
     description=description,
-    examples=[
-        "What makes a good knowledge base article?",
-        "How could a KB assistant help agents?",
-        "Why is self-service important for customer support?",
-    ],
+    examples=examples,
     cache_examples=False,
+    type="messages",  # Use modern messages format for proper history passing
 )
 
-# Launch
+# ----------------------------- 
+# LAUNCH
+# ----------------------------- 
+
 if __name__ == "__main__":
-    # Detect environment and launch appropriately
-    is_huggingface = os.getenv('SPACE_ID') is not None
-    is_container = os.path.exists('/.dockerenv') or os.getenv('KUBERNETES_SERVICE_HOST') is not None
-    
-    if is_huggingface:
-        print("🤗 Launching on HuggingFace Spaces...")
+    is_huggingface = os.getenv("SPACE_ID") is not None
+    is_container   = os.path.exists("/.dockerenv") or os.getenv("KUBERNETES_SERVICE_HOST") is not None
+
+    if is_huggingface or is_container:
+        env_label = "🤗 HuggingFace Spaces" if is_huggingface else "🐳 container"
+        logger.info(f"Launching on {env_label}...")
         chat_interface.launch(server_name="0.0.0.0", server_port=7860)
-    elif is_container:
-        print("🐳 Launching in container environment...")
-        chat_interface.launch(server_name="0.0.0.0", server_port=7860, share=False)
     else:
-        print("💻 Launching locally...")
+        logger.info("💻 Launching locally...")
         chat_interface.launch(share=False)