File size: 1,321 Bytes
8cf4b8e
 
36cac86
8cf4b8e
 
 
 
36cac86
 
 
 
 
 
 
 
 
8cf4b8e
 
 
 
 
 
 
 
 
 
 
 
36cac86
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
import os
from langchain_community.vectorstores import FAISS
from langchain_huggingface import HuggingFaceEmbeddings

# Hàm khởi tạo retriever
def load_retriever():
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
    vectorstore_path = "vectorstore"

    # Kiểm tra nếu tồn tại index.faiss mới load
    if os.path.exists(vectorstore_path) and os.path.exists(os.path.join(vectorstore_path, "index.faiss")):
        retriever = FAISS.load_local(vectorstore_path, embeddings, allow_dangerous_deserialization=True).as_retriever(search_kwargs={"k": 5})
        print("✅ Đã load vectorstore thành công.")
    else:
        retriever = None
        print("⚠️ Vectorstore chưa tồn tại. Hãy upload tài liệu trước khi tìm kiếm.")
    return retriever

# Lần đầu load retriever
retriever = load_retriever()

# Hàm reload retriever khi thêm tài liệu
def reload_retriever():
    global retriever
    retriever = load_retriever()

# Hàm retrieve_docs để lấy tài liệu
def retrieve_docs(query):
    if retriever:
        return retriever.get_relevant_documents(query)
    else:
        print("⚠️ Chưa có vectorstore để tìm kiếm. Bạn cần tải tài liệu trước.")
        return []