Check check

app.py

@@ -3,7 +3,6 @@ import streamlit as st
 from together import Together
 from langchain.vectorstores import Chroma
 from langchain.embeddings import HuggingFaceBgeEmbeddings
-from langchain.chains import ConversationalRetrievalChain
 
 # --- Configuration ---
 TOGETHER_API_KEY = os.environ.get("pilotikval")
@@ -43,8 +42,8 @@ cols = {
     'OldMedicine': 'mrcppassmednotes'
 }
 
-persist_directory = dirs.get(collection)
-collection_name = cols.get(collection)
+persist_directory = dirs[collection]
+collection_name = cols[collection]
 
 # Load Chroma vector store
 vectorstore = Chroma(
@@ -55,47 +54,51 @@ vectorstore = Chroma(
 retriever = vectorstore.as_retriever(search_kwargs={"k":20})
 
 # System prompt template
-SYSTEM_PROMPT = (
-    "You are a helpful assistant for medical professionals. "
-    "Use the following context from medical documents to answer the question. "
-    "If you don't know, say you don't know.\n\nContext:\n{context}\n"
-)
+def build_system(context: str) -> dict:
+    return {"role": "system", "content": (
+        "You are a helpful assistant for medical professionals. "
+        "Use the following context from medical documents to answer the question. "
+        "If you don't know, say you don't know.\n\nContext:\n" + context
+    )}
 
 st.title("🩺 DocChatter RAG (Streaming)")
 
 # Initialize chat history
 if 'chat_history' not in st.session_state:
-    st.session_state.chat_history = []  # list of dicts {role, content}
+    st.session_state.chat_history = []
+
+# Get user input (must be at top-level)
+prompt = st.chat_input("Ask anything about your docs...")
 
-# Tabs
+# Tabs for display
 chat_tab, clear_tab = st.tabs(["Chat", "Clear History"])
+
 with chat_tab:
-    # Display history
+    # Display conversation so far
     for msg in st.session_state.chat_history:
         st.chat_message(msg['role']).write(msg['content'])
 
-    # User input
-    if prompt := st.chat_input("Ask anything about your docs..."):
-        # User message
+    # If user has just submitted
+    if prompt:
+        # Show user message
         st.chat_message("user").write(prompt)
         st.session_state.chat_history.append({"role": "user", "content": prompt})
 
-        # Retrieve relevant docs
+        # Retrieve context
         docs = retriever.get_relevant_documents(prompt)
         context = "\n---\n".join([d.page_content for d in docs])
 
-        # Build messages for TogetherAI
-        system_msg = {"role": "system", "content": SYSTEM_PROMPT.format(context=context)}
-        messages = [system_msg]
-        # include prior conversation
-        for msg in st.session_state.chat_history:
-            if msg['role'] in ('user', 'assistant'):
-                messages.append(msg)
-        # Prepare streaming response
+        # Assemble messages for streaming
+        messages = [build_system(context)]
+        # Include full history
+        for m in st.session_state.chat_history:
+            if m['role'] in ('user','assistant'):
+                messages.append(m)
+
+        # Stream assistant response
         response_container = st.chat_message("assistant")
         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,
@@ -105,10 +108,13 @@ with chat_tab:
                 delta = token.choices[0].delta.content
                 answer += delta
                 placeholder.write(answer)
-        # Save assistant message
+        # Save response
         st.session_state.chat_history.append({"role": "assistant", "content": answer})
 
 with clear_tab:
     if st.button("🗑️ Clear chat history"):
         st.session_state.chat_history = []
         st.experimental_rerun()
+
+# Optionally persist new embeddings
+# vectorstore.persist()