Add Refs

app.py

@@ -5,7 +5,7 @@ import streamlit as st
 import os
 from urllib.parse import urlencode, urlparse, parse_qs
 
-st.set_page_config(page_title="ViBidLQA - Trợ lý AI hỗ trợ hỏi đáp luật Việt Nam", page_icon="./app/static/ai.jpg", layout="centered", initial_sidebar_state="collapsed")
+st.set_page_config(page_title="ViBidLQA - Trợ lý AI văn bản pháp luật Việt Nam", page_icon="./app/static/ai.jpg", layout="centered", initial_sidebar_state="collapsed")
 
 # ==== MÔI TRƯỜNG OAuth ====
 FB_APP_ID = os.getenv("FB_APP_ID")
@@ -28,6 +28,7 @@ url_api_introduce_system_model = f"{routing_response_module}/about_me"
 url_api_retrieval_model = f"{retrieval_module}/search"
 url_api_reranker_model = f"{reranker_module}/rerank"
 url_api_generation_model = f"{abs_QA_module}/answer"
+url_api_extract_reference_model = f"{routing_response_module}/extract_references_unstream"
 
 # ========== STREAMLIT UI ==========
 
@@ -117,17 +118,9 @@ def rerank_context(url_rerank_module, question, relevant_docs, top_k=5):
     else:
         return f"Lỗi tại Rerank module: {response.status_code} - {response.text}"
 
-def get_abstractive_answer(question):
-    retrieved_context = retrieve_context(question=question)
-    retrieved_context = [item['text'] for item in retrieved_context]
-
-    reranked_context = rerank_context(url_rerank_module=url_api_reranker_model,
-                                      question=question,
-                                      relevant_docs=retrieved_context,
-                                      top_k=5)[0]
-    
+def get_abstractive_answer(context, question):
     data = {
-        "context": reranked_context,
+        "context": context,
         "question": question
     }
 
@@ -138,6 +131,20 @@ def get_abstractive_answer(question):
     else:
         return f"Lỗi: {response.status_code} - {response.text}"
 
+def get_references(context, question, answer):
+    data = {
+        "context": context,
+        "question": question,
+        "answer": answer
+    }
+
+    response = requests.post(url_api_extract_reference_model, json=data)
+
+    if response.status_code == 200:
+        return response.json()["refs"]
+    else:
+        return f"Lỗi tại module Reference Extractor: {response.status_code} - {response.text}"
+
 def generate_text_effect(answer):
     words = answer.split()
     for i in range(len(words)):
@@ -171,120 +178,129 @@ if prompt := st.chat_input(placeholder='Tôi có thể giúp được gì cho b
     message_placeholder = st.empty()
     
     full_response = ""
-    try:
-        classify_result = classify_question(question=prompt).json()
+    classify_result = classify_question(question=prompt).json()
 
-        print(f"The type of user query: {classify_result}")
-    
-        if classify_result == "BIDDING_RELATED":
-            abs_answer = get_abstractive_answer(question=prompt)
-    
-            if isinstance(abs_answer, str):
-                full_response = abs_answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in abs_answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+    print(f"The type of user query: {classify_result}")
+
+    if classify_result == "BIDDING_RELATED":
+        retrieved_context = retrieve_context(question=prompt, top_k=10)
+        retrieved_context = [item['text'] for item in retrieved_context]
+        reranked_context = rerank_context(url_rerank_module=url_api_reranker_model,
+                                      question=prompt,
+                                      relevant_docs=retrieved_context,
+                                      top_k=5)[0]
+        
+        abs_answer = get_abstractive_answer(context=reranked_context, question=prompt)
+
+        if isinstance(abs_answer, str):
+            full_response = abs_answer
+            message_placeholder.markdown(f"""
+            <div class="assistant-message">
+                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                <div class="stMarkdown">{full_response}</div>
+            </div>
+            """, unsafe_allow_html=True)
+        else:
+            full_response = ""
+            for line in abs_answer.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        data_str = line[6:]
+                        if data_str == '[DONE]':
+                            break
+                        
+                        try:
+                            data = json.loads(data_str)
+                            token = data.get('token', '')
+                            full_response += token
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass
-    
-        elif classify_result == "ABOUT_CHATBOT":
-            answer = introduce_system(question=prompt)
-    
-            if isinstance(answer, str):
-                full_response = answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+                            message_placeholder.markdown(f"""
+                            <div class="assistant-message">
+                                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                                <div class="stMarkdown">{full_response}●</div>
+                            </div>
+                            """, unsafe_allow_html=True)
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass
+                        except json.JSONDecodeError:
+                            pass
         
+            refs = st.expander("Tài liệu tham khảo", expanded=False)
+            refs_list = get_references(context=reranked_context, question=prompt, answer=full_response)
+            print(refs_list)
+            refs.write(f"{refs_list}")
+
+    elif classify_result == "ABOUT_CHATBOT":
+        answer = introduce_system(question=prompt)
+
+        if isinstance(answer, str):
+            full_response = answer
+            message_placeholder.markdown(f"""
+            <div class="assistant-message">
+                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                <div class="stMarkdown">{full_response}</div>
+            </div>
+            """, unsafe_allow_html=True)
         else:
-            answer = response_unrelated_question(question=prompt)
+            full_response = ""
+            for line in answer.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        data_str = line[6:]
+                        if data_str == '[DONE]':
+                            break
+                        
+                        try:
+                            data = json.loads(data_str)
+                            token = data.get('token', '')
+                            full_response += token
+                            
+                            message_placeholder.markdown(f"""
+                            <div class="assistant-message">
+                                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                                <div class="stMarkdown">{full_response}●</div>
+                            </div>
+                            """, unsafe_allow_html=True)
+                            
+                        except json.JSONDecodeError:
+                            pass
     
-            if isinstance(answer, str):
-                full_response = answer
-                message_placeholder.markdown(f"""
-                <div class="assistant-message">
-                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                    <div class="stMarkdown">{full_response}</div>
-                </div>
-                """, unsafe_allow_html=True)
-            else:
-                full_response = ""
-                for line in answer.iter_lines():
-                    if line:
-                        line = line.decode('utf-8')
-                        if line.startswith('data: '):
-                            data_str = line[6:]
-                            if data_str == '[DONE]':
-                                break
+    else:
+        answer = response_unrelated_question(question=prompt)
+
+        if isinstance(answer, str):
+            full_response = answer
+            message_placeholder.markdown(f"""
+            <div class="assistant-message">
+                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                <div class="stMarkdown">{full_response}</div>
+            </div>
+            """, unsafe_allow_html=True)
+        else:
+            full_response = ""
+            for line in answer.iter_lines():
+                if line:
+                    line = line.decode('utf-8')
+                    if line.startswith('data: '):
+                        data_str = line[6:]
+                        if data_str == '[DONE]':
+                            break
+                        
+                        try:
+                            data = json.loads(data_str)
+                            token = data.get('token', '')
+                            full_response += token
+                            
+                            message_placeholder.markdown(f"""
+                            <div class="assistant-message">
+                                <img src="./app/static/ai.jpg" class="assistant-avatar" />
+                                <div class="stMarkdown">{full_response}●</div>
+                            </div>
+                            """, unsafe_allow_html=True)
                             
-                            try:
-                                data = json.loads(data_str)
-                                token = data.get('token', '')
-                                full_response += token
-                                
-                                message_placeholder.markdown(f"""
-                                <div class="assistant-message">
-                                    <img src="./app/static/ai.jpg" class="assistant-avatar" />
-                                    <div class="stMarkdown">{full_response}●</div>
-                                </div>
-                                """, unsafe_allow_html=True)
-                                
-                            except json.JSONDecodeError:
-                                pass
-    except Exception as e:
-        full_response = "Hiện tại trợ lý AI đang nghỉ xíu để sạc pin 🔌. Bạn hãy quay lại sau nhé!"
+                        except json.JSONDecodeError:
+                            pass   
 
     message_placeholder.markdown(f"""
     <div class="assistant-message">