Spaces:
Sleeping
Sleeping
| import os | |
| import glob | |
| from typing import List, Tuple | |
| import gradio as gr | |
| import numpy as np | |
| from sentence_transformers import SentenceTransformer | |
| # ----------------------------- | |
| # CONFIG | |
| # ----------------------------- | |
| KB_DIR = "./kb" # optional: folder with .txt or .md files | |
| EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2" | |
| TOP_K = 3 # how many chunks to show per answer | |
| CHUNK_SIZE = 500 # characters | |
| CHUNK_OVERLAP = 100 # characters | |
| # ----------------------------- | |
| # 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.""" | |
| if not text: | |
| return [] | |
| chunks = [] | |
| start = 0 | |
| length = len(text) | |
| 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 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). | |
| """ | |
| 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: | |
| 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}") | |
| # If no files found, fall back to some built-in demo content | |
| if not texts: | |
| print("No KB files found. Using built-in demo content.") | |
| demo_text = """ | |
| Welcome to the Self-Service KB Assistant. | |
| This assistant is meant to help you find information inside a knowledge base. | |
| In a real setup, it would be connected to your own articles, procedures, | |
| troubleshooting guides and FAQs. | |
| Good knowledge base content is: | |
| - Clear and structured with headings, steps and expected outcomes. | |
| - Written in a customer-friendly tone. | |
| - Easy to scan, with short paragraphs and bullet points. | |
| - Maintained regularly to reflect product and process changes. | |
| Example use cases for a KB assistant: | |
| - Agents quickly searching for internal procedures. | |
| - Customers asking “how do I…” style questions. | |
| - Managers analyzing gaps in documentation based on repeated queries. | |
| """ | |
| texts.append(("demo_content.txt", demo_text)) | |
| return texts | |
| # ----------------------------- | |
| # KB INDEX | |
| # ----------------------------- | |
| class KBIndex: | |
| def __init__(self, model_name: str = EMBEDDING_MODEL_NAME): | |
| print("Loading embedding model...") | |
| self.model = SentenceTransformer(model_name) | |
| print("Model loaded.") | |
| self.chunks: List[str] = [] | |
| self.chunk_sources: List[str] = [] | |
| self.embeddings: np.ndarray | None = 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) | |
| all_chunks = [] | |
| all_sources = [] | |
| for source_name, content in texts: | |
| for chunk in chunk_text(content): | |
| all_chunks.append(chunk) | |
| all_sources.append(source_name) | |
| if not all_chunks: | |
| print("⚠️ 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) | |
| self.chunks = all_chunks | |
| self.chunk_sources = all_sources | |
| self.embeddings = embeddings | |
| print("KB index ready.") | |
| 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): | |
| return [] | |
| 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) | |
| 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 results | |
| kb_index = KBIndex() | |
| # ----------------------------- | |
| # CHAT LOGIC | |
| # ----------------------------- | |
| def build_answer(query: str) -> str: | |
| """Use the KB index to build a human-readable answer.""" | |
| results = kb_index.search(query, top_k=TOP_K) | |
| if not results: | |
| return ( | |
| "I couldn't find anything relevant in the knowledge base for this query yet.\n\n" | |
| "If this were connected to your real KB, this would be a good moment to:\n" | |
| "- Create a new article, or\n" | |
| "- Improve the existing documentation for this topic." | |
| ) | |
| intro = "Here’s what I found in the knowledge base:\n" | |
| bullets = [] | |
| for i, (chunk, source, score) in enumerate(results, start=1): | |
| bullets.append(f"{i}. From **{source}**:\n{chunk.strip()}\n") | |
| guidance = ( | |
| "\nYou can ask follow-up questions, or try a more specific query if this doesn't fully answer your question." | |
| ) | |
| return intro + "\n".join(bullets) + guidance | |
| def chat_respond(message: str, history): | |
| """ | |
| Gradio ChatInterface (type='messages') calls this with: | |
| - message: latest user message (str) | |
| - history: list of previous messages (handled internally by Gradio) | |
| We only need to return the assistant's reply as a string. | |
| """ | |
| answer = build_answer(message) | |
| return answer | |
| # ----------------------------- | |
| # GRADIO UI | |
| # ----------------------------- | |
| description = """ | |
| Ask questions as if you were talking to a knowledge base assistant. | |
| In a real scenario, this assistant would be connected to your own | |
| help center or internal documentation. Here, it's using a small demo | |
| knowledge base to show how retrieval-based self-service can work. | |
| """ | |
| chat = gr.ChatInterface( | |
| fn=chat_respond, | |
| title="Self-Service KB Assistant", | |
| description=description, | |
| type="messages", # use new-style message format | |
| examples=[ | |
| "What makes a good knowledge base article?", | |
| "How could a KB assistant help agents?", | |
| "Why is self-service important for customer support?", | |
| ], | |
| cache_examples=False, # avoid example pre-caching issues on HF Spaces | |
| ) | |
| if __name__ == "__main__": | |
| chat.launch() | |