cehck

app.py

@@ -3,7 +3,6 @@ import streamlit as st
 from together import Together
 from langchain_community.vectorstores import Chroma
 from langchain_huggingface import HuggingFaceEmbeddings
-import streamlit.components.v1 as components
 
 # --- Configuration ---
 # TogetherAI API key (env var name pilotikval)
@@ -13,21 +12,14 @@ if not TOGETHER_API_KEY:
     st.stop()
 
 # Initialize TogetherAI client
-os.environ['TOGETHER_API_KEY'] = TOGETHER_API_KEY  # Set for Together to read
-client = Together()
+client = Together(api_key=TOGETHER_API_KEY)
 
 # Embeddings setup
 EMBED_MODEL_NAME = "BAAI/bge-base-en"
-
-@st.cache_resource
-def load_embeddings():
-    return HuggingFaceEmbeddings(
-        model_name=EMBED_MODEL_NAME,
-        model_kwargs={"device": "cpu"},
-        encode_kwargs={"normalize_embeddings": True},
-    )
-
-embeddings = load_embeddings()
+embeddings = HuggingFaceEmbeddings(
+    model_name=EMBED_MODEL_NAME,
+    encode_kwargs={"normalize_embeddings": True},
+)
 
 # Sidebar: select collection
 st.sidebar.title("DocChatter RAG")
@@ -55,16 +47,12 @@ persist_directory = dirs[collection]
 collection_name = cols[collection]
 
 # Load Chroma vector store
-@st.cache_resource
-def load_vectorstore(_embeddings, persist_dir, coll_name):
-    vectorstore = Chroma(
-        collection_name=coll_name,
-        persist_directory=persist_dir,
-        embedding_function=_embeddings
-    )
-    return vectorstore.as_retriever(search_kwargs={"k": 20})
-
-retriever = load_vectorstore(embeddings, persist_directory, collection_name)
+vectorstore = Chroma(
+    collection_name=collection_name,
+    persist_directory=persist_directory,
+    embedding_function=embeddings
+)
+retriever = vectorstore.as_retriever(search_kwargs={"k": 20})  # k=20
 
 # System prompt template
 
@@ -80,7 +68,7 @@ def build_system(context: str) -> dict:
     """
     prompt = f"""
 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.
+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.
@@ -91,43 +79,6 @@ Retain memory of previous user messages to support follow-up interactions.
 """
     return {"role": "system", "content": prompt}
 
-
-def copy_button(text_to_copy: str, button_id: str):
-    """
-    Creates a copy button using HTML and JavaScript.
-    """
-    # Escape text for JavaScript
-    escaped_text = text_to_copy.replace('\\', '\\\\').replace('`', '\\`').replace('$', '\\$')
-    
-    html_code = f"""
-    <div style="display: inline-block;">
-        <button id="copy-btn-{button_id}" onclick="copyToClipboard{button_id}()" 
-                style="background: none; border: none; cursor: pointer; font-size: 20px; 
-                       padding: 5px; margin-left: 10px; vertical-align: middle;"
-                title="Copy to clipboard">
-            📋
-        </button>
-        <span id="copied-msg-{button_id}" style="color: green; font-size: 12px; margin-left: 5px; display: none;">
-            ✓ Copied!
-        </span>
-    </div>
-    <script>
-    function copyToClipboard{button_id}() {{
-        const text = `{escaped_text}`;
-        navigator.clipboard.writeText(text).then(function() {{
-            document.getElementById('copied-msg-{button_id}').style.display = 'inline';
-            setTimeout(function() {{
-                document.getElementById('copied-msg-{button_id}').style.display = 'none';
-            }}, 2000);
-        }}, function(err) {{
-            console.error('Failed to copy: ', err);
-        }});
-    }}
-    </script>
-    """
-    components.html(html_code, height=40)
-
-
 st.title("🩺 DocChatter RAG (Streaming & Memory)")
 
 # Initialize chat history
@@ -142,16 +93,8 @@ chat_tab, clear_tab = st.tabs(["Chat", "Clear History"])
 
 with chat_tab:
     # Display existing chat
-    for idx, msg in enumerate(st.session_state.chat_history):
-        if msg['role'] == 'user':
-            st.chat_message("user").write(msg['content'])
-        else:  # assistant
-            with st.chat_message("assistant"):
-                col1, col2 = st.columns([0.95, 0.05])
-                with col1:
-                    st.write(msg['content'])
-                with col2:
-                    copy_button(msg['content'], f"msg{idx}")
+    for msg in st.session_state.chat_history:
+        st.chat_message(msg['role']).write(msg['content'])
 
     # Handle new user input
     if user_prompt:
@@ -160,7 +103,10 @@ with chat_tab:
         st.session_state.chat_history.append({"role": "user", "content": user_prompt})
 
         # Retrieve top-k documents
-        docs = retriever.invoke(user_prompt)
+        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 TogetherAI message sequence
@@ -191,11 +137,8 @@ with chat_tab:
 
         # Save assistant response
         st.session_state.chat_history.append({"role": "assistant", "content": answer})
-        
-        # Rerun to show copy button for the new message
-        st.rerun()
 
 with clear_tab:
     if st.button("🗑️ Clear chat history"):
         st.session_state.chat_history = []
-        st.rerun()
+        st.experimental_rerun()