app.py

# import os
# import subprocess
# import gradio as gr
# import json
# from tqdm import tqdm
# from langchain_community.vectorstores import FAISS
# from langchain_google_genai import GoogleGenerativeAIEmbeddings
# import google.generativeai as genai
# # from playwright._impl._driver import get_driver_dir

# from helpers import (
#     list_docx_files, get_splits, get_json_splits_only, prompt_order, log_message
# )

# from file_loader import get_vectorstore
# # import asyncio

# key = os.getenv("GOOGLE_API_KEY")
# # Cấu hình API key cho Google GenAI
# genai.configure(api_key=key)

# vectorstore = get_vectorstore()
# last_vector_docs = None  # Lưu kết quả docs từ vectorstore.invoke trong lần gọi get_answer gần nhất

# def augment_prompt(query: str, k: int = 10):
#     global last_vector_docs
#     retriever = vectorstore.as_retriever(search_kwargs={"k": k})
#     results = retriever.invoke(query)
#     # Lưu kết quả để dùng cho log và lọc sau này
#     last_vector_docs = results  
#     if results:
#         source_knowledge = "\n\n".join([doc.page_content for doc in results])
#         return f"""Dữ liệu dưới đây liên quan đến Trường Công Nghệ (NCT) thuộc Đại học Kinh tế Quốc dân (NEU), dựa vào đó trả lời câu hỏi.
# Dữ liệu:
# {source_knowledge}
# """
#     else:
#         return "Không có thông tin liên quan.\n."

# def get_answer(query, queries_list=None):
#     if queries_list is None:
#         queries_list = []
    
#     messages = [
#         {"role": "user", "parts": [{"text": "IMPORTANT: You are a super helpful, polite, Vietnamese-speaking assistant to give information of an university. If you cannot see the answer in contexts, try to search it up online by yourself OR tell user to make a more detailed question."}]},
#         {"role": "user", "parts": [{"text": augment_prompt(query=query, k=100)}]}
#     ]
    
#     queries_list.append(query)
#     queries = {"role": "user", "parts": [{"text": prompt_order(queries_list)}]}
#     messages_with_queries = messages.copy()
#     messages_with_queries.append(queries)
    
#     # Cấu hình API key và khởi tạo model Gemini
#     genai.configure(api_key=key)
#     model = genai.GenerativeModel("gemini-2.0-flash")
    
#     response = model.generate_content(contents=messages_with_queries, stream=True)
#     response_text = ""
#     for chunk in response:
#         response_text += chunk.text
#         yield response_text
#     messages.append({"role": "model", "parts": [{"text": response_text}]})
#     log_message(messages)

# def filter_vector_docs(keyword: str):
#     global last_vector_docs
#     if last_vector_docs is None:
#         return "Chưa có dữ liệu vectorstore được gọi từ get_answer."
#     else:
#         if not keyword.strip():
#             # Nếu không nhập gì, trả về tất cả
#             filtered = [doc.page_content for doc in last_vector_docs]
#         else:
#             # Lọc các chunk chứa từ khoá (không phân biệt chữ hoa thường)
#             filtered = [doc.page_content for doc in last_vector_docs if keyword.lower() in doc.page_content.lower()]
#         if not filtered:
#             return f"Không có kết quả chứa từ khoá '{keyword}'."
#         return "\n\n".join(filtered)

# institutions = ['Tất cả', 'Trường Công Nghệ']
# categories = ['Tất cả', 'Đề án', 'Chương trình đào tạo']

# print("Launching on space... This may take some time...")

# with gr.Blocks() as demo:
#     with gr.Row():
#         # Dropdown category nếu cần
#         category1 = gr.Dropdown(choices=institutions, label="Trường", value=None)
#         category2 = gr.Dropdown(choices=categories, label="Bạn quan tâm tới", value=None)
    
#     # Chat Interface sử dụng ô nhập chung
#     shared_query = gr.Textbox(placeholder="Đặt câu hỏi tại đây", container=False, autoscroll=True, scale=7)
#     chat_interface = gr.ChatInterface(get_answer, textbox=shared_query, type="messages")
    
#     # Phần lọc các chunk: ô prompt nhập từ khoá và nút "Tìm trích xuất" nằm cùng hàng,
#     # kết quả sẽ hiển thị ở ô bên dưới. Nếu để trống, hiển thị toàn bộ.
#     with gr.Row():
#         filter_prompt = gr.Textbox(label="Nhập từ khoá", placeholder="Nhập từ khoá", interactive=True)
#         filter_button = gr.Button("Đọc trích xuất")
#     filter_output = gr.Textbox(label="Content", interactive=False)
#     filter_button.click(fn=filter_vector_docs, inputs=filter_prompt, outputs=filter_output)

# if __name__ == "__main__":
#     demo.launch()

import os
import subprocess
import gradio as gr
import json
from tqdm import tqdm
from langchain_community.vectorstores import FAISS
from langchain_google_genai import GoogleGenerativeAIEmbeddings
import google.generativeai as genai
# from playwright._impl._driver import get_driver_dir

from helpers import (
    list_docx_files, get_splits, get_json_splits_only, prompt_order, log_message, extract_metadata, update_metadata
)

from file_loader import get_vectorstore
# import asyncio

# key = "AIzaSyDJ4vIKuIBIPNHATLxnoHlagXWbsAz-vRs"
key = os.getenv("GOOGLE_API_KEY")
# Cấu hình API key cho Google GenAI
genai.configure(api_key=key)

# vectorstore = get_vectorstore()
vectorstore = get_vectorstore()
last_vector_docs = None  # Lưu kết quả docs từ vectorstore.invoke trong lần gọi get_answer gần nhất
# see_metadata = None
# see_metadata2 = None
# see_metadata_fin = None
queries_list = []

metadata = {}

def augment_prompt(query: str, queries_list=None, k: int = 10):
    #define metadata, make the button to clear metadata
    global metadata
    messages = [
        {"role": "user", "parts": [{"text": """
'Tai lieu ve': ['Đề án', 'Chương trình đào tạo', 'Đề cương']
'Nganh': ['Trí tuệ nhân tạo',
                  'Toán kinh tế (TOKT)',
                  'Thống kê kinh tế',
                  'Phân tích dữ liệu trong Kinh tế (DSEB)',
                  'Kỹ thuật phần mềm',
                  'Khoa học máy tính',
                  'Khoa học dữ liệu',
                  'Hệ thống thông tin quản lý',
                  'Hệ thống thông tin',
                  'Định phí bảo hiểm và Quản trị rủi ro (Actuary)',
                  'Công nghệ thông tin',
                  'An toàn thông tin']
'Khoa': ['Công nghệ thông tin (FIT)',
        'Toán Kinh tế (MFE)',
        'Khoa học dữ liệu và Trí tuệ nhân tạo (FDA)',
        'Thống kê',
        'Hệ thống thông tin quản lý (MIS)']

        """}]},
        {"role": "user", "parts": [{"text": f'''Câu hỏi là: "{query}"'''}]},
        {"role": "user", "parts": [{"text": """Bạn là một mô hình NER, nhiệm vụ của bạn là đọc và hiểu câu hỏi, sau đó đưa ra metadata dưới dạng dict, metadata có dạng: {'Tai lieu ve': '<write here>', 'Nganh': <write here>, 'Khoa': <write here>}, nếu không có theo 3 lists bên trên thì trả về dict rỗng {}"""}]},
    ]
    genai.configure(api_key=key)
    model = genai.GenerativeModel("gemini-2.0-flash")

    response = model.generate_content(contents=messages)
    response_text = "\n\n".join([chunk.text for chunk in response])
    # global see_metadata
    metadata_child = extract_metadata(response_text)
    # see_metadata = str(response_text)

    messages2 = [
        {"role": "user", "parts": [{"text": f'''Metadata cũ: "{str(metadata_child)}"'''}]},
        {"role": "user", "parts": [{"text": """
            Bạn là một mô hình NER, bạn sẽ update metadata cũ thỏa mãn các tiêu chí sau:
            1. Tôi đã đưa metadata ở phía trên, nếu bạn thấy metadata cũ không có hoặc 1 dict rỗng {} thì output của bạn là dict rỗng {}.
            2. Nếu không có biến 'Nganh' thì output của bạn là dict rỗng {}.
            3. Nếu thỏa mãn các tiêu chí trên thì output ra giống hệt như cũ, không thay đổi gì cả, output dưới dạng dict {}.
        ]}"""}]},
    ]
    # genai.configure(api_key=key)
    model = genai.GenerativeModel("gemini-2.0-flash")
    response = model.generate_content(contents=messages2)
    response_text = "\n\n".join([chunk.text for chunk in response])
    # global see_metadata2
    metadata_child = extract_metadata(response_text)

    # see_metadata2 = response_text
    #retrieve
    if metadata_child == None:
        metadata_child = {}
    # metadata = {}

    #hàm update metadata
    # global see_metadata_fin
    metadata = update_metadata(metadata, metadata_child)
    # see_metadata_fin = metadata
    global last_vector_docs
    if metadata == {}:
        # retriever = vectorstore.as_retriever(search_kwargs={"k": 25, "fetch_k": 50, "filter": {'Tai lieu ve': 'Đề án'}})
        retriever = vectorstore.as_retriever(search_kwargs={"k": 30})
    else:
        retriever = vectorstore.as_retriever(search_kwargs={"k": k, "fetch_k": 20, "filter": metadata})

    results = retriever.invoke(query)
    # Lưu kết quả để dùng cho log và lọc sau này
    last_vector_docs = results
    if results:
        source_knowledge = "\n\n".join([doc.page_content for doc in results])
        return f"""Dữ liệu dưới đây liên quan đến Trường Công Nghệ (NCT) thuộc Đại học Kinh tế Quốc dân (NEU), dựa vào đó trả lời câu hỏi.
Dữ liệu:
{source_knowledge}
"""
    else:
        return "Không có thông tin liên quan.\n."

def get_answer(query, queries_listt = None):
    # if history is None:
    #     history = [] # tạm thời không lưu lịch sử
    global queries_list
    queries_list.append(query)

    queries_listt = queries_list.copy()

    messages = [
        {"role": "user", "parts": [{"text": "IMPORTANT: You are a super helpful, polite, Vietnamese-speaking assistant to give information of an university. If you cannot see the answer in contexts, tell user to make a more detailed question."}]},
        {"role": "user", "parts": [{"text": augment_prompt(query=query, queries_list=queries_listt)}]}
    ]

    queries = [
        {"role": "user", "parts": [{"text": prompt_order(queries_list)}]},
        {"role": "user", "parts": [{"text": "Lưu ý: Nếu câu hỏi về môn học, hãy hỏi người dùng về ngành dạy bộ môn đó là ngành nào! Nếu câu hỏi người dùng không liên quan tới nội dung được cung cấp, bảo họ bấm vào nút 'Hỏi cái khác'"}]}
    ]
    messages_with_queries = messages.copy()
    messages_with_queries.extend(queries)

    # Cấu hình API key và khởi tạo model Gemini
    genai.configure(api_key=key)
    model = genai.GenerativeModel("gemini-2.0-flash")

    response = model.generate_content(contents=messages_with_queries, stream=False)
    response_text = ""
    response_text = "".join(chunk.text for chunk in response)
    messages.append({"role": "model", "parts": [{"text": response_text}]})
    # log_message(messages)
    return response_text

def reset_metadata():  #reset metadata button
    global metadata
    metadata = {}
    global queries_list
    queries_list = {}
    return "{}"

def filter_vector_docs(keyword: str): #check retrived content button
    global last_vector_docs
    if last_vector_docs is None:
        return "Chưa có dữ liệu vectorstore được gọi từ get_answer."
    else:
        if not keyword.strip():
            # Nếu không nhập gì, trả về tất cả
            filtered = [doc.page_content for doc in last_vector_docs]
        else:
            # Lọc các chunk chứa từ khoá (không phân biệt chữ hoa thường)
            filtered = [doc.page_content for doc in last_vector_docs if keyword.lower() in doc.page_content.lower()]
        if not filtered:
            return f"Không có kết quả chứa từ khoá '{keyword}'."
        return "\n\n".join(filtered)

def get_metadata():
    # Hàm này dùng để cập nhật ô metadata sau mỗi prompt
    global metadata
    return str(metadata) if metadata is not None else "{}"

institutions = ['Tất cả', 'Trường Công Nghệ']
categories = ['Tất cả', 'Đề án', 'Chương trình đào tạo']

print("Launching on space... This may take some time...")

with gr.Blocks() as demo:
    # with gr.Row():
    #     # Dropdown các trường và danh mục
    #     category1 = gr.Dropdown(choices=institutions, label="Trường", value=None)
    #     category2 = gr.Dropdown(choices=categories, label="Bạn quan tâm tới", value=None)

        # Phần Metadata: một hàng chứa hộp hiển thị metadata và cột chứa 2 nút (Show Metadata - Reset Metadata)
    with gr.Row():
        metadata_box = gr.Textbox(label="Metadata", value="{}", interactive=False)
        with gr.Column():
            show_meta_button = gr.Button("Show Metadata")
            reset_meta_button = gr.Button("Hỏi cái khác")
    
    # Kết nối các nút với các hàm tương ứng
    show_meta_button.click(fn=get_metadata, inputs=[], outputs=metadata_box)
    reset_meta_button.click(fn=reset_metadata, inputs=[], outputs=metadata_box)
    
    # Chat Interface sử dụng ô nhập chung và callback get_answer
    shared_query = gr.Textbox(placeholder="Đặt câu hỏi tại đây", container=False, autoscroll=True, scale=7)
    chat_interface = gr.ChatInterface(get_answer, textbox=shared_query, type="messages")
    
    # Phần lọc các chunk: ô nhập từ khoá và nút "Đọc trích xuất" nằm cùng hàng, kết quả hiển thị bên dưới
    with gr.Row():
        filter_prompt = gr.Textbox(label="Nhập từ khoá", placeholder="Nhập từ khoá", interactive=True)
        filter_button = gr.Button("Đọc trích xuất")
    filter_output = gr.Textbox(label="Content", interactive=False)
    filter_button.click(fn=filter_vector_docs, inputs=filter_prompt, outputs=filter_output)

if __name__ == "__main__":
    demo.launch()