app.py

import streamlit as st
import os
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_community.document_loaders import TextLoader
from langchain_huggingface import HuggingFaceEmbeddings
from langchain.prompts import PromptTemplate

from langchain_core.output_parsers import StrOutputParser

from langchain_core.runnables import RunnablePassthrough
from langchain_chroma import Chroma
import Raptor

page = st.title("Chat with AskUSTH")

if "gemini_api" not in st.session_state:
    st.session_state.gemini_api = None

if "rag" not in st.session_state:
    st.session_state.rag = None
    
if "llm" not in st.session_state:
    st.session_state.llm = None

@st.cache_resource
def get_chat_google_model(api_key):
    os.environ["GOOGLE_API_KEY"] = api_key
    return ChatGoogleGenerativeAI(
        model="gemini-1.5-flash",
        temperature=0,
        max_tokens=None,
        timeout=None,
        max_retries=2,
    )

@st.cache_resource
def get_embedding_model():
    model_name = "bkai-foundation-models/vietnamese-bi-encoder"
    model_kwargs = {'device': 'cpu'}
    encode_kwargs = {'normalize_embeddings': False}
    
    model = HuggingFaceEmbeddings(
        model_name=model_name,
        model_kwargs=model_kwargs,
        encode_kwargs=encode_kwargs
    )   
    return model

if "embd" not in st.session_state:
    st.session_state.embd = get_embedding_model()

if "model" not in st.session_state:
    st.session_state.model = None

if "save_dir" not in st.session_state:
    st.session_state.save_dir = None 

if "uploaded_files" not in st.session_state:
    st.session_state.uploaded_files = set()
    
@st.dialog("Setup Gemini")
def vote():
    st.markdown(
        """
        Để sử dụng Google Gemini, bạn cần cung cấp API key. Tạo key của bạn [tại đây](https://ai.google.dev/gemini-api/docs/get-started/tutorial?lang=python&hl=vi) và dán vào bên dưới.
        """
    )
    key = st.text_input("Key:", "")
    if st.button("Save") and key != "":
        st.session_state.gemini_api = key
        st.rerun()  

if st.session_state.gemini_api is None:
    vote()

if st.session_state.gemini_api and st.session_state.model is None:
    st.session_state.model = get_chat_google_model(st.session_state.gemini_api)

if st.session_state.save_dir is None:
    save_dir = "./Documents"
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    st.session_state.save_dir = save_dir
    
def load_txt(file_path):
    loader_sv = TextLoader(file_path=file_path, encoding="utf-8")
    doc = loader_sv.load()
    return doc

with st.sidebar:
    uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"])   
    if st.session_state.gemini_api:
        if uploaded_files:
            documents = []
            uploaded_file_names = set()
            new_docs = False
            for uploaded_file in uploaded_files:
                uploaded_file_names.add(uploaded_file.name)
                if uploaded_file.name not in st.session_state.uploaded_files:
                    file_path = os.path.join(st.session_state.save_dir, uploaded_file.name)
                    with open(file_path, mode='wb') as w:
                        w.write(uploaded_file.getvalue())
                else:
                    continue

                new_docs = True
                
                doc = load_txt(file_path)
                
                documents.extend([*doc])
                
            if new_docs:  
                st.session_state.uploaded_files = uploaded_file_names
                st.session_state.rag = None
        else:
            st.session_state.uploaded_files = set()
            st.session_state.rag = None

def format_docs(docs):
    return "\n\n".join(doc.page_content for doc in docs)

@st.cache_resource
def compute_rag_chain(_model, _embd, docs_texts):
    results = Raptor.recursive_embed_cluster_summarize(_model, _embd, docs_texts, level=1, n_levels=3)
    all_texts = docs_texts.copy()
    i = 0
    for level in sorted(results.keys()):
        summaries = results[level][1]["summaries"].tolist()
        all_texts.extend(summaries)
        print(f"summary {i} -------------------------------------------------")
        print(summaries)
        i += 1
    print("all_texts ______________________________________")
    print(all_texts)
    vectorstore = Chroma.from_texts(texts=all_texts, embedding=_embd)
    retriever = vectorstore.as_retriever()
    template = """
        Bạn là một trợ lí AI hỗ trợ tuyển sinh và sinh viên. \n
        Hãy trả lời câu hỏi chính xác, tập trung vào thông tin liên quan đến câu hỏi. \n
        Nếu bạn không biết câu trả lời, hãy nói không biết, đừng cố tạo ra câu trả lời.\n
        Dưới đây là thông tin liên quan mà bạn cần sử dụng tới:\n
        {context}\n
        hãy trả lời:\n
        {question}
        """
    prompt = PromptTemplate(template=template, input_variables=["context", "question"])
    rag_chain = (
        {"context": retriever | format_docs, "question": RunnablePassthrough()}
        | prompt
        | _model
        | StrOutputParser()
    )
    return rag_chain

@st.dialog("Setup RAG")
def load_rag():
    docs_texts = [d.page_content for d in documents]
    st.session_state.rag = compute_rag_chain(st.session_state.model, st.session_state.embd, docs_texts)
    st.rerun()  

if st.session_state.uploaded_files and st.session_state.model is not None:
    if st.session_state.rag is None:
        load_rag()

if st.session_state.model is not None:
    if st.session_state.llm is None:
        mess = ChatPromptTemplate.from_messages(
            [
                (
                    "system",
                    "Bản là một trợ lí AI hỗ trợ tuyển sinh và sinh viên",
                ),
                ("human", "{input}"),
            ]
        )
        chain = mess | st.session_state.model
        st.session_state.llm = chain

if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

for message in st.session_state.chat_history:
    with st.chat_message(message["role"]):
        st.write(message["content"])

prompt = st.chat_input("Bạn muốn hỏi gì?")
if st.session_state.model is not None:
    if prompt:
        st.session_state.chat_history.append({"role": "user", "content": prompt})
            
        with st.chat_message("user"):
            st.write(prompt)
            
        with st.chat_message("assistant"):
            if st.session_state.rag is not None:
                respone = st.session_state.rag.invoke(prompt)
                st.write(respone)
            else: 
                ans = st.session_state.llm.invoke(prompt)
                respone = ans.content
                st.write(respone)
                
        st.session_state.chat_history.append({"role": "assistant", "content": respone})