from pymilvus import (
    MilvusClient,
    DataType,
    AnnSearchRequest,
    RRFRanker
)
from langchain_community.embeddings import DashScopeEmbeddings
import dashscope

# 替换为您的 DashScope API-KEY
dashscope_api_key = "<YOUR_DASHSCOPE_API_KEY>"

# 初始化 Embedding 模型
embeddings = DashScopeEmbeddings(
    model="text-embedding-v2",  # 使用text-embedding-v2模型。
    dashscope_api_key=dashscope_api_key
)

# 替换为您的 Milvus 实例信息
milvus_url = "<YOUR_MILVUS_URL>"
user_name = "<YOUR_USERNAME>"
password = "<YOUR_PASSWORD>"

# 创建 Milvus Client
client = MilvusClient(
    uri="http://" + milvus_url + ":19530",
    token=user_name + ":" + password
)

# 定义 Collection 名称
collection_name = "milvus_bm25_hybrid"

# 定义文本内容
text_contents = [
    "Milvus is a highly scalable vector database designed for AI applications. It offers high performance and hybrid search capabilities.",
    "Milvus supports multiple indexing algorithms, including IVF, HNSW, and DiskANN, to ensure efficient similarity search.",
    "Milvus is compatible with various ecosystems, including LangChain, LlamaIndex, and ChatGLM.",
    "Milvus provides a unified platform for data management, including data ingestion, processing, and querying.",
    "Milvus offers enterprise-grade security features, including authentication, authorization, and encryption.",
    "Milvus is optimized for cloud environments and supports Kubernetes deployment for easy scaling.",
    "Milvus has a comprehensive set of APIs and SDKs for various programming languages, including Python, Java, and Go.",
    "Milvus supports both sparse and dense vector search, enabling hybrid search capabilities.",
    "Milvus provides a web-based management interface for easy monitoring and administration.",
    "Milvus is an open-source project with a vibrant community and extensive documentation."
]

# 生成向量
vectors = embeddings.embed_documents(text_contents)
dense_dim = len(vectors[0])

# 定义 Schema
schema = MilvusClient.create_schema(
    auto_id=False,
    enable_dynamic_field=True,
)

schema.add_field(
    field_name="id",
    datatype=DataType.INT64,
    is_primary=True,
)

schema.add_field(
    field_name="dense",
    datatype=DataType.FLOAT_VECTOR,
    dim=dense_dim
)

# 定义索引参数
index_params = client.prepare_index_params()

index_params.add_index(
    field_name="dense",
    index_type="HNSW",
    metric_type="IP",
    params={"M": 8, "efConstruction": 96}
)

# 定义分词器参数
analyzer_params = {
    "type": "english"  # 指定分词器类型为英文
}

# 添加文本字段到 Schema，并启用分词器
schema.add_field(
    field_name="text",                      # 字段名称
    datatype=DataType.VARCHAR,              # 数据类型：字符串（VARCHAR）
    max_length=65535,                       # 最大长度：65535 字符
    enable_analyzer=True,                   # 启用分词器
    analyzer_params=analyzer_params         # 分词器参数
)

index_params.add_index(
    field_name="text",
    index_name="text_index",
    index_type="SPARSE_INVERTED_INDEX",
    metric_type="BM25",
    params={"drop_ratio_build": 0.2}
)

# 创建 collection
client.create_collection(
    collection_name=collection_name,
    schema=schema,
    index_params=index_params
)

data = [
    {"id": idx, "dense": vectors[idx], "text": doc}
    for idx, doc in enumerate(text_contents)
]

# 插入数据
res = client.insert(
    collection_name=collection_name,
    data=data
)

print("生成 " + str(len(vectors)) + " 个向量，维度：" + str(len(vectors[0])))

from pymilvus import MilvusClient
from pymilvus import AnnSearchRequest, RRFRanker
from langchain_community.embeddings import DashScopeEmbeddings
from dashscope import Generation

# 创建Milvus Client。
client = MilvusClient(
    uri="http://c-xxxx.milvus.aliyuncs.com:19530",  # Milvus实例的公网地址。
    token="<yourUsername>:<yourPassword>",  # 登录Milvus实例的用户名和密码。
    db_name="default"  # 待连接的数据库名称，本文示例为默认的default。
)

collection_name = "milvus_bm25_hybrid"

# 替换为您的 DashScope API-KEY
dashscope_api_key = "<YOUR_DASHSCOPE_API_KEY>"

# 初始化 Embedding 模型
embeddings = DashScopeEmbeddings(
    model="text-embedding-v2",  # 使用text-embedding-v2模型。
    dashscope_api_key=dashscope_api_key
)

# Define the query
query = "Why does Milvus run so scalable?"

# Embed the query and generate the corresponding vector representation
query_embeddings = embeddings.embed_documents([query])

# Set the top K result count
top_k = 5  # Get the top 5 docs related to the query

# Define the parameters for the dense vector search
search_params_dense = {
    "metric_type": "IP",
    "params": {"nprobe": 2}
}

# Create a dense vector search request
request_dense = AnnSearchRequest([query_embeddings[0]], "dense", search_params_dense, limit=top_k)

# Define the parameters for the BM25 text search
search_params_bm25 = {
    "metric_type": "BM25"
}

# Create a BM25 text search request
request_bm25 = AnnSearchRequest([query], "text", search_params_bm25, limit=top_k)

# Combine the two requests
reqs = [request_dense, request_bm25]

# Initialize the RRF ranking algorithm
ranker = RRFRanker(100)

# Perform the hybrid search
hybrid_search_res = client.hybrid_search(
    collection_name=collection_name,
    reqs=reqs,
    ranker=ranker,
    limit=top_k,
    output_fields=["text"]
)

# Extract the context from hybrid search results
context = []
print("Top K Results:")
for hits in hybrid_search_res:  # Use the correct variable here
    for hit in hits:
        context.append(hit['entity']['text'])  # Extract text content to the context list
        print(hit['entity']['text'])  # Output each retrieved document

# Define a function to get an answer based on the query and context
def getAnswer(query, context):
    # Build the prompt using the query and context
    prompt_text = "Please answer my question based on the content within:\n\n"
    prompt_text += str(context) + "\n\n"
    prompt_text += "My question is: " + query
    
    # Call the generation module to get an answer
    rsp = Generation.call(model='qwen-turbo', prompt=prompt_text)
    return rsp.output.text

# Get the answer
answer = getAnswer(query, context)

print(answer)

# Expected output excerpt

Milvus is highly scalable due to its cloud-native and highly decoupled system architecture. This architecture allows the system to continuously expand as data grows. Additionally, Milvus supports three deployment modes that cover a wide range of application scenarios.