import streamlit as st import os from langchain_google_genai import ChatGoogleGenerativeAI from langchain_core.prompts import ChatPromptTemplate 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_qdrant import QdrantVectorStore import Raptor from io import StringIO from qdrant_client import QdrantClient from qdrant_client.models import Distance, VectorParams 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() @st.cache_resource def load_chromadb(collection_name): client = QdrantClient( url="https://da9fadd2-dc5a-4481-ac79-4e2677a2354b.europe-west3-0.gcp.cloud.qdrant.io", api_key="X_-IVToBM07Mot4Mmzg5xNjYzc1DlIgl0VQDUNmGhI_Z-WA5FJ2ETA" ) client.recreate_collection( collection_name=collection_name, vectors_config=VectorParams(size=768, distance=Distance.COSINE) ) db = QdrantVectorStore( client=client, collection_name=collection_name, embedding=st.session_state.embd, ) return db @st.cache_resource def update_chromadb(collection_name): client = QdrantClient( url="https://da9fadd2-dc5a-4481-ac79-4e2677a2354b.europe-west3-0.gcp.cloud.qdrant.io", api_key="X_-IVToBM07Mot4Mmzg5xNjYzc1DlIgl0VQDUNmGhI_Z-WA5FJ2ETA" ) try: client.delete_collection(collection_name=collection_name) except Exception as e: print(f"Warning: {e}") client.recreate_collection( collection_name=collection_name, vectors_config=VectorParams(size=768, distance=Distance.COSINE) ) db = QdrantVectorStore( client=client, collection_name=collection_name, embedding=st.session_state.embd, ) return db if "vector_store" not in st.session_state: st.session_state.vector_store = load_chromadb("data") 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) def format_docs(docs): return "\n\n".join(doc.page_content for doc in docs) @st.cache_resource def rag_chain(_model, _vectorstore): 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 = ( {"context": retriever | format_docs, "question": RunnablePassthrough()} | prompt | _model | StrOutputParser() ) return rag if st.session_state.model is not None and st.session_state.vector_store is not None: st.session_state.rag = rag_chain(st.session_state.model, st.session_state.vector_store) if "new_docs" not in st.session_state: st.session_state.new_docs = False with st.sidebar: uploaded_files = st.file_uploader("Chọn file txt", accept_multiple_files=True, type=["txt"]) if st.session_state.model: documents = [] uploaded_file_names = set() if uploaded_files: for uploaded_file in uploaded_files: uploaded_file_names.add(uploaded_file.name) if uploaded_file_names != st.session_state.uploaded_files and not st.session_state.new_docs: st.session_state.uploaded_files = uploaded_file_names st.session_state.new_docs = True if uploaded_files: for uploaded_file in uploaded_files: stringio=StringIO(uploaded_file.getvalue().decode('utf-8')) read_data=str(stringio.read()) documents.append(read_data) def update_rag_chain(_model, _embd, _vectorstore, docs_texts): results = Raptor.recursive_embed_cluster_summarize(_model, _embd, docs_texts, level=1, n_levels=3) all_texts = docs_texts.copy() for level in sorted(results.keys()): summaries = results[level][1]["summaries"].tolist() all_texts.extend(summaries) _vectorstore.add_texts(texts=all_texts) rag = rag_chain(_model, _vectorstore) return rag def reset_rag_chain(_model, _vectorstore): rag = rag_chain(_model, _vectorstore) return rag if "query_router" not in st.session_state: st.session_state.query_router = None @st.cache_resource def query_router(_model): mess = ChatPromptTemplate.from_messages( [ ( "system", """Bạn là một chatbot hỗ trợ giải đáp về đại học, nhiệm vụ của bạn là phân loại câu hỏi. Nếu câu hỏi về đại học thì trả về 'university', nếu không liên quan tới tuyển sinh và sinh viên thì trả về 'other'. Bắt buộc Kết quả chỉ trả về với một trong hai lựa chọn trên. Không được trả lời thêm bất kỳ thông tin nào khác.""", ), ("human", "{input}"), ] ) chain = mess | _model return chain if st.session_state.model is not None: st.session_state.query_router = query_router(st.session_state.model) @st.dialog("Update DB") def update_vectorstore(_model, _embd, _vectorstore, docs): docs_texts = [d for d in docs] st.session_state.rag = update_rag_chain(_model, _embd, _vectorstore, docs_texts) st.rerun() @st.dialog("Reset DB") def reset_vectorstore(_model, _vectorstore): st.session_state.rag = reset_rag_chain(_model, _vectorstore) st.rerun() if st.session_state.new_docs: st.session_state.new_docs = False st.session_state.vector_store = update_chromadb("data") if st.session_state.uploaded_files: update_vectorstore(st.session_state.model, st.session_state.embd, st.session_state.vector_store, documents) else: reset_vectorstore(st.session_state.model, st.session_state.vector_store) 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"): router = st.session_state.query_router.invoke(prompt) switch = router.content if "university" in switch: respone = st.session_state.rag.invoke(prompt) f_response = f"RAG: {respone}" st.write(f_response) else: respone = st.session_state.llm.invoke(prompt) f_response = f"other: {respone.content}" st.write(f_response) st.session_state.chat_history.append({"role": "assistant", "content": f_response})