Update app.py

app.py

@@ -6,13 +6,17 @@ from chromadb.utils import embedding_functions
 from fastapi import FastAPI, Query
 import gradio as gr
 
-# 1. Load model via llama-cpp-python
+# === Globals ===
+TOKEN_LIMIT = 256  # Default, overridden by slider
+
+# === Load LLM ===
 model = Llama.from_pretrained(
     repo_id="openbmb/MiniCPM-V-2_6-gguf",
-    filename="ggml-model-Q4_K_M.gguf", # Choose a specific GGUF file
+    filename="ggml-model-Q4_K_M.gguf",
     n_ctx=4096,
 )
-# 2. Setup RAG
+
+# === RAG Setup ===
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 client = chromadb.PersistentClient(path="chroma_db")
 col = client.get_or_create_collection(
@@ -21,7 +25,6 @@ col = client.get_or_create_collection(
         model_name="all-MiniLM-L6-v2"
     )
 )
-# Seed with example context
 seed_texts = [
     "MiniCPM‑V‑2_6‑gguf runs well on CPU via llama.cpp.",
     "This model supports RAG with Chromadb and FastAPI + Gradio UI."
@@ -29,44 +32,52 @@ seed_texts = [
 for t in seed_texts:
     col.add(documents=[t], ids=[str(hash(t))])
 
-def rag_query(q: str) -> str:
+# === Query Function ===
+def rag_query(q: str, max_tokens: int) -> str:
     results = col.query(
         query_embeddings=[embedder.encode(q)],
         n_results=3
     )
     context = "\n".join(results["documents"][0])
     prompt = f"Context:\n{context}\n\nUser: {q}\nAssistant:"
-    out = model.create_completion(prompt=prompt, max_tokens=256, temperature=0.7)
+    out = model.create_completion(prompt=prompt, max_tokens=max_tokens, temperature=0.7)
     return out["choices"][0]["text"]
 
-# 3. FastAPI app
+# === FastAPI App ===
 app = FastAPI()
 
 @app.get("/ask")
-def ask(q: str = Query(...)):
-    return {"answer": rag_query(q)}
+def ask(q: str = Query(...), tokens: int = Query(TOKEN_LIMIT)):
+    return {"answer": rag_query(q, tokens)}
 
 @app.post("/ask")
 def ask_post(body: dict):
-    return ask(q=body.get("q",""))
+    return ask(q=body.get("q", ""), tokens=body.get("tokens", TOKEN_LIMIT))
 
-# 4. Gradio UI
-def chat_fn(message, history):
-    reply = rag_query(message)
+# === Gradio UI ===
+def chat_fn(message, history, max_tokens):
+    reply = rag_query(message, max_tokens)
     history = history or []
-    history.append(("User", message))
-    history.append(("Assistant", reply))
-    return history, history
+    history.append((f"🧑 You", message))
+    history.append((f"🤖 Bot", reply))
+    return history, history, ""
 
 with gr.Blocks() as demo:
-    chatbot = gr.Chatbot()
-    txt = gr.Textbox(placeholder="Ask me...", show_label=False)
-    txt.submit(chat_fn, [txt, chatbot], [chatbot, chatbot])
-    gr.Markdown("### 🧠 MiniCPM‑V‑2_6‑gguf RAG Chat (GET & POST API support)")
+    gr.Markdown("### 🧠 MiniCPM‑V‑2_6‑gguf RAG Chat")
+
+    chatbot = gr.Chatbot(label="Chat", bubble_full_width=False)
+    with gr.Row():
+        txt = gr.Textbox(placeholder="Ask me...", show_label=False, scale=8)
+        send_btn = gr.Button("Send", scale=1)
+
+    token_slider = gr.Slider(64, 1024, value=256, step=16, label="Max tokens")
+
+    txt.submit(chat_fn, [txt, chatbot, token_slider], [chatbot, chatbot, txt])
+    send_btn.click(chat_fn, [txt, chatbot, token_slider], [chatbot, chatbot, txt])
 
 @app.on_event("startup")
 def startup():
-    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT",7860)))
+    demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))
 
 if __name__ == "__main__":
     import uvicorn