Update app.py

app.py

@@ -1,98 +1,211 @@
 import os
 from llama_cpp import Llama
-from sentence_transformers import SentenceTransformer
-import chromadb
-from chromadb.utils import embedding_functions
+# from sentence_transformers import SentenceTransformer # Keep commented for now due to RAM/complexity
+# import chromadb
+# from chromadb.utils import embedding_functions
 from fastapi import FastAPI, Query
 import gradio as gr
-from functools import lru_cache
 
-# === Config ===
-TOKEN_LIMIT = 256
-USE_RAG = True  # Toggle RAG mode
+# === Globals ===
+TOKEN_LIMIT = 256  # Default, overridden by slider
+
+# --- System Message for Bella ---
+SYSTEM_MESSAGE = """You are Bella, an expert AI assistant dedicated to supporting users across diverse domains such as coding, academic assignments (homework, computer science projects), and professional document creation. Your responses should always be accurate, comprehensive, and tailored to the user's needs, whether they are beginners or advanced learners. Prioritize clear explanations, practical advice, and step-by-step guidance to ensure user success. Do not engage in conversational filler; focus strictly on providing direct and valuable assistance."""
 
 # === Load LLM ===
-model = Llama.from_pretrained(
-    repo_id="openbmb/MiniCPM-V-2_6-gguf",
-    filename="ggml-model-Q4_K_M.gguf",
-    n_ctx=4096,
-    n_threads=os.cpu_count(),
+llm = None # Initialize llm to None
+try:
+    # MiniCPM-V models are generally used with `create_chat_completion`
+    # Llama.from_pretrained automatically handles downloading the GGUF from HF Hub
+    print("Loading MiniCPM-V-2_6-gguf model...")
+    llm = Llama.from_pretrained(
+        repo_id="openbmb/MiniCPM-V-2_6-gguf",
+        filename="ggml-model-Q4_K_M.gguf",
+        n_ctx=4096,
+        n_threads=os.cpu_count(),
+        verbose=False, # Set to True for more debug output
+        # `chat_format` can sometimes be inferred from model, but explicitly setting for safety:
+        # MiniCPM-V-2_6-gguf uses a specific chat template.
+        # Check the model card or a GGUF viewer for its precise chat template.
+        # This one is a common pattern for MiniCPM:
+        chat_format="chatml" # Or "llama-2" if that's what it uses, but chatml is more common
+                             # For MiniCPM specifically, it's <|im_start|>role\ncontent<|im_end|>
+                             # which is a variant of ChatML. Llama.cpp handles it if metadata exists.
+    )
+    print("MiniCPM-V-2_6-gguf model loaded successfully.")
+except Exception as e:
+    print(f"Error loading MiniCPM-V-2_6-gguf model: {e}")
+    # Consider raising an error or exiting if the model is critical
+    # sys.exit(1) # You might want to import sys for this
+
+# === RAG Setup (Commented out for free Space compatibility and initial focus) ===
+"""
+embedder = SentenceTransformer("all-MiniLM-L6-v2")
+client = chromadb.PersistentClient(path="chroma_db")
+col = client.get_or_create_collection(
+    "docs",
+    embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
+        model_name="all-MiniLM-L6-v2"
+    )
 )
+seed_texts = [
+    "MiniCPM‑V‑2_6‑gguf runs well on CPU via llama.cpp.",
+    "This model supports RAG with Chromadb and FastAPI + Gradio UI."
+]
+for t in seed_texts:
+    col.add(documents=[t], ids=[str(hash(t))])
+"""
+
+# === Query Function (Modified to use chat_completion) ===
+def llm_query(messages_history: list, max_tokens: int) -> str:
+    if llm is None:
+        return "Error: LLM model not loaded. Cannot generate response."
+
+    # context = "" # If RAG were active, you'd insert context here
+    # prompt = f"Context:\n{context}\n\nUser: {q}\nAssistant:" # Not needed with chat_completion
 
-# === Optional: RAG Setup ===
-if USE_RAG:
-    embedder = SentenceTransformer("all-MiniLM-L6-v2")
-    client = chromadb.PersistentClient(path="chroma_db")
-    col = client.get_or_create_collection(
-        "docs",
-        embedding_function=embedding_functions.SentenceTransformerEmbeddingFunction(
-            model_name="all-MiniLM-L6-v2"
+    try:
+        # Use create_chat_completion for streaming responses
+        response_generator = llm.create_chat_completion(
+            messages=messages_history, # Pass the entire prepared history
+            stream=True,
+            max_tokens=max_tokens,
+            temperature=0.7,
+            top_p=0.9,
+            # Add stop tokens if known for MiniCPM-V-2_6, e.g., for ChatML:
+            # The model's chat_format (if set correctly during Llama init) will often handle these.
+            # MiniCPM-V-2_6 uses <|im_end|>
+            stop=["<|im_end|>"]
         )
-    )
 
-    seed_texts = [
-        "MiniCPM‑V‑2_6‑gguf runs well on CPU via llama.cpp.",
-        "This model supports RAG with Chromadb and FastAPI + Gradio UI."
-    ]
-    for t in seed_texts:
-        try:
-            col.add(documents=[t], ids=[str(hash(t))])
-        except:
-            pass  # Avoid duplicates on restart
-
-    @lru_cache(maxsize=128)
-    def embed_query(q: str):
-        return embedder.encode(q)
-
-# === RAG or Vanilla Query ===
-def rag_query(q: str, max_tokens: int) -> str:
-    try:
-        context = ""
-        if USE_RAG:
-            results = col.query(
-                query_embeddings=[embed_query(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=max_tokens, temperature=0.7)
-        return out["choices"][0]["text"]
+        full_response = ""
+        for chunk in response_generator:
+            # 'delta' contains the new token
+            token = chunk["choices"][0]["delta"].get("content", "")
+            full_response += token
+            yield full_response # Yield partial response for streaming
+
     except Exception as e:
-        return f"[Error] {e}"
+        print(f"Error during LLM inference: {e}")
+        yield f"An error occurred during generation: {e}"
+
 
-# === FastAPI ===
+# === FastAPI App ===
+# Keep FastAPI part if you intend to expose an API endpoint.
+# Note: For Gradio-only Spaces, you don't strictly need FastAPI, but it's fine to keep.
 app = FastAPI()
 
 @app.get("/ask")
-def ask(q: str = Query(...), tokens: int = Query(TOKEN_LIMIT)):
-    tokens = min(max(32, tokens), 1024)
-    return {"answer": rag_query(q, tokens)}
+def ask_api(q: str = Query(...), tokens: int = Query(TOKEN_LIMIT)):
+    # This FastAPI endpoint will now use the chat history format internally,
+    # but for a single query it's just the system message and user message.
+    messages_for_api = [
+        {"role": "system", "content": SYSTEM_MESSAGE},
+        {"role": "user", "content": q}
+    ]
+    # For a non-streaming API, you'd run it to completion and return the final text
+    # Note: llm_query is a generator now, so you'd need to consume it for an API.
+    # For simplicity, if this API is purely for the Gradio frontend, it might not be necessary.
+    # If it is for external use and non-streaming, you'd adapt llm_query or call llm.create_chat_completion directly here.
+    try:
+        response = llm.create_chat_completion(
+            messages=messages_for_api,
+            max_tokens=tokens,
+            temperature=0.7,
+            top_p=0.9,
+            stop=["<|im_end|>"]
+        )
+        return {"answer": response["choices"][0]["message"]["content"]}
+    except Exception as e:
+        return {"error": str(e)}, 500
 
 @app.post("/ask")
-def ask_post(body: dict):
-    return ask(q=body.get("q", ""), tokens=body.get("tokens", TOKEN_LIMIT))
+def ask_post_api(body: dict):
+    return ask_api(q=body.get("q", ""), tokens=body.get("tokens", TOKEN_LIMIT))
+
 
 # === Gradio UI ===
 def chat_fn(message, history, max_tokens):
-    history = history or []
-    reply = rag_query(message, max_tokens)
-    new_history = history + [(message, reply)]
-    yield new_history, new_history, ""
+    # Gradio `history` for gr.Chatbot(type="messages") is already in OpenAI format:
+    # list of dictionaries like [{"role": "user", "content": "hello"}, {"role": "assistant", "content": "hi"}]
+
+    # 1. Add user message to history immediately for display
+    # This creates a new history list with the user's message, for immediate display
+    new_history = history + [{"role": "user", "content": message}]
+    yield new_history, gr.update(value="") # Clear textbox and update chatbot with user message
+
+    # 2. Prepare full message list for LLM, including the system message
+    messages_for_llm = [{"role": "system", "content": SYSTEM_MESSAGE}] + new_history
+
+    # 3. Call LLM for response (streaming)
+    full_bot_response = ""
+    for chunk in llm_query(messages_for_llm, max_tokens):
+        full_bot_response = chunk # `llm_query` now yields the full_response string
+        # Update the last assistant message in the history with the streaming content
+        # Note: new_history[-1] is the user's message. We need to add a new assistant message.
+        # This implies modifying the `new_history` list in place for streaming to work on the UI.
+        if len(new_history) > 0 and new_history[-1]["role"] == "user":
+            if len(new_history) == len(history) + 1: # First chunk after user message
+                new_history.append({"role": "assistant", "content": full_bot_response})
+            else: # Subsequent chunks for the same assistant message
+                new_history[-1]["content"] = full_bot_response
+        else: # Fallback if history state is unexpected (shouldn't happen with Chatbot type="messages")
+            new_history.append({"role": "assistant", "content": full_bot_response})
+
+        yield new_history, gr.update(value="") # Keep textbox cleared, update chatbot
+
+    # After generation is complete, ensure the final history state is sent
+    # (though the last yield in the loop should cover this)
+    # yield new_history, gr.update(value="") # This might be redundant but harmless
+
 
 with gr.Blocks() as demo:
-    gr.Markdown("### 🧠 MiniCPM‑V‑2_6‑gguf RAG Chat")
+    gr.Markdown(
+        """
+        # 🧠 Bella: MiniCPM-V-2_6-gguf AI Assistant
+        Welcome! I'm Bella, designed to assist you with coding, homework, computer science projects,
+        and document writing. I provide accurate, comprehensive, and tailored guidance.
+        """
+    )
 
-    chatbot = gr.Chatbot(label="Bella Lite", show_copy_button=True)
-    txt = gr.Textbox(placeholder="Ask me anything...", show_label=False)
-    token_slider = gr.Slider(64, 1024, value=TOKEN_LIMIT, step=16, label="Max Tokens")
+    # Use type="messages" for OpenAI-like chat history format
+    chatbot = gr.Chatbot(
+        height=500,
+        label="Bella's Responses",
+        type="messages", # Important for the history format
+        autoscroll=True,
+        resizable=True,
+        show_copy_button=True
+    )
+
+    # Simplified input section
+    msg = gr.Textbox(placeholder="Ask Bella a question...", show_label=False, submit_btn="Ask")
+
+    token_slider = gr.Slider(64, 1024, value=256, step=16, label="Max tokens")
 
-    txt.submit(chat_fn, [txt, chatbot, token_slider], [chatbot, chatbot, txt])
+    # Clear button
+    clear_btn = gr.ClearButton([msg, chatbot])
+
+    # Gradio submit event for streaming.
+    # The `outputs` here are: chatbot (for history updates) and msg (to clear it).
+    msg.submit(
+        fn=chat_fn,
+        inputs=[msg, chatbot, token_slider],
+        outputs=[chatbot, msg], # Order: chatbot first for history, msg second to clear input
+        queue=True # Set to True for streaming to work correctly in Gradio
+    )
 
+# When using FastAPI, Gradio is launched via FastAPI's startup event.
 @app.on_event("startup")
-def startup():
+async def startup_event(): # Use async def for FastAPI startup events
+    print("Starting Gradio app within FastAPI startup event...")
+    # This will launch Gradio within the Uvicorn server started by FastAPI
+    # The `share=True` is not needed in Hugging Face Spaces; it's handled automatically.
     demo.queue().launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))
+    print("Gradio app launch initiated.")
 
 if __name__ == "__main__":
     import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    # This block is for local testing. On Hugging Face Spaces, `app` is run by Gunicorn/Uvicorn.
+    print("Running FastAPI app locally (if not in Hugging Face Space)...")
+    uvicorn.run(app, host="0.0.0.0", port=int(os.getenv("PORT", 7860)))