Update app.py

app.py

@@ -1,4 +1,5 @@
 import os
+import re
 import gradio as gr
 import faiss
 import numpy as np
@@ -29,20 +30,37 @@ doc_names = []
 index = None
 
 
-def chunk_text(text, chunk_size=300):
+# -------------------------
+# Smarter Chunking
+# -------------------------
+def chunk_text(text, chunk_size=500):
+    sentences = re.split(r'(?<=[.!?]) +', text)
     chunks = []
-    for i in range(0, len(text), chunk_size):
-        chunks.append(text[i:i+chunk_size])
+    current_chunk = ""
+
+    for sentence in sentences:
+        if len(current_chunk) + len(sentence) < chunk_size:
+            current_chunk += " " + sentence
+        else:
+            chunks.append(current_chunk.strip())
+            current_chunk = sentence
+
+    if current_chunk:
+        chunks.append(current_chunk.strip())
+
     return chunks
 
 
+# -------------------------
+# Upload Document
+# -------------------------
 def add_document(file):
     global index, documents, doc_chunks, doc_names
 
     if file is None:
-        return "No file uploaded."
+        return "Please upload a .txt file."
 
-    # Reset everything (clean workspace)
+    # Reset workspace for clean demo behavior
     index = None
     documents = []
     doc_chunks = []
@@ -56,57 +74,87 @@ def add_document(file):
 
     chunks = chunk_text(content)
 
+    if len(chunks) == 0:
+        return "Uploaded file is empty."
+
     embeddings = embedder.encode(chunks)
     embeddings = np.array(embeddings).astype("float32")
 
     index = faiss.IndexFlatL2(embeddings.shape[1])
     index.add(embeddings)
 
+    clean_name = os.path.basename(file.name)
+
     for chunk in chunks:
         doc_chunks.append(chunk)
-        doc_names.append(file.name)
+        doc_names.append(clean_name)
 
-    documents.append(file.name)
+    documents.append(clean_name)
 
-    return f"Workspace reset. Uploaded: {file.name}"
+    return f"Workspace reset. Uploaded: {clean_name}"
 
 
+# -------------------------
+# List Documents
+# -------------------------
 def list_documents():
     if not documents:
         return "No documents uploaded."
     return "\n".join(documents)
 
 
+# -------------------------
+# Ask Question
+# -------------------------
 def ask_question(question):
+    global index
+
     if index is None:
-        return "No documents available.", "", ""
+        return "Please upload a document first.", "", ""
 
-    if not question.strip():
-        return "Empty question.", "", ""
+    if question is None or question.strip() == "":
+        return "Please enter a valid question.", "", ""
 
     query_embedding = embedder.encode([question])
     query_embedding = np.array(query_embedding).astype("float32")
 
-    D, I = index.search(query_embedding, k=3)
+    D, I = index.search(query_embedding, k=5)
+
+    retrieved_chunks = []
+    retrieved_sources = []
 
-    retrieved_chunks = [doc_chunks[i] for i in I[0]]
-    retrieved_sources = [doc_names[i] for i in I[0]]
+    for idx in I[0]:
+        if idx < len(doc_chunks):
+            retrieved_chunks.append(doc_chunks[idx])
+            retrieved_sources.append(doc_names[idx])
+
+    if not retrieved_chunks:
+        return "No relevant content found.", "", ""
 
     context = "\n\n".join(retrieved_chunks)
 
     prompt = f"""
-    Answer the question using only the context below.
-    If the answer is not in the context, say you don't know.
+You are a strict document-based question answering system.
+
+Use ONLY the provided context.
+Do NOT use outside knowledge.
+If the answer is not clearly present in the context, say exactly:
+"I don't know based on the provided documents."
 
-    Context:
-    {context}
+Context:
+{context}
 
-    Question:
-    {question}
-    """
+Question:
+{question}
 
-    response = llm.generate_content(prompt)
-    answer = response.text
+Answer clearly and concisely:
+"""
+
+    try:
+        response = llm.generate_content(prompt)
+        answer = response.text.strip()
+    except Exception as e:
+        return f"LLM Error: {str(e)}", "", ""
 
     source_display = "\n".join(set(retrieved_sources))
     chunk_display = "\n\n---\n\n".join(retrieved_chunks)
@@ -114,17 +162,31 @@ def ask_question(question):
     return answer, source_display, chunk_display
 
 
+# -------------------------
+# System Status
+# -------------------------
 def system_status():
-    llm_status = "OK"
+    backend_status = "OK"
     db_status = "OK" if index is not None else "No documents loaded"
-    return f"Backend: OK\nVector DB: {db_status}\nLLM: {llm_status}"
+
+    try:
+        llm.generate_content("Say OK")
+        llm_status = "OK"
+    except Exception:
+        llm_status = "LLM connection failed"
+
+    return f"Backend: {backend_status}\nVector DB: {db_status}\nLLM: {llm_status}"
 
 
 # -------------------------
 # Gradio UI
 # -------------------------
 with gr.Blocks() as demo:
-    gr.Markdown("# Mini Private Knowledge Q&A")
+    gr.Markdown("""
+# Mini Private Knowledge Q&A Workspace
+
+Upload a text document, ask questions, and see exactly which document and text snippet the answer comes from.
+""")
 
     with gr.Tab("Upload"):
         file_input = gr.File(file_types=[".txt"])