chat_app.py

import os
import time
import random
import torch
import dotenv
import ollama
import logging
import requests
import streamlit as st

from typing import Optional, List
dotenv.load_dotenv()
logger = logging.getLogger(__name__)

# Default prompt context (unchanged)
DEFAULT_PROMPT_CONTEXT = """Bạn là một trợ lí AI pháp luật Việt Nam, có kiến thức về pháp luật Việt Nam.
Dựa vào ngữ cảnh hoặc tài liệu sau hãy trả lời câu hỏi người dùng
Ngữ cảnh: {context}
Câu hỏi: {question}
"""

class Config:
   # URL_RETRIEVE = str(os.getenv("URL_RETRIEVE", "http://202.191.56.254:9002/getPredictionOutput"))
    URL_RETRIEVE = str(os.getenv("URL_RETRIEVE", "http://0.0.0.0:9002/getPredictionOutput"))
    TOP_K = int(os.getenv("TOP_K", 5))


def prompt_model(top_k_chunks: Optional[List]):
    res = []
    for context in top_k_chunks:
        text = ''
        if context.get("diem_id", ""):
            text += "điểm " + context.get("diem_id", "") + " "
            
        if context.get("khoan_id", ""):
            text += "khoản " + str(int(context.get("khoan_id"))) + " "
        
        if context.get("diem_id", ""):
            text += context.get("diem_id", "") + " "
        
        if context.get("law_id", ""):
            if 'ttlt' in context.get("law_id", ""):
                text += "thông tư liên tịch " + context.get("law_id", "").upper() + " "
            elif 'tt' in context.get("law_id", ""):
                text += "thông tư " + context.get("law_id", "").upper() + " "
            elif 'nđ' in context.get("law_id", ""):
                text += "nghị định " + context.get("law_id", "").upper() + " "
            elif 'tb' in context.get("law_id", ""):
                text += "thông báo " + context.get("law_id", "").upper() + " "
            else:
                text += "luật số " + context.get("law_id", "").upper() + " "
            
        if context.get("title", "") and random.choice([1, 0]):
            text += context.get("title", "") + " "

        text += context.get("text", "") + " "
        res.append(text)
    return res[:Config.TOP_K]


def get_retrieval(query, config):
    retrieve = {"query": [query]}
    response = requests.post(config.URL_RETRIEVE, json=retrieve)
    if response.status_code == 200:
        return response.json()['predict'][0][0][0]['top_relevant_chunks']
    else:
        return []


def model_res_generator(prompt_template, retrieval_results, question):
    context = '\n'.join(prompt_model(retrieval_results))  
    input_model = prompt_template.format(context=context, question=question)
    
    stream = ollama.chat(
        model=st.session_state["model"],
        messages=[{'role': 'user', 'content': input_model}],
        options={
            'temperature': 0.0
        },
        stream=True,
    )

    full_response = ""
    for chunk in stream:
        chunk_content = chunk.get("message", {}).get("content", "")
        if chunk_content:
            full_response += chunk_content
            yield full_response


def process_input(prompt_template, user_input, config):
    if user_input:
        st.session_state.messages.append({"role": "user", "content": user_input})
        retrieval_results = get_retrieval(user_input, config)
        retrieval_results.sort(key=lambda x: x['bi_score'], reverse=True)
        st.session_state.retrieval_results = retrieval_results[:st.session_state.top_k]
        st.session_state.queries_and_results[user_input] = st.session_state.retrieval_results
        st.session_state.selected_query = user_input
        st.session_state.query_list = list(st.session_state.queries_and_results.keys())
        st.session_state.messages.append({"role": "assistant", "content": ""})

        for message in st.session_state.messages[-2:-1]:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])

        with st.chat_message("assistant"):
            message_placeholder = st.empty()
            for chunk in model_res_generator(prompt_template, retrieval_results, user_input):
                message_placeholder.markdown(chunk + "▌")
                st.session_state.messages[-1]["content"] = chunk
            message_placeholder.markdown(st.session_state.messages[-1]["content"])

        st.rerun() 


def reset_session_state(config):
    st.session_state.messages = [] 
    st.session_state.retrieval_results = []
    st.session_state.queries_and_results = {}
    st.session_state.selected_query = None
    st.session_state.top_k = config.TOP_K
    st.session_state.query_list = []  


def update_all_queries_results(config):
    """Cập nhật lại kết quả của tất cả các truy vấn khi top_k thay đổi"""
    for query in st.session_state.queries_and_results.keys():
        retrieval_results = get_retrieval(query, config)
        retrieval_results.sort(key=lambda x: x['bi_score'], reverse=True)
        st.session_state.queries_and_results[query] = retrieval_results[:st.session_state.top_k]
    st.rerun()



if __name__ == "__main__":
    st.set_page_config(page_title="AsklexAI", page_icon="🧊", layout="centered")
    config = Config()
    models = [model["name"] for model in ollama.list()["models"]]

    # Initialize session state variables if not already initialized
    if 'query_list' not in st.session_state:
        st.session_state.query_list = []  
    if 'messages' not in st.session_state:
        st.session_state.messages = [] 
    if 'queries_and_results' not in st.session_state:
        st.session_state.queries_and_results = {}
    if 'model' not in st.session_state:
        st.session_state.model = "" 
    if 'retrieval_results' not in st.session_state:
        st.session_state.retrieval_results = []
    if 'selected_query' not in st.session_state:
        st.session_state.selected_query = None
    if 'top_k' not in st.session_state:
        st.session_state.top_k = 5
    if 'custom_prompt' not in st.session_state:
        st.session_state.custom_prompt = DEFAULT_PROMPT_CONTEXT  # Set default custom prompt if not yet set

    st.markdown("""
    <style>
    .css-1aumxhk {
        background-color: #F0F2F6;  /* Light blue-gray background */
    }
    .css-1aumxhk .stMarkdown {
        color: #333;  /* Darker text for better readability */
    }
    .css-1aumxhk .stButton>button {
        background-color: #4A90E2;  /* Blue button color */
        color: white;
    }
    .css-1aumxhk .stSelectbox>div {
        background-color: white;
        border-color: #4A90E2;
    }
    </style>
    """, unsafe_allow_html=True)


    with st.sidebar:
        if st.button("New Chat", use_container_width=True):
            reset_session_state(config)

        st.subheader("Danh sách các truy vấn:")
        selected_query = st.selectbox(
            "Chọn truy vấn",
            options=st.session_state.query_list[::-1] + ["--Chọn truy vấn--"],
            key="query_selectbox"
        )

        if selected_query != "--Chọn truy vấn--":
            st.session_state.selected_query = selected_query

        if st.session_state.selected_query:
            selected_results = st.session_state.queries_and_results[st.session_state.selected_query]
            for i, result in enumerate(selected_results):
                with st.expander(f"Top {i+1}: {result['title'][0].upper() + result['title'][1:50]}... (Score: {result['bi_score']:.2f})"):
                    st.markdown(f"**Văn bản:** {result['law_id']}")
                    st.markdown(f"**Tiêu đề:** {result['title'][0].upper() + result['title'][1:]}")
                    st.markdown(f"**Nội dung:** {result['text']}")

        st.subheader("Cài đặt Retrieve")
        new_top_k = st.slider("Chọn số lượng kết quả top-k:", min_value=1, max_value=30, value=st.session_state.top_k)
        if new_top_k != st.session_state.top_k:
            st.session_state.top_k = new_top_k
            update_all_queries_results(config) 

        st.session_state.model = st.selectbox("Chọn mô hình Ollama", models)
        with st.expander("Custom Prompt", expanded=False):
            custom_prompt = st.text_area("Prompt Template", value=st.session_state.custom_prompt, height=300)
            if st.button("Save"):
                st.session_state.custom_prompt = custom_prompt
                st.success("Prompt template updated.")
                st.caption(st.session_state.custom_prompt)
                
    
    # Use the custom prompt if available, otherwise use the default one
    prompt_template = st.session_state.custom_prompt

    if len(st.session_state.messages) != 0:
        for message in st.session_state.messages:
            with st.chat_message(message["role"]):
                st.markdown(message["content"])
    else:
        with st.chat_message("assistant"):
            intro_text = "Chào bạn! Tôi là trợ lý pháp luật Việt Nam. Tôi có thể giúp bạn trả lời các câu hỏi và tìm kiếm về pháp luật Việt Nam. Nếu có câu hỏi gì xin vui lòng nhắn bên dưới!"
            message_placeholder = st.empty()
            for i in range(len(intro_text) + 1):
                message_placeholder.markdown(intro_text[:i+1] + "▌")
                time.sleep(0.005)
            message_placeholder.markdown(intro_text)
            st.session_state.messages.append({
                "role": "assistant",
                "content": intro_text
            })

    col1, col2, col3 = st.columns(3)
    with col1:
        sg_1 = st.button("Điều kiện áp dụng hợp đồng trong pháp luật Việt Nam?", use_container_width=True)

    with col2:
        sg_2 = st.button("Quy định về bảo vệ quyền lợi người tiêu dùng?", use_container_width=True)

    with col3:
        sg_3 = st.button("Quy trình khiếu nại trong pháp luật Việt Nam?", use_container_width=True)

    
    if sg_1:
        user_input = "Điều kiện áp dụng hợp đồng trong pháp luật Việt Nam?"
        process_input(prompt_template, user_input, config)
    elif sg_2:
        user_input = "Quy định về bảo vệ quyền lợi người tiêu dùng?"
        process_input(prompt_template, user_input, config)
    elif sg_3:
        user_input = "Quy trình khiếu nại trong pháp luật Việt Nam?"
        process_input(prompt_template, user_input, config)

    user_input = st.chat_input("Nhập tin nhắn của bạn")
    if user_input:
        process_input(prompt_template, user_input, config)