Update services/kb_creation.py

services/kb_creation.py

@@ -8,15 +8,15 @@ from docx import Document
 from sentence_transformers import SentenceTransformer
 import chromadb
 
-# ----------------------- ChromaDB setup -----------------------
+# ---------------------------- ChromaDB setup ----------------------------
 CHROMA_PATH = os.path.join(os.getcwd(), "chroma_db")
 client = chromadb.PersistentClient(path=CHROMA_PATH)
 collection = client.get_or_create_collection(name="knowledge_base")
 
-# ----------------------- Embedding model -----------------------
+# ---------------------------- Embedding model ----------------------------
 model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
 
-# ----------------------- BM25 (lightweight) -----------------------
+# ---------------------------- BM25 (lightweight) ----------------------------
 BM25_INDEX_FILE = os.path.join(CHROMA_PATH, "bm25_index.pkl")
 bm25_docs: List[Dict[str, Any]] = []
 bm25_inverted: Dict[str, List[int]] = {}
@@ -26,23 +26,25 @@ bm25_ready: bool = False
 BM25_K1 = 1.5
 BM25_B = 0.75
 
-# ----------------------- Utilities -----------------------
+# ---------------------------- Utilities ----------------------------
 def _tokenize(text: str) -> List[str]:
     if not text:
         return []
     text = text.lower()
     return re.findall(r"[a-z0-9]+", text)
 
+
 def _normalize_query(q: str) -> str:
     q = (q or "").strip().lower()
     q = re.sub(r"[^\w\s]", " ", q)
     q = re.sub(r"\s+", " ", q).strip()
     return q
 
+
 def _tokenize_meta_value(val: Optional[str]) -> List[str]:
     return _tokenize(val or "")
 
-# ----------------------- DOCX parsing & chunking -----------------------
+# ---------------------------- DOCX parsing & chunking ----------------------------
 def _split_by_sections(doc: Document) -> List[Tuple[str, List[str]]]:
     sections: List[Tuple[str, List[str]]] = []
     current_title = None
@@ -66,6 +68,7 @@ def _split_by_sections(doc: Document) -> List[Tuple[str, List[str]]]:
         sections = [("Document", all_text)]
     return sections
 
+
 def _chunk_text_with_context(doc_title: str, section_title: str, paragraphs: List[str], max_words: int = 900) -> List[str]:
     body = "\n".join(paragraphs).strip()
     if not body:
@@ -80,9 +83,9 @@ def _chunk_text_with_context(doc_title: str, section_title: str, paragraphs: Lis
         chunks = [body]
     return chunks
 
-# ----------------------- Intent & Module tagging -----------------------
+# ---------------------------- Intent & Module tagging ----------------------------
 SECTION_STEPS_HINTS = ["process steps", "procedure", "how to", "workflow", "instructions", "steps"]
-SECTION_ERRORS_HINTS = ["common errors", "resolution", "troubleshooting", "known issues", "escalation", "escalation path", "permissions", "access"]
+SECTION_ERRORS_HINTS = ["common errors", "resolution", "troubleshooting", "known issues", "common issues", "escalation", "escalation path", "permissions", "access"]
 
 PERMISSION_TERMS = [
     "permission", "permissions", "access", "access right", "authorization", "authorisation",
@@ -90,7 +93,6 @@ PERMISSION_TERMS = [
     "not allowed", "not authorized", "denied", "restrict"
 ]
 ERROR_TERMS = ["error", "issue", "fail", "failure", "not working", "cannot", "can't"]
-
 STEP_VERBS = ["navigate", "select", "scan", "verify", "confirm", "print", "move", "complete", "click", "open", "choose", "enter", "update", "save", "delete", "create", "attach", "assign"]
 
 MODULE_VOCAB = {
@@ -102,6 +104,7 @@ MODULE_VOCAB = {
     "replenishment": ["replenishment", "replenish"],
 }
 
+
 def _infer_intent_tag(section_title: str) -> str:
     st = (section_title or "").lower()
     if any(k in st for k in SECTION_STEPS_HINTS):
@@ -114,6 +117,7 @@ def _infer_intent_tag(section_title: str) -> str:
         return "purpose"
     return "neutral"
 
+
 def _derive_semantic_intent_from_text(text: str) -> Tuple[str, List[str]]:
     t = (text or "").lower()
     tags: List[str] = []
@@ -121,8 +125,10 @@ def _derive_semantic_intent_from_text(text: str) -> Tuple[str, List[str]]:
     if any(term in t for term in PERMISSION_TERMS):
         intent = "errors"
         tags.append("permissions")
-        if "role" in t: tags.append("role_access")
-        if "security" in t: tags.append("security")
+    if "role" in t:
+        tags.append("role_access")
+    if "security" in t:
+        tags.append("security")
     if intent == "neutral" and any(term in t for term in ERROR_TERMS):
         intent = "errors"
         tags.append("errors")
@@ -131,6 +137,7 @@ def _derive_semantic_intent_from_text(text: str) -> Tuple[str, List[str]]:
         tags.append("procedure")
     return intent, list(set(tags))
 
+
 def _derive_module_tags(text: str, filename: str, section_title: str) -> List[str]:
     tokens = " ".join([filename or "", section_title or "", text or ""]).lower()
     found = []
@@ -142,7 +149,7 @@ def _derive_module_tags(text: str, filename: str, section_title: str) -> List[st
             found = ["inventory"]
     return list(sorted(set(found)))
 
-# ----------------------- Ingestion -----------------------
+# ---------------------------- Ingestion ----------------------------
 def ingest_documents(folder_path: str) -> None:
     print(f"[KB] Checking folder: {folder_path}")
     files = [f for f in os.listdir(folder_path) if f.lower().endswith('.docx')]
@@ -150,19 +157,24 @@ def ingest_documents(folder_path: str) -> None:
     if not files:
         print("[KB] WARNING: No .docx files found. Please check the folder path.")
         return
+
     global bm25_docs, bm25_inverted, bm25_df, bm25_avgdl, bm25_ready
     bm25_docs, bm25_inverted, bm25_df = [], {}, {}
     bm25_avgdl, bm25_ready = 0.0, False
+
     for file in files:
         file_path = os.path.join(folder_path, file)
         doc_title = os.path.splitext(file)[0]
         doc = Document(file_path)
         sections = _split_by_sections(doc)
         total_chunks = 0
+
         for s_idx, (section_title, paras) in enumerate(sections):
             chunks = _chunk_text_with_context(doc_title, section_title, paras, max_words=900)
             total_chunks += len(chunks)
+
             base_intent = _infer_intent_tag(section_title)
+
             for c_idx, chunk in enumerate(chunks):
                 derived_intent, topic_tags = _derive_semantic_intent_from_text(chunk)
                 final_intent = base_intent
@@ -170,8 +182,8 @@ def ingest_documents(folder_path: str) -> None:
                     final_intent = "errors"
                 elif base_intent == "neutral" and derived_intent in ("steps", "prereqs"):
                     final_intent = derived_intent
-                module_tags = _derive_module_tags(chunk, file, section_title)
 
+                module_tags = _derive_module_tags(chunk, file, section_title)
                 embedding = model.encode(chunk).tolist()
                 doc_id = f"{file}:{s_idx}:{c_idx}"
                 meta = {
@@ -180,9 +192,9 @@ def ingest_documents(folder_path: str) -> None:
                     "chunk_index": c_idx,
                     "title": doc_title,
                     "collection": "SOP",
-                    "intent_tag": final_intent,                                      # str
-                    "topic_tags": ", ".join(topic_tags) if topic_tags else "",       # str (NOT list)
-                    "module_tags": ", ".join(module_tags) if module_tags else "",    # str (NOT list)
+                    "intent_tag": final_intent,  # str
+                    "topic_tags": ", ".join(topic_tags) if topic_tags else "",  # str (NOT list)
+                    "module_tags": ", ".join(module_tags) if module_tags else "",  # str (NOT list)
                 }
                 try:
                     collection.add(ids=[doc_id], embeddings=[embedding], documents=[chunk], metadatas=[meta])
@@ -192,10 +204,12 @@ def ingest_documents(folder_path: str) -> None:
                         collection.add(ids=[doc_id], embeddings=[embedding], documents=[chunk], metadatas=[meta])
                     except Exception as e2:
                         print(f"[KB] ERROR: Upsert failed for {doc_id}: {e2}")
+
                 tokens = _tokenize(chunk)
                 tf: Dict[str, int] = {}
                 for tkn in tokens:
                     tf[tkn] = tf.get(tkn, 0) + 1
+
                 idx = len(bm25_docs)
                 bm25_docs.append({
                     "id": doc_id,
@@ -211,11 +225,14 @@ def ingest_documents(folder_path: str) -> None:
                     if term not in seen:
                         bm25_df[term] = bm25_df.get(term, 0) + 1
                         seen.add(term)
+
         print(f"[KB] Ingested {file} → {total_chunks} chunks")
+
     N = len(bm25_docs)
     if N > 0:
         bm25_avgdl = sum(d["length"] for d in bm25_docs) / float(N)
         bm25_ready = True
+
     payload = {
         "bm25_docs": bm25_docs,
         "bm25_inverted": bm25_inverted,
@@ -230,7 +247,7 @@ def ingest_documents(folder_path: str) -> None:
     print(f"[KB] BM25 index saved: {BM25_INDEX_FILE}")
     print(f"[KB] Documents ingested. Total entries in Chroma: {collection.count()}")
 
-# ----------------------- BM25 load -----------------------
+# ---------------------------- BM25 load ----------------------------
 def _load_bm25_index() -> None:
     global bm25_docs, bm25_inverted, bm25_df, bm25_avgdl, bm25_ready
     if not os.path.exists(BM25_INDEX_FILE):
@@ -248,9 +265,10 @@ def _load_bm25_index() -> None:
     except Exception as e:
         print(f"[KB] WARNING: Could not load BM25 index: {e}")
 
+
 _load_bm25_index()
 
-# ----------------------- BM25 search -----------------------
+# ---------------------------- BM25 search ----------------------------
 def _bm25_score_for_doc(query_terms: List[str], doc_idx: int) -> float:
     if not bm25_ready or doc_idx < 0 or doc_idx >= len(bm25_docs):
         return 0.0
@@ -272,9 +290,10 @@ def _bm25_score_for_doc(query_terms: List[str], doc_idx: int) -> float:
         except Exception:
             idf = 1.0
         denom = tf + BM25_K1 * (1 - BM25_B + BM25_B * (dl / (bm25_avgdl or 1.0)))
-        score += idf * ((tf * (BM25_K1 + 1)) / (denom or 1.0))
+        score += idf * (((tf * (BM25_K1 + 1)) / (denom or 1.0)))
     return score
 
+
 def bm25_search(query: str, top_k: int = 50) -> List[Tuple[int, float]]:
     if not bm25_ready:
         return []
@@ -296,22 +315,19 @@ def bm25_search(query: str, top_k: int = 50) -> List[Tuple[int, float]]:
     scored.sort(key=lambda x: x[1], reverse=True)
     return scored[:top_k]
 
-# ----------------------- Semantic-only -----------------------
+# ---------------------------- Semantic-only ----------------------------
 def search_knowledge_base(query: str, top_k: int = 10) -> dict:
     query_embedding = model.encode(query).tolist()
     res = collection.query(
         query_embeddings=[query_embedding],
         n_results=top_k,
-        include=['documents', 'metadatas', 'distances']
+        include=['documents', 'metadatas', 'distances', 'ids']
     )
-    docs_ll = res.get("documents", [[ ]]) or [[ ]]
-    metas_ll = res.get("metadatas", [[ ]]) or [[ ]]
-    dists_ll = res.get("distances", [[ ]]) or [[ ]]
-    ids_ll = res.get("ids", [[ ]]) or [[ ]]
-    documents = docs_ll[0] if docs_ll else []
-    metadatas = metas_ll[0] if metas_ll else []
-    distances = dists_ll[0] if dists_ll else []
-    ids = ids_ll[0] if ids_ll else []
+    documents = (res.get("documents", [[]]) or [[]])[0]
+    metadatas = (res.get("metadatas", [[]]) or [[]])[0]
+    distances = (res.get("distances", [[]]) or [[]])[0]
+    ids = (res.get("ids", [[]]) or [[]])[0]
+
     if not ids and documents:
         synthesized = []
         for i, m in enumerate(metadatas):
@@ -320,6 +336,7 @@ def search_knowledge_base(query: str, top_k: int = 10) -> dict:
             idx = (m or {}).get("chunk_index", i)
             synthesized.append(f"{fn}:{sec}:{idx}")
         ids = synthesized
+
     print(f"[KB] search → {len(documents)} docs (top_k={top_k}); first distance: {distances[0] if distances else 'n/a'}; ids={len(ids)}")
     return {
         "documents": documents,
@@ -328,7 +345,7 @@ def search_knowledge_base(query: str, top_k: int = 10) -> dict:
         "ids": ids,
     }
 
-# ----------------------- Hybrid search (intent + module + action) -----------------------
+# ---------------------------- Hybrid search (intent + module + action) ----------------------------
 ACTION_SYNONYMS = {
     "create": ["create", "creation", "add", "new", "generate"],
     "update": ["update", "modify", "change", "edit"],
@@ -342,6 +359,7 @@ ERROR_INTENT_TERMS = [
     "escalation", "escalation path", "access right"
 ]
 
+
 def _detect_user_intent(query: str) -> str:
     q = (query or "").lower()
     if any(k in q for k in ERROR_INTENT_TERMS):
@@ -354,6 +372,7 @@ def _detect_user_intent(query: str) -> str:
         return "purpose"
     return "neutral"
 
+
 def _extract_actions(query: str) -> List[str]:
     q = (query or "").lower()
     found = []
@@ -362,6 +381,7 @@ def _extract_actions(query: str) -> List[str]:
             found.append(act)
     return found or []
 
+
 def _extract_modules_from_query(query: str) -> List[str]:
     q = (query or "").lower()
     found = []
@@ -372,6 +392,7 @@ def _extract_modules_from_query(query: str) -> List[str]:
         found = ["inventory"]
     return list(sorted(set(found)))
 
+
 def _intent_weight(meta: dict, user_intent: str) -> float:
     tag = (meta or {}).get("intent_tag", "neutral")
     if user_intent == "neutral":
@@ -380,13 +401,15 @@ def _intent_weight(meta: dict, user_intent: str) -> float:
         return 1.0
     if tag in ["purpose", "prereqs"] and user_intent in ["steps", "errors"]:
         return -0.6
-    st = (meta or {}).get("section", "").lower()
+    st = ((meta or {}).get("section", "") or "").lower()
     topics = (meta or {}).get("topic_tags", "") or ""
     topic_list = [t.strip() for t in topics.split(",") if t.strip()]
-    if user_intent == "errors" and (any(k in st for k in ["escalation", "permissions", "access", "known issues"]) or ("permissions" in topic_list)):
-        return 0.7
+    # Strongly prefer errors/escalation/permissions when the user intent is errors
+    if user_intent == "errors" and (any(k in st for k in ["escalation", "permissions", "access", "known issues", "common issues"]) or ("permissions" in topic_list)):
+        return 0.95
     return -0.2
 
+
 def _module_weight(meta: Dict[str, Any], user_modules: List[str]) -> float:
     if not user_modules:
         return 0.0
@@ -394,9 +417,10 @@ def _module_weight(meta: Dict[str, Any], user_modules: List[str]) -> float:
     doc_modules = [m.strip() for m in raw.split(",") if m.strip()] if isinstance(raw, str) else (raw or [])
     overlap = len(set(user_modules) & set(doc_modules))
     if overlap == 0:
-        return -0.4   # demote different modules to avoid wrong SOP
+        return -0.4  # demote different modules to avoid wrong SOP
     return 0.6 * overlap
 
+
 def _meta_overlap(meta: Dict[str, Any], q_terms: List[str]) -> float:
     fn_tokens = _tokenize_meta_value(meta.get("filename"))
     title_tokens = _tokenize_meta_value(meta.get("title"))
@@ -410,6 +434,7 @@ def _meta_overlap(meta: Dict[str, Any], q_terms: List[str]) -> float:
     inter = len(meta_tokens & qset)
     return inter / max(1, len(qset))
 
+
 def _action_weight(text: str, actions: List[str]) -> float:
     if not actions:
         return 0.0
@@ -427,6 +452,7 @@ def _action_weight(text: str, actions: List[str]) -> float:
                     score -= 0.8
     return score
 
+
 def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6, beta: float = 0.4) -> dict:
     norm_query = _normalize_query(query)
     q_terms = _tokenize(norm_query)
@@ -460,12 +486,13 @@ def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6
         bm25_id_to_text[d["id"]] = d["text"]
         bm25_id_to_meta[d["id"]] = d["meta"]
 
+    # Union of IDs from semantic and BM25
     union_ids = set(sem_ids) | set(bm25_id_to_norm.keys())
 
     gamma = 0.30  # meta overlap
     delta = 0.45  # intent boost (stronger for errors)
     epsilon = 0.30  # action weight
-    zeta = 0.50    # module weight (new)
+    zeta = 0.50  # module weight (new)
 
     combined_records_ext: List[Tuple[str, float, float, str, Dict[str, Any], float, float, float, float]] = []
     for cid in union_ids:
@@ -477,15 +504,19 @@ def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6
             sem_meta = sem_metas[pos] if pos < len(sem_metas) else {}
         else:
             sem_sim, sem_dist, sem_text, sem_meta = 0.0, None, "", {}
+
         bm25_sim = bm25_id_to_norm.get(cid, 0.0)
         bm25_text = bm25_id_to_text.get(cid, "")
         bm25_meta = bm25_id_to_meta.get(cid, {})
+
         text = sem_text if sem_text else bm25_text
         meta = sem_meta if sem_meta else bm25_meta
+
         m_overlap = _meta_overlap(meta, q_terms)
         intent_boost = _intent_weight(meta, user_intent)
         act_wt = _action_weight(text, actions)
         mod_wt = _module_weight(meta, user_modules)
+
         final_score = alpha * sem_sim + beta * bm25_sim + gamma * m_overlap + delta * intent_boost + epsilon * act_wt + zeta * mod_wt
         combined_records_ext.append(
             (cid, final_score, (sem_dist if sem_dist is not None else 999.0), text, meta, m_overlap, intent_boost, act_wt, mod_wt)
@@ -505,8 +536,8 @@ def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6
         total_action = sum(max(0.0, r[7]) for r in recs)
         total_module = sum(r[8] for r in recs)
         total_penalty = sum(min(0.0, r[6]) for r in recs) + sum(min(0.0, r[7]) for r in recs)
-        esc_weight = 0.3 if any("escalation" in (r[4] or {}).get("section", "").lower() for r in recs) else 0.0
-        perm_weight = 0.3 if any("permissions" in ((r[4] or {}).get("topic_tags") or []) for r in recs) else 0.0
+        esc_weight = 0.3 if any("escalation" in ((r[4] or {}).get("section", "")).lower() for r in recs) else 0.0
+        perm_weight = 0.3 if any("permissions" in (((r[4] or {}).get("topic_tags") or [])) for r in recs) else 0.0
         return total_score + 0.4 * total_overlap + 0.7 * total_intent + 0.5 * total_action + 0.6 * total_module + 0.3 * total_penalty + esc_weight + perm_weight
 
     best_doc, best_doc_prior = None, -1.0
@@ -516,7 +547,7 @@ def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6
             best_doc_prior, best_doc = p, fn
 
     best_recs = sorted(doc_groups.get(best_doc, []), key=lambda x: x[1], reverse=True)
-    other_recs = []
+    other_recs: List[Tuple[str, float, float, str, Dict[str, Any], float, float, float, float]] = []
     for fn, recs in doc_groups.items():
         if fn == best_doc:
             continue
@@ -525,6 +556,7 @@ def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6
 
     reordered = best_recs + other_recs
     top = reordered[:top_k]
+
     documents = [t[3] for t in top]
     metadatas = [t[4] for t in top]
     distances = [t[2] for t in top]
@@ -543,7 +575,7 @@ def hybrid_search_knowledge_base(query: str, top_k: int = 10, alpha: float = 0.6
         "actions": actions,
     }
 
-# ----------------------- Section fetch helpers -----------------------
+# ---------------------------- Section fetch helpers ----------------------------
 def get_section_text(filename: str, section: str) -> str:
     texts: List[str] = []
     for d in bm25_docs:
@@ -554,6 +586,7 @@ def get_section_text(filename: str, section: str) -> str:
                 texts.append(t)
     return "\n\n".join(texts).strip()
 
+
 def get_best_steps_section_text(filename: str) -> str:
     texts: List[str] = []
     for d in bm25_docs:
@@ -564,6 +597,7 @@ def get_best_steps_section_text(filename: str) -> str:
                 texts.append(t)
     return "\n\n".join(texts).strip()
 
+
 def get_best_errors_section_text(filename: str) -> str:
     texts: List[str] = []
     for d in bm25_docs:
@@ -577,6 +611,8 @@ def get_best_errors_section_text(filename: str) -> str:
             or "escalation" in sec
             or "permission" in sec
             or "access" in sec
+            or "known issues" in sec
+            or "common issues" in sec
             or ("permissions" in topic_list)
         ):
             t = (d.get("text") or "").strip()
@@ -584,7 +620,7 @@ def get_best_errors_section_text(filename: str) -> str:
                 texts.append(t)
     return "\n\n".join(texts).strip()
 
-# ----------------------- Admin helpers -----------------------
+# ---------------------------- Admin helpers ----------------------------
 def get_kb_runtime_info() -> Dict[str, Any]:
     return {
         "chroma_path": CHROMA_PATH,
@@ -595,6 +631,7 @@ def get_kb_runtime_info() -> Dict[str, Any]:
         "bm25_ready": bm25_ready,
     }
 
+
 def reset_kb(folder_path: str) -> Dict[str, Any]:
     result = {"status": "OK", "message": "KB reset and re-ingested"}
     try:
@@ -604,14 +641,16 @@ def reset_kb(folder_path: str) -> Dict[str, Any]:
             pass
         global collection
         collection = client.get_or_create_collection(name="knowledge_base")
+
         try:
             if os.path.isfile(BM25_INDEX_FILE):
                 os.remove(BM25_INDEX_FILE)
         except Exception as e:
             result.setdefault("warnings", []).append(f"bm25 index delete: {e}")
+
         os.makedirs(CHROMA_PATH, exist_ok=True)
         ingest_documents(folder_path)
         result["info"] = get_kb_runtime_info()
         return result
     except Exception as e:
-        return {"status": "ERROR", "error": f"{e}", "info": get_kb_runtime_info()}
+        return {"status": "ERROR", "error": f"{e}", "info": get_kb_runtime_info()}