Update prompt and streaming

app.py

@@ -2,7 +2,8 @@ import os
 import streamlit as st
 from together import Together
 from langchain_community.vectorstores import Chroma
-from langchain_community.embeddings import HuggingFaceBgeEmbeddings
+# Use the updated HuggingFace Embeddings class
+from langchain_huggingface import HuggingFaceEmbeddings
 
 # --- Configuration ---
 # TogetherAI API key (env var name pilotikval)
@@ -14,9 +15,9 @@ if not TOGETHER_API_KEY:
 # Initialize TogetherAI client
 client = Together(api_key=TOGETHER_API_KEY)
 
-# Embeddings setup
+# Embeddings setup (new huggingface integration)
 EMBED_MODEL_NAME = "BAAI/bge-base-en"
-embeddings = HuggingFaceBgeEmbeddings(
+embeddings = HuggingFaceEmbeddings(
     model_name=EMBED_MODEL_NAME,
     encode_kwargs={"normalize_embeddings": True},
 )
@@ -54,17 +55,33 @@ vectorstore = Chroma(
 )
 retriever = vectorstore.as_retriever(search_kwargs={"k": 20})  # k=20
 
-# System prompt template with instruction for detailed long answers
+# System prompt template for long, detailed answers
 def build_system(context: str) -> dict:
-    return {
-        "role": "system",
-        "content": (
-            "You are an expert medical assistant. Provide a thorough, detailed, and complete answer. "
-            "If you don't know, say you don't know.\n"
-            "Use the following context from medical docs to answer.\n\n"
-            "Context:\n" + context
-        )
-    }
+    """
+    Build the system prompt that instructs the model to:
+    - Act as an expert medical assistant and attentive listener.
+    - Leverage all retrieved context to craft detailed, accurate, and empathetic responses.
+    - Ask clarifying follow-up questions if the user’s query is ambiguous.
+    - Structure answers clearly, using headings, bullet points, and step-by-step explanations.
+    - Cite relevant context sections when appropriate.
+    - Maintain conversational memory so follow-up queries build upon prior discussion.
+    """
+    prompt = (
+        "You are a world-class medical assistant and conversational partner. "
+        "Listen carefully to the user’s questions, reference the context below, and provide a thorough, evidence-based response. "
+        "If any part of the question is unclear, ask a clarifying question before proceeding. "
+        "Organize your answer with clear headings or bullet points, and refer back to specific context snippets as needed. "
+        "Always be empathetic, concise, and precise in your medical explanations. "
+        "Retain memory of previous user messages to support follow-up interactions.
+
+"
+        "=== Retrieved Context Start ===
+" 
+        + context +
+        "
+=== Retrieved Context End ==="
+    )
+    return {"role": "system", "content": prompt}
 
 st.title("🩺 DocChatter RAG (Streaming & Memory)")
 
@@ -83,27 +100,29 @@ with chat_tab:
     for msg in st.session_state.chat_history:
         st.chat_message(msg['role']).write(msg['content'])
 
-    # On new input
+    # Handle new user input
     if user_prompt:
         # Echo user
         st.chat_message("user").write(user_prompt)
         st.session_state.chat_history.append({"role": "user", "content": user_prompt})
 
-        # Retrieve top-k docs
-        docs = retriever.get_relevant_documents(user_prompt)
+        # Retrieve top-k documents
+        try:
+            docs = retriever.invoke({"query": user_prompt})
+        except Exception:
+            docs = retriever.get_relevant_documents(user_prompt)
         context = "\n---\n".join([d.page_content for d in docs])
 
-        # Build message sequence: system + full history
+        # Build TogetherAI message sequence
         messages = [build_system(context)]
         for m in st.session_state.chat_history:
             messages.append(m)
 
-        # Prepare streaming response
+        # Stream assistant response
         response_container = st.chat_message("assistant")
         stream_placeholder = response_container.empty()
         answer = ""
 
-        # Stream tokens
         for token in client.chat.completions.create(
             model="meta-llama/Llama-4-Scout-17B-16E-Instruct",
             messages=messages,
@@ -111,10 +130,15 @@ with chat_tab:
             temperature=0.1,
             stream=True
         ):
-            if hasattr(token, 'choices') and token.choices[0].delta.content:
-                delta = token.choices[0].delta.content
-                answer += delta
-                stream_placeholder.write(answer)
+            try:
+                choice = token.choices[0]
+                delta = getattr(choice.delta, 'content', '')
+                if delta:
+                    answer += delta
+                    stream_placeholder.write(answer)
+            except (IndexError, AttributeError):
+                # Skip empty or malformed token
+                continue
 
         # Save assistant response
         st.session_state.chat_history.append({"role": "assistant", "content": answer})