Initialize app

.gitignore

@@ -0,0 +1,2 @@
+venv
+.env

Dockerfile

@@ -1,13 +1,11 @@
-FROM python:3.9-slim
+FROM python:3.12-alpine
 
 WORKDIR /app
 
-RUN apt-get update && apt-get install -y \
-    build-essential \
+RUN apk add --no-cache \
+    build-base \
     curl \
-    software-properties-common \
-    git \
-    && rm -rf /var/lib/apt/lists/*
+    git
 
 COPY requirements.txt ./
 COPY src/ ./src/

src/rag.py

@@ -0,0 +1,280 @@
+import os
+from dotenv import load_dotenv
+import langchain_google_genai as genai
+import streamlit as st
+from sentence_transformers import SentenceTransformer
+import os
+import pymongo
+from langchain_google_genai import ChatGoogleGenerativeAI
+from sentence_transformers import CrossEncoder
+from typing import List, Dict, Any, Optional
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.prompts import ChatPromptTemplate
+from langchain_core.runnables import RunnablePassthrough, RunnableMap
+import time
+
+def safe_log_info(message):
+    print(f"INFO: {message}")
+
+def safe_log_warning(message):
+    print(f"WARNING: {message}")
+
+def safe_log_error(message, exc_info=False):
+    print(f"ERROR: {message}")
+    if exc_info:
+        import traceback
+        traceback.print_exc()
+        safe_log_error("Error occurred during logging", exc_info=True)
+
+load_dotenv()
+google_api_key = os.environ.get("GOOGLE_API_KEY")
+mongo_uri = os.environ.get("MONGODB_URI")
+
+@st.cache_resource
+def load_generative_model():
+    llm = ChatGoogleGenerativeAI(
+    model = 'models/gemini-2.0-flash',
+    temperature=0.2,
+    max_tokens = None,
+    timeout = 180,
+    max_retries = 2,
+    convert_system_message_to_human= True,
+    api_key = google_api_key
+    )
+    return llm
+
+@st.cache_resource
+def load_embedding_model():
+    embedding_model = SentenceTransformer("namdp-ptit/Vidense")
+    return embedding_model
+
+@st.cache_resource
+def load_mongo_collection():
+    client = pymongo.MongoClient(mongo_uri)
+    db = client['vietnamese-llms']
+    collection = db['vietnamese-llms-data']
+    return collection
+
+@st.cache_resource
+def load_reranker():
+    reranker = CrossEncoder("namdp-ptit/ViRanker")
+    return reranker
+
+def get_embedding(text: str) -> list[float]:
+    embedding_model = load_embedding_model()
+    embedding = embedding_model.encode(text).tolist()
+    return embedding
+
+def find_similar_documents_hybrid_search(
+    query_vector: list[float],
+    search_query: str,
+    limit: int = 10,
+    candidates: int = 20,
+    vector_search_index: str = "embedding_search", 
+    atlas_search_index: str = "header_text"
+) -> list[dict]:
+    """
+    Hybrid search combining vector and text search with parallel execution.
+    """
+    all_results = []
+    collection = load_mongo_collection()
+    def perform_vector_search():
+        """Perform vector search in parallel."""
+        try:
+            vector_pipeline = [
+                {
+                    "$vectorSearch": {
+                        "index": vector_search_index,
+                        "path": "embedding",
+                        "queryVector": query_vector,
+                        "limit": limit,
+                        "numCandidates": candidates
+                    }
+                },
+                {
+                    "$project": {
+                        '_id': 1,
+                        'header' : 1,
+                        'content': 1,
+                        "vector_score": {"$meta": "vectorSearchScore"}
+                    }
+                }
+            ]
+
+            vector_results = list(collection.aggregate(vector_pipeline))
+            safe_log_info(f"Vector search returned {len(vector_results)} results")
+            for doc in vector_results:
+                doc['search_type'] = 'vector'
+                doc['combined_score'] = doc.get('vector_score', 0) * 0.6  # Weight vector score
+            return vector_results
+        except Exception as e:
+           safe_log_warning(f"Vector search failed: {e}")
+           return []
+
+    def perform_text_search():
+        """Perform text search in parallel."""
+        if not search_query or not search_query.strip():
+            return []
+
+        try:
+            text_pipeline = [
+                {
+                    "$search": {
+                        "index": atlas_search_index,
+                        "compound": {
+                            "must": [
+                                {
+                                    "text": {
+                                        "query": search_query,
+                                        "path": ["header", "content"]
+                                    }
+                                }
+                            ]
+                        }
+                    }
+                },
+                {
+                    "$project": {
+                        '_id': 1,
+                        'header': 1,
+                        'content': 1,
+                        "text_score": {"$meta": "searchScore"}
+                    }
+                }
+            ]
+
+            text_results = list(collection.aggregate(text_pipeline))
+            safe_log_info(f"Text search returned {len(text_results)} results")
+            for doc in text_results:
+                doc['search_type'] = 'text'
+                doc['combined_score'] = doc.get('text_score', 0) * 0.4  # Weight text score
+            return text_results
+        except Exception as e:
+            safe_log_warning(f"Text search failed: {e}")
+            return []
+
+    try:
+        # Run both searches in parallel
+        start_time = time.time()
+        with ThreadPoolExecutor(max_workers=2) as executor:
+            vector_future = executor.submit(perform_vector_search)
+            text_future = executor.submit(perform_text_search)
+
+            # Collect results as they complete
+            for future in as_completed([vector_future, text_future]):
+                try:
+                    results = future.result()
+                    all_results.extend(results)
+                except Exception as e:
+                    safe_log_error(f"Error in parallel search: {e}")
+
+        search_time = time.time() - start_time
+        safe_log_info(f"Parallel search completed in {search_time:.3f}s")
+
+        # 3. Merge và deduplicate results
+        seen_ids = set()
+        merged_results = []
+
+        for doc in all_results:
+            doc_id = str(doc['_id'])
+            if doc_id not in seen_ids:
+                seen_ids.add(doc_id)
+                # Clean up the document for final result
+                final_doc = {
+                    '_id': doc['_id'],
+                    'content': doc.get('content', ''),
+                    # 'uploader_username': doc.get('uploader_username', ''), # Removed
+                    'header': doc.get('header', ''),
+                    'score': doc.get('combined_score', 0)
+                }
+                merged_results.append(final_doc)
+            else:
+                # If document already exists, boost its score
+                for existing_doc in merged_results:
+                    if str(existing_doc['_id']) == doc_id:
+                        existing_doc['score'] += doc.get('combined_score', 0) * 0.5
+                        break
+
+        # Sort by combined score
+        merged_results.sort(key=lambda x: x.get('score', 0), reverse=True)
+
+        # Return top results
+        final_results = merged_results[:limit]
+        safe_log_info(f"Hybrid search final results: {len(final_results)} documents")
+
+        return final_results
+
+    except Exception as e:
+        safe_log_error(f"Error in hybrid search: {e}", exc_info=True)
+
+def rerank_documents(query: str, documents: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+    """
+    Reranks a list of documents based on their relevance to the query using a reranker model.
+
+    Args:
+        query: The original search query.
+        documents: A list of dictionaries, where each dictionary represents a document
+                   and contains a 'content' key with the document's text.
+
+    Returns:
+        A list of dictionaries representing the reranked documents, sorted by relevance score.
+    """
+    if not documents:
+        return []
+    reranker_model = load_reranker()
+    # Prepare pairs for the reranker model
+    sentence_pairs = [[query, doc.get('content', '')] for doc in documents]
+
+    # Get reranking scores
+    rerank_scores = reranker_model.predict(sentence_pairs)
+
+    # Add reranking scores to the documents
+    for i, doc in enumerate(documents):
+        doc['rerank_score'] = float(rerank_scores[i]) # Convert to float for potential serialization
+
+    # Sort documents by reranking score in descending order
+    reranked_documents = sorted(documents, key=lambda x: x.get('rerank_score', -1), reverse=True)
+
+    return reranked_documents
+
+def format_docs(docs):
+    return "\n\n".join([doc.get('header', '') + doc.get('content', '') for doc in docs if isinstance(doc, dict) and 'content' in doc and 'header' in doc])
+
+def get_answer_with_rag(query:str) -> str:
+
+    revised_template = ChatPromptTemplate.from_messages([
+    ('system', """bạn là một trợ lý AI thân thiện, được thiết kế để giúp khám phá mọi điều về Học viện Bưu chính Viễn thông (PTIT).
+        Bạn sẽ sử dụng thông tin được cung cấp để trả lời các câu hỏi của người dùng một cách chi tiết và dễ hiểu nhất.
+        Hãy nhớ rằng, bạn chỉ có thể trả lời dựa trên thông tin bạn cung cấp. Nếu câu hỏi nằm ngoài phạm vi thông tin đó, bạn sẽ cho người dùng biết."""),
+    ('human', "Thông tin tham khảo:\n```\n{context}\n```\n\nCâu hỏi của tôi:\n{question}")
+    ])
+    llm = load_generative_model()
+    query_embedding = get_embedding(query)
+
+    context_docs = find_similar_documents_hybrid_search(
+        query_vector=query_embedding,
+        search_query=query,
+        limit=10,
+        candidates=20, 
+        vector_search_index="embedding_search",
+        atlas_search_index="header_text"
+    )
+
+    reranked_docs = rerank_documents(query, context_docs)
+    top_n_docs = reranked_docs[:10]
+    context = format_docs(top_n_docs)
+
+    chain = (
+        RunnableMap({
+                "context": RunnablePassthrough(),
+                "question": RunnablePassthrough()
+            })
+        | revised_template
+        | llm
+        | StrOutputParser()
+    )
+    response = chain.invoke({
+        "context": context,
+        "question": query})
+    return response

src/streamlit_app.py

@@ -1,40 +1,64 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from rag import get_answer_with_rag 
+from rag import get_answer_with_rag, load_generative_model, load_reranker, load_embedding_model, load_mongo_collection
 
-"""
-# Welcome to Streamlit!
-
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
-
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
-
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+load_generative_model()
+load_reranker()
+load_embedding_model()
+load_mongo_collection()
+st.set_page_config(
+    page_title="PTIT RAG Chatbot",
+    page_icon="🤖",
+    layout="wide"
+)
+
+# --- GIAO DIỆN CHÍNH ---
+st.title("🤖 PTIT RAG Chatbot")
+st.caption("Trợ lý ảo thông minh về Học viện Bưu chính Viễn thông")
+
+# Khởi tạo session state để lưu trữ lịch sử trò chuyện
+if "messages" not in st.session_state:
+    st.session_state.messages = [
+        {"role": "assistant", "content": "Xin chào! Tôi có thể giúp gì cho bạn về các thông tin tại PTIT?"}
+    ]
+
+# --- SIDEBAR ---
+with st.sidebar:
+    st.header("Tùy chọn")
+    if st.button("🗑️ Xóa cuộc trò chuyện", use_container_width=True):
+        st.session_state.messages = [
+            {"role": "assistant", "content": "Cuộc trò chuyện đã được xóa. Hãy bắt đầu lại nhé!"}
+        ]
+        st.rerun()
+
+    st.markdown("---")
+    st.markdown("### Về ứng dụng")
+    st.info("Ứng dụng này sử dụng RAG để trả lời câu hỏi dựa trên tài liệu về PTIT.")
+
+# Hiển thị lịch sử trò chuyện
+for message in st.session_state.messages:
+    avatar = "🧑‍💻" if message["role"] == "user" else "🤖"
+    with st.chat_message(message["role"], avatar=avatar):
+        st.write(message["content"])
+
+def submit_question(question: str):
+    st.session_state.messages.append({"role": "user", "content": question})
+    st.rerun()
+
+# Khu vực nhập liệu của người dùng
+if prompt := st.chat_input("Nhập câu hỏi của bạn..."):
+    # Thêm tin nhắn của người dùng vào session state và hiển thị ngay
+    st.session_state.messages.append({"role": "user", "content": prompt})
+    with st.chat_message("user", avatar="🧑‍💻"):
+        st.write(prompt)
+
+    with st.chat_message("assistant", avatar="🤖"):
+        with st.spinner("🤖 Tôi đang suy nghĩ, bạn chờ chút nhé..."):
+            try:
+                response = get_answer_with_rag(prompt)
+                st.markdown(response)
+                st.session_state.messages.append({"role": "assistant", "content": response})
+            except Exception as e:
+                error_message = "Rất tiếc, đã có lỗi xảy ra. Vui lòng thử lại sau!"
+                st.error(error_message)
+                st.session_state.messages.append({"role": "assistant", "content": error_message})