import os
import dotenv
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex, StorageContext,get_response_synthesizer
from llama_index.core.response_synthesizers import ResponseMode
from llama_index.core.postprocessor import LLMRerank, SimilarityPostprocessor
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.embeddings.dashscope import DashScopeEmbedding,DashScopeTextEmbeddingModels
from llama_index.llms.dashscope import DashScope,DashScopeGenerationModels
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.vector_stores.qdrant import QdrantVectorStore
from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance
dotenv.load_dotenv()
os.environ["DASHSCOPE_API_KEY"] = os.getenv("DASHSCOPE_API_KEY")
PATH = "./qdrant_db"
EMBEDDING_DIM = 1536
COLLECTION_NAME = "full_demo"
client = QdrantClient(path=PATH)

# 指定全局的LLM
Settings.llm = DashScope(DashScopeGenerationModels.QWEN_MAX,api_key=os.getenv("DASHSCOPE_API_KEY"))
Settings.embed_model = DashScopeEmbedding(DashScopeTextEmbeddingModels.TEXT_EMBEDDING_V1,api_key=os.getenv("DASHSCOPE_API_KEY"))

documents = SimpleDirectoryReader('./data').load_data()

if client.collection_exists(collection_name=COLLECTION_NAME):
    client.delete_collection(collection_name=COLLECTION_NAME)

collection = client.create_collection(collection_name=COLLECTION_NAME, vectors_config=VectorParams(size=EMBEDDING_DIM, distance=Distance.COSINE))

vector_store = QdrantVectorStore(
    client=client,
    collection_name=COLLECTION_NAME
)

storage_context = StorageContext.from_defaults(vector_store=vector_store)

index = VectorStoreIndex(documents, storage_context=storage_context)

# 定义检索后的排序模型
reranker = LLMRerank(top_n=2)
sp = SimilarityPostprocessor(similarity_cutoff=0.6)

# 定义 RAG Fusion 检索器
fusion_retriever = QueryFusionRetriever(
    retrievers=[index.as_retriever()],
    similarity_top_k=5,
    num_queries=3,
    use_async=False
)

query_engine = RetrieverQueryEngine.from_args(
    fusion_retriever,
    node_postprocessors=[reranker],
    response_synthesizer=get_response_synthesizer(
        response_mode = ResponseMode.REFINE
    )
)


# 10. 对话引擎
chat_engine = CondenseQuestionChatEngine.from_defaults(
    query_engine=query_engine
)

while True:
    question = input("User: ")
    if question.strip() == "":
        break
    response = chat_engine.chat(question)
    print(f"AI: {response}")