import os
import time
from typing import List
from qdrant_client import QdrantClient, models
from langchain_core.documents import Document
from semantic_cache.main import SemanticCache
from langchain_community.embeddings.huggingface import HuggingFaceEmbeddings
from Router.router import Evaluator
from langchain_openai import ChatOpenAI
# from utils.pipelines.main import get_last_user_message, add_or_update_system_message, pop_system_message
from blueprints.rag_utils import format_docs, translate
from blueprints.prompts import QUERY_PROMPT, evaluator_intent, basic_template, chitchat_prompt, safe_prompt, cache_prompt
from SafetyChecker import SafetyChecker
from langchain.retrievers import EnsembleRetriever
from BM25 import BM25SRetriever
# from database_Routing import DB_Router
from langchain.retrievers.multi_query import MultiQueryRetriever
# import cohere
from langchain.retrievers.document_compressors import LLMChainExtractor
from langchain_core.output_parsers import StrOutputParser
from langchain_core.output_parsers import BaseOutputParser
# from langchain_cohere import CohereRerank
from langchain.retrievers.contextual_compression import ContextualCompressionRetriever
from langchain_groq import ChatGroq
from langchain_core.runnables import RunnablePassthrough
import time
from qdrant_client import QdrantClient
from langchain_community.vectorstores import Qdrant
from langchain.retrievers.document_compressors import LLMChainFilter
from langchain.retrievers.document_compressors import EmbeddingsFilter
from langchain.retrievers.document_compressors import LLMListwiseRerank
from dotenv import load_dotenv
from langchain_openai import OpenAIEmbeddings
load_dotenv()
os.environ["OPENAI_API_KEY"] = os.getenv('OPENAI_KEY')
os.environ["COHERE_API_KEY"]
# HF_EMBEDDING = HuggingFaceEmbeddings(model_name="dangvantuan/vietnamese-embedding")
HF_EMBEDDING = OpenAIEmbeddings(model='text-embedding-3-small', api_key = os.getenv('OPENAI_KEY'))
class LineListOutputParser(BaseOutputParser[List[str]]):
"""Output parser for a list of lines."""
def parse(self, text: str) -> List[str]:
lines = text.strip().split("\n")
return list(filter(None, lines)) # Remove empty lines
def add_or_update_system_message(content: str, messages: List[dict]):
"""
Adds a new system message at the beginning of the messages list
:param msg: The message to be added or appended.
:param messages: The list of message dictionaries.
:return: The updated list of message dictionaries.
"""
if messages and messages[0].get("role") == "system":
messages[0]["content"] += f"{content}\n"
else:
# Insert at the beginning
messages.insert(0, {"role": "system", "content": content})
return messages
def split_context( context):
split_index = context.find("User question")
system_prompt = context[:split_index].strip()
user_question = context[split_index:].strip()
user_split_index = user_question.find("<context>")
f_system_prompt = str(system_prompt) +"\n" + str(user_question[user_split_index:])
return f_system_prompt
def extract_metadata(docs, headers=('Header_1', 'Header_2', 'Header_3')):
meta_data_docs = []
for doc in docs:
meta_data_doc = [doc.metadata[header] for header in headers if doc.metadata.get(header)]
meta_data_docs.append(meta_data_doc)
return meta_data_docs
def search_with_filter(query, vector_store, k, headers):
conditions = []
# Xử lý điều kiện theo số lượng headers
if len(headers) == 1:
conditions.append(
models.FieldCondition(
key="metadata.Header_1",
match=models.MatchValue(
value=headers[0]
),
)
)
elif len(headers) == 2:
conditions.append(
models.FieldCondition(
key="metadata.Header_1",
match=models.MatchValue(
value=headers[0]
),
)
)
conditions.append(
models.FieldCondition(
key="metadata.Header_2",
match=models.MatchValue(
value=headers[1]
),
)
)
elif len(headers) == 3:
conditions.append(
models.FieldCondition(
key="metadata.Header_1",
match=models.MatchValue(
value=headers[0]
),
)
)
conditions.append(
models.FieldCondition(
key="metadata.Header_2",
match=models.MatchValue(
value=headers[1]
),
)
)
conditions.append(
models.FieldCondition(
key="metadata.Header_3",
match=models.MatchValue(
value=headers[2]
),
)
)
# Thực hiện truy vấn với các điều kiện
single_result = vector_store.similarity_search(
query=query,
k=k,
filter=models.Filter(
must=conditions
),
)
return single_result
def get_relevant_documents(documents: List[Document], limit: int) -> List[Document]:
result = []
seen = set()
for doc in documents:
if doc.page_content not in seen:
result.append(doc)
seen.add(doc.page_content)
if len(result) == limit:
break
return result
if __name__ == "__main__":
client = QdrantClient(
url="http://localhost:6333"
)
stsv = Qdrant(client, collection_name="sotaysinhvien_filter", embeddings= HF_EMBEDDING)
stsv_db = stsv.as_retriever(search_kwargs={'k': 10})
gthv = Qdrant(client, collection_name="gioithieuhocvien_filter", embeddings= HF_EMBEDDING)
gthv_db = gthv.as_retriever(search_kwargs={'k': 10})
ttts = Qdrant(client, collection_name="thongtintuyensinh_filter", embeddings= HF_EMBEDDING)
ttts_db = ttts.as_retriever(search_kwargs={'k': 10})
import pickle
with open('data/sotaysinhvien_filter.pkl', 'rb') as f:
sotaysinhvien = pickle.load(f)
with open('data/thongtintuyensinh_filter.pkl', 'rb') as f:
thongtintuyensinh = pickle.load(f)
with open('data/gioithieuhocvien_filter.pkl', 'rb') as f:
gioithieuhocvien = pickle.load(f)
retriever_bm25_tuyensinh = BM25SRetriever.from_documents(thongtintuyensinh, k= 10, save_directory = "data/bm25s/ttts")
retriever_bm25_sotay = BM25SRetriever.from_documents(sotaysinhvien, k= 10, save_directory = "data/bm25s/stsv")
retriever_bm25_hocvien = BM25SRetriever.from_documents(gioithieuhocvien, k= 10, save_directory = "data/bm25s/gthv" )
# reranker = CohereRerank(model = "rerank-multilingual-v3.0", top_n = 5)
llm = ChatGroq(model_name="llama3-70b-8192", temperature=0.1,api_key= os.getenv('llm_api_3'))
llm2 = ChatGroq(model_name="llama-3.1-70b-versatile", temperature=1,api_key= os.getenv('llm_api_8'))
output_parser = LineListOutputParser()
llm_chain = QUERY_PROMPT | llm | output_parser
# messages = [
# {"role": "system", "content": "Dựa vào thông tin sau, trả lời câu hỏi bằng tiếng việt"}
# ]
# ###########################
cache = SemanticCache()
another_chain = ( chitchat_prompt | llm2 | StrOutputParser())
safe_chain = ( safe_prompt | llm2 | StrOutputParser())
cache_chain = ( cache_prompt | llm2 | StrOutputParser())
# def duy_phen():
while 1:
body = {}
user_message = input("Nhập câu hỏi nào!: ")
checker = SafetyChecker()
safety_result = checker.check_safety(translate(user_message))
print("Safety check :" ,safety_result)
if safety_result != 'safe' :
print("UNSAFE")
response = safe_chain.invoke({'meaning': f'{safety_result}'})
print(response)
exit()
evaluator = Evaluator(llm="llama3-70b", prompt=evaluator_intent)
output = evaluator.classify_text(user_message)
print(output.result)
retriever = None # or assign a specific default retriever if applicable
db = None # initialize db as well if it is used later in the code
# print(output.result)
source = None
cache_result =cache.checker(user_message)
if cache_result is not None:
print("###Cache hit!###")
response = cache_chain.invoke({"question": f'{user_message}', "content": f"{cache_result}"})
print(response)
if output and output.result == 'OUT_OF_SCOPE' :
print('OUT OF SCOPE')
# print(body)
response = another_chain.invoke({"question": f"{user_message}"})
print(response)
elif output and output.result == 'ASK_QUYDINH' :
print('SO TAY SINH VIEN DB')
retriever = stsv_db
retriever_bm25 = retriever_bm25_sotay
source = stsv
# db = sotaysinhvien
elif output and output.result == 'ASK_HOCVIEN' :
print('GIOI THIEU HOC VIEN DB')
retriever = gthv_db
retriever_bm25 = retriever_bm25_hocvien
source = gthv
# db = gioithieuhocvien
elif output and output.result == 'ASK_TUYENSINH' :
print('THONG TIN TUYEN SINH DB')
retriever = ttts_db
retriever_bm25 = retriever_bm25_tuyensinh
source = ttts
# db = thongtintuyensinh
if retriever is not None:
# retriever_multi = MultiQueryRetriever(
# retriever=retriever, llm_chain=llm_chain, parser_key="lines"
# )
start_time = time.time()
ensemble_retriever = EnsembleRetriever(
retrievers=[retriever_bm25, retriever], weights=[0.5, 0.5])
# compressor = LLMChainExtractor.from_llm(llm)
# _filter = LLMChainFilter.from_llm(llm)
# embeddings_filter = EmbeddingsFilter(embeddings=HF_EMBEDDING, similarity_threshold=0.5)
# compression = ContextualCompressionRetriever(
# base_compressor=_filter2, base_retriever=ensemble_retriever
# )
reranker = LLMListwiseRerank.from_llm(
llm=llm, top_n=5
)
tailieu = ensemble_retriever.invoke(f"{user_message}")
docs = reranker.compress_documents(tailieu, user_message)
end_time = time.time()
#################### Filter lại ở đây -> add more documents liên quan hơn #########################
# docs = compression.invoke(f"{user_message}")
# print(docs)
meta_data_docs = extract_metadata(docs)
full_result = []
for meta_data_doc in meta_data_docs:
result = search_with_filter(user_message, source, 10, meta_data_doc)
for i in result:
full_result.append(i)
print("Context liên quan" + '\n')
print(full_result)
# rag_chain = (
# {"context": compression | format_docs, "question": RunnablePassthrough()}
# | basic_template | llm2 | StrOutputParser()
# )
result_final = get_relevant_documents(full_result, 10)
context = format_docs(result_final)
best_chain = ( basic_template | llm2 | StrOutputParser())
best_result = best_chain.invoke({"question": f'{user_message}', "context": f"{context}"})
print(f'Câu trả lời tối ưu nhất: {best_result}')
print(f'TIME USING : {end_time -start_time}')
else:
print('Retriever is not defined. Check output results and ensure retriever is assigned correctly.')
# duy_phen()