Update app.py

app.py

@@ -13,15 +13,16 @@ from sentence_transformers import SentenceTransformer
 
 # -----------------------------
 # 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 = {
@@ -40,7 +41,6 @@ def load_config(path: str = _CONFIG_PATH) -> dict:
         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)
@@ -53,6 +53,7 @@ def load_config(path: str = _CONFIG_PATH) -> dict:
         logger.warning(f"{path} not found. Using default configuration.")
     return defaults
 
+
 CONFIG = load_config()
 
 KB_DIR               = CONFIG["kb"]["directory"]
@@ -64,15 +65,15 @@ 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
+MAX_QUERY_CHARS      = 1000
 
 # 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]:
     """
@@ -82,12 +83,14 @@ def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVE
     if not text:
         return []
 
-    # Split into sentences using punctuation as delimiters
+    # FIX: clamp overlap so chunk_size - overlap is always >= 1, preventing infinite loop
+    overlap = min(overlap, chunk_size - 1)
+
     sentence_endings = re.compile(r'(?<=[.!?])\s+')
     sentences = sentence_endings.split(text.strip())
 
     chunks = []
-    current_chunk = []
+    current_chunk: List[str] = []
     current_len = 0
 
     for sentence in sentences:
@@ -102,7 +105,6 @@ def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVE
                 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)
@@ -114,7 +116,7 @@ def chunk_text(text: str, chunk_size: int = CHUNK_SIZE, overlap: int = CHUNK_OVE
             if current_chunk:
                 chunks.append(" ".join(current_chunk))
             # Carry over last sentence(s) for overlap
-            overlap_words = []
+            overlap_words: List[str] = []
             overlap_len = 0
             for s in reversed(current_chunk):
                 if overlap_len + len(s) + 1 > overlap:
@@ -138,19 +140,25 @@ 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()
 
 
 def kb_fingerprint(kb_dir: str) -> str:
-    """Return an MD5 fingerprint of all KB filenames + modification times for cache invalidation."""
+    """
+    Return an MD5 fingerprint of all KB filenames + modification times.
+    Returns a sentinel string when the directory is empty or missing,
+    so cache comparisons never falsely match a populated index.
+    """
     paths = sorted(
         glob.glob(os.path.join(kb_dir, "*.txt")) +
         glob.glob(os.path.join(kb_dir, "*.md"))
     )
+    # FIX: return distinct sentinel when no files found, not md5('') which
+    # would match any other empty-directory state from a different run.
+    if not paths:
+        return "NO_KB_FILES"
     sig = "|".join(f"{p}:{os.path.getmtime(p)}" for p in paths)
     return hashlib.md5(sig.encode()).hexdigest()
 
@@ -202,15 +210,16 @@ def load_kb_texts(kb_dir: str = KB_DIR) -> List[Tuple[str, str]]:
     return texts
 
 
-# ----------------------------- 
+# -----------------------------
 # KB INDEX
-# ----------------------------- 
+# -----------------------------
 
 class KBIndex:
     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
+        # FIX: _lock removed from KBIndex — it was defined but never acquired.
+        # LLM concurrency is handled by the module-level _llm_lock instead.
 
         logger.info("Loading embedding model...")
         self.model = SentenceTransformer(model_name)
@@ -236,7 +245,10 @@ class KBIndex:
         if not force and os.path.exists(_CACHE_FILE):
             try:
                 cached = np.load(_CACHE_FILE, allow_pickle=True)
-                if str(cached.get("fingerprint")) == fingerprint:
+                # FIX: NpzFile stores arrays; extract scalar string correctly
+                # using item() instead of str() which would produce "array('...')"
+                cached_fp = str(cached["fingerprint"].item())
+                if cached_fp == fingerprint:
                     self.chunks = list(cached["chunks"])
                     self.chunk_sources = list(cached["chunk_sources"])
                     self.embeddings = cached["embeddings"]
@@ -246,8 +258,8 @@ class KBIndex:
             except Exception as e:
                 logger.warning(f"Cache load failed: {e}. Rebuilding index.")
 
-        all_chunks = []
-        all_sources = []
+        all_chunks: List[str] = []
+        all_sources: List[str] = []
 
         for source_name, content in texts:
             for chunk in chunk_text(content):
@@ -306,11 +318,10 @@ class KBIndex:
             [query], show_progress_bar=False, convert_to_numpy=True
         )[0]
 
-        # 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
+        denom = np.where(denom == 0, 1e-10, denom)
         scores = np.dot(self.embeddings, query_vec) / denom
 
         top_idx = np.argsort(scores)[::-1][:top_k]
@@ -320,9 +331,16 @@ class KBIndex:
         ]
 
 
-# ----------------------------- 
+# -----------------------------
 # LLM SETUP
-# ----------------------------- 
+# -----------------------------
+
+# Single import of torch at module level (not duplicated inside functions)
+try:
+    import torch
+    _torch_available = True
+except ImportError:
+    _torch_available = False
 
 logger.info("Initializing KB index...")
 kb_index = KBIndex(kb_dir=KB_DIR, model_name=EMBEDDING_MODEL_NAME)
@@ -331,34 +349,41 @@ logger.info("Loading LLM for answer generation...")
 llm_available = False
 llm_tokenizer = None
 llm_model = None
-_llm_lock = threading.Lock()  # Serialize LLM generate calls
+_llm_lock = threading.Lock()
 
-try:
-    from transformers import AutoTokenizer, AutoModelForCausalLM
-    import torch
+if _torch_available:
+    try:
+        from transformers import AutoTokenizer, AutoModelForCausalLM
 
-    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,
-    )
+        logger.info(f"Loading {LLM_MODEL_NAME}...")
+        llm_tokenizer = AutoTokenizer.from_pretrained(LLM_MODEL_NAME)
+
+        # FIX: ensure pad_token is set so eos_token_id is never None at generation time
+        if llm_tokenizer.pad_token is None:
+            llm_tokenizer.pad_token = llm_tokenizer.eos_token
 
-    if not torch.cuda.is_available():
-        llm_model = llm_model.to("cpu")
+        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()
-    llm_available = True
-    logger.info(f"✅ LLM loaded on {'GPU' if torch.cuda.is_available() else 'CPU'}")
+        llm_model.eval()
+        llm_available = True
+        logger.info(f"✅ LLM loaded on {'GPU' if torch.cuda.is_available() else 'CPU'}")
 
-except Exception as e:
-    logger.warning(f"⚠️ Could not load LLM: {e}. Will use fallback retrieval mode.")
+    except Exception as e:
+        logger.warning(f"⚠️ Could not load LLM: {e}. Will use fallback retrieval mode.")
+else:
+    logger.warning("⚠️ torch not available. Will use fallback retrieval mode.")
 
 
-# ----------------------------- 
+# -----------------------------
 # CHAT LOGIC
-# ----------------------------- 
+# -----------------------------
 
 def normalise_history(history: list) -> List[Tuple[str, str]]:
     """
@@ -369,7 +394,6 @@ def normalise_history(history: list) -> List[Tuple[str, str]]:
     if not history:
         return []
     if isinstance(history[0], dict):
-        # Gradio 6: walk role-by-role and pair up user + assistant turns
         pairs = []
         i = 0
         while i < len(history):
@@ -391,7 +415,7 @@ def normalise_history(history: list) -> List[Tuple[str, str]]:
 
 def build_llm_prompt(query: str, context: str, history: list) -> str:
     """
-    Build a TinyLlama chat prompt that includes recent conversation history
+    Build a TinyLlama chat prompt including recent conversation history
     so the model can resolve follow-up questions.
     """
     system = (
@@ -401,7 +425,6 @@ def build_llm_prompt(query: str, context: str, history: list) -> str:
         "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 user_msg, asst_msg in normalise_history(history)[-3:]:
         if user_msg:
@@ -409,7 +432,7 @@ def build_llm_prompt(query: str, context: str, history: list) -> str:
         if asst_msg:
             history_text += f"<|assistant|>\n{asst_msg}\n</s>\n"
 
-    prompt = (
+    return (
         f"<|system|>\n{system}\n</s>\n"
         f"{history_text}"
         f"<|user|>\n"
@@ -418,7 +441,6 @@ def build_llm_prompt(query: str, context: str, history: list) -> str:
         f"</s>\n"
         f"<|assistant|>\n"
     )
-    return prompt
 
 
 def generate_answer_with_llm(
@@ -431,11 +453,11 @@ def generate_answer_with_llm(
     if not llm_available:
         return None
 
-    import torch
-
     prompt = build_llm_prompt(query, context, history)
 
     try:
+        # FIX: keep decode() inside the lock so llm_tokenizer can't be
+        # reassigned/cleared by a reload() between generate() and decode()
         with _llm_lock:
             inputs = llm_tokenizer(
                 prompt,
@@ -454,10 +476,11 @@ def generate_answer_with_llm(
                     temperature=0.7,
                     top_p=0.9,
                     do_sample=True,
-                    pad_token_id=llm_tokenizer.eos_token_id,
+                    pad_token_id=llm_tokenizer.pad_token_id,  # FIX: use pad_token_id (always set above)
                 )
 
-        full_response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
+            # FIX: decode inside the lock — keeps tokenizer reference stable
+            full_response = llm_tokenizer.decode(outputs[0], skip_special_tokens=True)
 
         # Extract only the final assistant turn
         if "<|assistant|>" in full_response:
@@ -466,7 +489,6 @@ def generate_answer_with_llm(
             answer = full_response.strip()
 
         answer = answer.replace("</s>", "").strip()
-
         sources_text = ", ".join(sources)
         return f"{answer}\n\n---\n📚 **Sources:** {sources_text}"
 
@@ -481,21 +503,19 @@ def format_fallback_answer(results: List[Tuple[str, str, float]]) -> str:
         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
+            "High"   if score >= 0.6 else
             "Medium" if score >= 0.4 else
             "Low"
         )
         parts.append(
             f"{label} **From {source}** *(relevance: {confidence} — {score:.2f})*\n\n{cleaned}"
         )
-        if source not in seen_sources:
-            seen_sources.append(source)
+        # FIX: removed seen_sources — it was built but never used for anything
 
     return "\n\n---\n\n".join(parts)
 
@@ -531,10 +551,9 @@ def build_answer(query: str, history: list) -> str:
             "Please try being more specific or using different keywords."
         )
 
-    # Build context from ALL filtered results (not just top 2), capped by CONTEXT_CHARS
     context_parts = []
-    sources = []
-    for chunk, source, _ in filtered:  # uses all TOP_K results
+    sources: List[str] = []
+    for chunk, source, _ in filtered:
         cleaned = clean_context(chunk)
         context_parts.append(cleaned)
         if source not in sources:
@@ -555,7 +574,6 @@ def chat_respond(message: str, history: list) -> str:
     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]
@@ -568,9 +586,9 @@ def chat_respond(message: str, history: list) -> str:
         return "Sorry, something went wrong while processing your question. Please try again."
 
 
-# ----------------------------- 
+# -----------------------------
 # GRADIO UI
-# ----------------------------- 
+# -----------------------------
 
 description = f"""
 🚀 **{CONFIG['client']['name']}**
@@ -599,12 +617,11 @@ chat_interface = gr.ChatInterface(
     description=description,
     examples=examples,
     cache_examples=False,
-    type="messages",  # Use modern messages format for proper history passing
 )
 
-# ----------------------------- 
+# -----------------------------
 # LAUNCH
-# ----------------------------- 
+# -----------------------------
 
 if __name__ == "__main__":
     is_huggingface = os.getenv("SPACE_ID") is not None