Update app.py

app.py

@@ -1,112 +1,177 @@
-import os
 import streamlit as st
-from langchain.prompts import ChatPromptTemplate, SystemMessagePromptTemplate, HumanMessagePromptTemplate
-from langchain_community.llms import HuggingFaceHub
 from langchain_huggingface import HuggingFaceEmbeddings
 from langchain_community.vectorstores import Chroma
-from langchain.memory import ConversationBufferMemory
-from langchain.chains import ConversationalRetrievalChain
+from langchain_community.llms import HuggingFaceHub
 from langchain.prompts import PromptTemplate
-
-# Environment
-os.environ["HUGGINGFACEHUB_API_TOKEN"] = os.getenv("HUGGINGFACEHUB_API_TOKEN")
-EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
-
-@st.cache_resource
-def get_response(question):
-    result = st.session_state.conversational_chain({"question": question})
-    response_text = result.get("answer", "Maaf, saya tidak mengetahui jawaban itu.")
-    
-    # Membersihkan jawaban dari teks yang tidak diperlukan
-    if "Answer:" in response_text:
-        response_text = response_text.split("Answer:")[1].strip()
-    return response_text
-
-
-def setup_vectorstore():
-    persist_directory = "./vector_db_dir"
-    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
-    return Chroma(persist_directory=persist_directory, embedding_function=embeddings)
-
-def chat_chain(vectorstore):
-    hf_hub_llm = HuggingFaceHub(
-        repo_id="SeaLLMs/SeaLLMs-v3-7B-Chat",
-        model_kwargs={"temperature": 1, "max_new_tokens": 1024},
+from langchain.chains import RetrievalQA, ConversationalRetrievalChain
+from langchain.memory import ConversationBufferMemory
+import warnings
+import os
+from dotenv import load_dotenv
+
+warnings.filterwarnings("ignore")
+load_dotenv()
+
+# Constants and configurations
+APP_TITLE = "💊 Asisten Kesehatan Feminacare"
+INITIAL_MESSAGE = """Halo! 👋 Saya adalah asisten kesehatan feminacare yang siap membantu Anda dengan informasi seputar kesehatan wanita. 
+Silakan ajukan pertanyaan apa saja dan saya akan membantu Anda dengan informasi yang akurat."""
+
+# Model configurations
+MODEL_NAME = "meta-llama/Meta-Llama-3-8B-Instruct"
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+TOP_K_DOCS = 5
+
+def initialize_models():
+    """Initialize the embedding model and vector store"""
+    data_directory = os.path.join(os.path.dirname(__file__), "vector_db_dir")
+    embedding_model = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+    vector_store = Chroma(
+        embedding_function=embedding_model, 
+        persist_directory=data_directory
+    )
+    return vector_store
+  
+def create_llm():
+    """Initialize the language model with optimized parameters"""
+    return HuggingFaceHub(
+        repo_id=MODEL_NAME,
+        model_kwargs={
+            "temperature": 0.7,  # Balanced between creativity and accuracy
+            "max_new_tokens": 1024,
+            "top_p": 0.9,
+            "frequency_penalty": 0.5
+        }
     )
 
-    prompt_template = """
-    You are an assistant specialized in women's health. Use the retrieved documents to answer the user's question.
-    If you don't know the answer or the information is not in the documents, reply with: "I'm sorry, I don't know."
-
-    Chat History:
-    {chat_history}
+# Improved prompt template with better context handling and response structure
+PROMPT_TEMPLATE = """
+Anda adalah asisten kesehatan profesional dengan nama Feminacare. 
+Berikan informasi yang akurat, jelas, dan bermanfaat berdasarkan konteks yang tersedia.
 
-    Question:
-    {question}
+Context yang tersedia:
+{context}
 
-    Answer:"""
-    prompt = PromptTemplate(input_variables=["chat_history", "question"], template=prompt_template)
+Chat historyt:
+{chat_history}
 
+Question: {question}
 
-    # qa_prompt = ChatPromptTemplate.from_messages(messages)
+Instruksi untuk menjawab:
+1. Berikan jawaban yang LENGKAP dan TERSTRUKTUR
+2. Selalu sertakan SUMBER informasi dari konteks yang diberikan
+3. Jika informasi tidak tersedia dalam konteks, katakan: "Maaf, saya tidak memiliki informasi yang cukup untuk menjawab pertanyaan tersebut secara akurat. Silakan konsultasi dengan tenaga kesehatan untuk informasi lebih lanjut."
+4. Gunakan bahasa yang mudah dipahami
+5. Jika relevan, berikan poin-poin penting menggunakan format yang rapi
+6. Akhiri dengan anjuran untuk konsultasi dengan tenaga kesehatan jika diperlukan
 
-    retriever = vectorstore.as_retriever(
-        search_type="similarity",
-        search_kwargs={"k": 2}
-    )
+Answer:
+"""
 
+def setup_qa_chain(vector_store):
+    """Set up the QA chain with improved configuration"""
     memory = ConversationBufferMemory(
-        llm=hf_hub_llm,
-        output_key="answer",
         memory_key="chat_history",
-        return_messages=True
+        return_messages=True,
+        output_key='answer'
     )
-
-    chain = ConversationalRetrievalChain.from_llm(
-        llm=hf_hub_llm,
-        retriever=retriever,
-        chain_type="stuff",
+    
+    custom_prompt = PromptTemplate(
+        template=PROMPT_TEMPLATE,
+        input_variables=["context", "question", "chat_history"]
+    )
+    
+    return ConversationalRetrievalChain.from_llm(
+        llm=create_llm(),
+        retriever=vector_store.as_retriever(
+            # search_type="mmr",  # Maximum Marginal Relevance for better diversity
+            # search_kwargs={"k": TOP_K_DOCS}
+        ),
         memory=memory,
-        verbose=True,
-        combine_docs_chain_kwargs={"prompt": prompt},
+        # combine_docs_chain_kwargs={"prompt": custom_prompt},
+        return_source_documents=True,
+        # return_generated_question=True,
     )
-    return chain
-
-# Streamlit App
-st.set_page_config(
-    page_title="Asisten Kesehatan Wanita",
-    page_icon="💊",
-    layout="centered"
-)
-
-st.title("💊 Asisten Kesehatan Wanita")
-
-if "chat_history" not in st.session_state:
-    st.session_state.chat_history = []
-
-if "vectorstore" not in st.session_state:
-    st.session_state.vectorstore = setup_vectorstore()
-
-if "conversational_chain" not in st.session_state:
-    st.session_state.conversational_chain = chat_chain(st.session_state.vectorstore)
-
-# Display Chat History
-for message in st.session_state.chat_history:
-    with st.chat_message(message["role"]):
-        st.markdown(message["content"])
-
-# User Input
-user_input = st.chat_input("Tanyakan sesuatu...")
 
-if user_input:
-    st.session_state.chat_history.append({"role": "user", "content": user_input})
-
-    with st.chat_message("user"):
-        st.markdown(user_input)
-
-
-    with st.chat_message("assistant"):
-        response = st.session_state.conversational_chain({"question": user_input})
-        assistant_response = response["answer"]
-        st.markdown(assistant_response)
-        st.session_state.chat_history.append({"role": "assistant", "content": assistant_response})
+def initialize_session_state():
+    """Initialize Streamlit session state"""
+    if "messages" not in st.session_state:
+        st.session_state.messages = [
+            {"role": "assistant", "content": INITIAL_MESSAGE}
+        ]
+    if "qa_chain" not in st.session_state:
+        vector_store = initialize_models()
+        st.session_state.qa_chain = setup_qa_chain(vector_store)
+
+def clear_chat():
+    """Clear chat history and memory"""
+    st.session_state.messages = [
+        {"role": "assistant", "content": INITIAL_MESSAGE}
+    ]
+    st.session_state.qa_chain.memory.clear()
+
+def create_ui():
+    """Create the Streamlit UI"""
+    st.set_page_config(page_title=APP_TITLE, page_icon="💊")
+    
+    # Custom CSS for better UI
+    st.markdown("""
+        <style>
+        .stApp {
+            max-width: 1200px;
+            margin: 0 auto;
+        }
+        .stChat {
+            border-radius: 10px;
+            padding: 20px;
+            margin: 10px 0;
+        }
+        </style>
+    """, unsafe_allow_html=True)
+    
+    st.title(APP_TITLE)
+    
+    # Sidebar
+    with st.sidebar:
+        st.title("ℹ️ Tentang Aplikasi")
+        st.markdown("""
+        Asisten digital ini dirancang untuk membantu Anda untuk berkonsultasi tentang kesehatan wanita.
+                    
+        _Catatan: Informasi yang diberikan bersifat umum. Selalu konsultasikan dengan tenaga kesehatan untuk saran yang lebih spesifik._
+        """)
+        
+        st.button('🗑️ Hapus Riwayat Chat', on_click=clear_chat)
+
+def handle_user_input(prompt):
+    """Handle user input and generate response"""
+    try:
+        with st.spinner("Sedang menyiapkan jawaban..."):
+            response = st.session_state.qa_chain({"question": prompt})
+            return response["answer"]
+    except Exception as e:
+        st.error("Maaf, terjadi kesalahan dalam memproses pertanyaan Anda. Silakan coba lagi.")
+        return None
+
+def main():
+    initialize_session_state()
+    create_ui()
+    
+    # Display chat messages
+    for message in st.session_state.messages:
+        with st.chat_message(message["role"]):
+            st.markdown(message["content"])
+    
+    # Handle user input
+    if prompt := st.chat_input("Ketik pertanyaan Anda di sini..."):
+        st.session_state.messages.append({"role": "user", "content": prompt})
+        with st.chat_message("user"):
+            st.markdown(prompt)
+        
+        response = handle_user_input(prompt)
+        if response:
+            with st.chat_message("assistant"):
+                st.markdown(response)
+            st.session_state.messages.append({"role": "assistant", "content": response})
+
+if __name__ == "__main__":
+    main()