# RAG离线链路-ElasticSearch引擎
# 环境需求：
# Python版本：3.7及以上
# ES集群版本：8.9及以上：如果是阿里云ES需要提前开通并设置访问ip白名单 https://help.aliyun.com/zh/es/user-guide/configure-a-public-or-private-ip-address-whitelist-for-an-elasticsearch-cluster
# 包需求：
# pip install alibabacloud_searchplat20240529
# pip install elasticsearch
# OpenSearch搜索开发工作台配置
aisearch_endpoint = "xxx.platform-cn-shanghai.opensearch.aliyuncs.com"
api_key = "OS-xxx"
workspace_name = "default"
service_id_config = {"extract": "ops-document-analyze-001",
                     "split": "ops-document-split-001",
                     "text_embedding": "ops-text-embedding-001",
                     "text_sparse_embedding": "ops-text-sparse-embedding-001",
                     "image_analyze": "ops-image-analyze-ocr-001"}
# ES配置
es_host = 'http://es-cn-xxx.public.elasticsearch.aliyuncs.com:9200'
es_auth = ('elastic', 'xxx')
# 输入文档url,示例文档为opensearch产品说明文档
document_url = "https://help.aliyun.com/zh/open-search/search-platform/product-overview/introduction-to-search-platform?spm=a2c4g.11186623.0.0.7ab93526WDzQ8z"
import asyncio
from typing import List
from elasticsearch import AsyncElasticsearch
from elasticsearch import helpers
from alibabacloud_tea_openapi.models import Config
from alibabacloud_searchplat20240529.client import Client
from alibabacloud_searchplat20240529.models import GetDocumentSplitRequest, CreateDocumentAnalyzeTaskRequest, \
    CreateDocumentAnalyzeTaskRequestDocument, GetDocumentAnalyzeTaskStatusRequest, \
    GetDocumentSplitRequestDocument, GetTextEmbeddingRequest, GetTextEmbeddingResponseBodyResultEmbeddings, \
    GetTextSparseEmbeddingRequest, GetTextSparseEmbeddingResponseBodyResultSparseEmbeddings, \
    CreateImageAnalyzeTaskRequestDocument, CreateImageAnalyzeTaskRequest, CreateImageAnalyzeTaskResponse, \
    GetImageAnalyzeTaskStatusRequest, GetImageAnalyzeTaskStatusResponse
async def poll_task_result(ops_client, task_id, service_id, interval=5):
    while True:
        request = GetDocumentAnalyzeTaskStatusRequest(task_id=task_id)
        response = await ops_client.get_document_analyze_task_status_async(workspace_name, service_id, request)
        status = response.body.result.status
        if status == "PENDING":
            await asyncio.sleep(interval)
        elif status == "SUCCESS":
            return response
        else:
            raise Exception("document analyze task failed")
def is_analyzable_url(url:str):
    if not url:
        return False
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
    return url.lower().endswith(tuple(image_extensions))
async def image_analyze(ops_client, url):
    try:
        print("image analyze :" + url)
        if url.startswith("//"):
            url = "https:" + url
        if not is_analyzable_url(url):
            print(url + " is not analyzable.")
            return url
        image_analyze_service_id = service_id_config["image_analyze"]
        document = CreateImageAnalyzeTaskRequestDocument(
            url=url,
        )
        request = CreateImageAnalyzeTaskRequest(document=document)
        response: CreateImageAnalyzeTaskResponse = ops_client.create_image_analyze_task(workspace_name, image_analyze_service_id, request)
        task_id = response.body.result.task_id
        while True:
            request = GetImageAnalyzeTaskStatusRequest(task_id=task_id)
            response: GetImageAnalyzeTaskStatusResponse = ops_client.get_image_analyze_task_status(workspace_name, image_analyze_service_id, request)
            status = response.body.result.status
            if status == "PENDING":
                await asyncio.sleep(5)
            elif status == "SUCCESS":
                return url + response.body.result.data.content
            else:
                print("image analyze error: " + response.body.result.error)
                return url
    except Exception as e:
        print(f"image analyze Exception : {e}")
def chunk_list(lst, chunk_size):
    for i in range(0, len(lst), chunk_size):
        yield lst[i:i + chunk_size]
async def write_to_es(doc_list):
    es = AsyncElasticsearch(
        [es_host],
        basic_auth=es_auth,
        verify_certs=False,  # 不使用SSL证书校验
        request_timeout=30,
        max_retries=10,
        retry_on_timeout=True
    )
    index_name = 'dense_vertex_index'
    # 删除已有索引
    if await es.indices.exists(index=index_name):
        await es.indices.delete(index=index_name)
    # 创建向量索引, 指定embedding字段为dense_vector, content字段为text, source字段为keyword
    index_mappings = {
        "mappings": {
            "properties": {
                "emb": {
                    "type": "dense_vector",
                    "index": True,
                    "similarity": "cosine",
                    "dims": 1536  # 根据embedding模型输出维度修改
                },
                "content": {
                    "type": "text"
                },
                "source_doc": {
                    "type": "keyword"
                }
            }
        }
    }
    await es.indices.create(index=index_name, body=index_mappings)
    # 上传embedding结果到上一步创建好的索引
    actions = []
    for i, doc in enumerate(doc_list):
        action = {
            "_index": index_name,
            "_id": doc['id'],
            "_source": {
                "emb": doc['embedding'],
                "content": doc['content'],
                "source_doc": document_url
            }
        }
        actions.append(action)
    try:
        await helpers.async_bulk(es, actions)
    except Exception as e:
        for error in e.errors:
            print(error)
    # 确认上传成功
    await asyncio.sleep(2)
    query = {
        "query": {
            "ids": {
                "values": [doc_list[0]["id"]]
            }
        }
    }
    res = await es.search(index=index_name, body=query)
    if len(res['hits']['hits']) > 0:
        print("ES write success")
    await es.close()
async def document_pipeline_execute(document_url: str = None, document_base64: str = None, file_name: str = None):
    # 生成opensearch开发平台client
    config = Config(bearer_token=api_key, endpoint=aisearch_endpoint, protocol="http")
    ops_client = Client(config=config)
    # Step 1: 文档解析
    document_analyze_request = CreateDocumentAnalyzeTaskRequest(
        document=CreateDocumentAnalyzeTaskRequestDocument(url=document_url, content=document_base64,
                                                          file_name=file_name, file_type='html'))
    document_analyze_response = await ops_client.create_document_analyze_task_async(workspace_name=workspace_name,
                                                                                    service_id=service_id_config[
                                                                                        "extract"],
                                                                                    request=document_analyze_request)
    print("document-analyze task_id:" + document_analyze_response.body.result.task_id)
    extraction_result = await poll_task_result(ops_client, document_analyze_response.body.result.task_id,
                                               service_id_config["extract"])
    print("document-analyze done")
    document_content = extraction_result.body.result.data.content
    content_type = extraction_result.body.result.data.content_type
    # Step 2: 文档切片
    document_split_request = GetDocumentSplitRequest(
        GetDocumentSplitRequestDocument(content=document_content, content_type=content_type))
    document_split_result = await ops_client.get_document_split_async(workspace_name, service_id_config["split"],
                                                                      document_split_request)
    print("document-split done, chunks count: " + str(len(document_split_result.body.result.chunks))
          + " rich text count:" + str(len(document_split_result.body.result.rich_texts)))
    # Step 3: 文本向量化
    # 提取切片结果。图片切片会通过图片解析服务提取出文本内容
    doc_list = ([{"id": chunk.meta.get("id"), "content": chunk.content} for chunk in
                 document_split_result.body.result.chunks]
                + [{"id": chunk.meta.get("id"), "content": chunk.content} for chunk in
                   document_split_result.body.result.rich_texts if chunk.meta.get("type") != "image"]
                + [{"id": chunk.meta.get("id"), "content": await image_analyze(ops_client, chunk.content)} for chunk in
                   document_split_result.body.result.rich_texts if chunk.meta.get("type") == "image"]
                )
    chunk_size = 32  # 一次最多允许计算32个embedding
    all_text_embeddings: List[GetTextEmbeddingResponseBodyResultEmbeddings] = []
    for chunk in chunk_list([text["content"] for text in doc_list], chunk_size):
        response = await ops_client.get_text_embedding_async(workspace_name, service_id_config["text_embedding"],
                                                             GetTextEmbeddingRequest(chunk))
        all_text_embeddings.extend(response.body.result.embeddings)
    all_text_sparse_embeddings: List[GetTextSparseEmbeddingResponseBodyResultSparseEmbeddings] = []
    for chunk in chunk_list([text["content"] for text in doc_list], chunk_size):
        response = await ops_client.get_text_sparse_embedding_async(workspace_name,
                                                                    service_id_config["text_sparse_embedding"],
                                                                    GetTextSparseEmbeddingRequest(chunk,
                                                                                                  input_type="document",
                                                                                                  return_token=True))
        all_text_sparse_embeddings.extend(response.body.result.sparse_embeddings)
    for i in range(len(doc_list)):
        doc_list[i]["embedding"] = all_text_embeddings[i].embedding
        doc_list[i]["sparse_embedding"] = all_text_sparse_embeddings[i].embedding
    print("text-embedding done")
    # Step 4: 写入ElasticSearch存储引擎
    await write_to_es(doc_list)
if __name__ == "__main__":
    # 运行异步任务
    #    import nest_asyncio # 如果在Jupyter notebook中运行，反注释这两行
    #    nest_asyncio.apply() # 如果在Jupyter notebook中运行，反注释这两行
    asyncio.run(document_pipeline_execute(document_url))
    # asyncio.run(document_pipeline_execute(document_base64="eHh4eHh4eHg...", file_name="attention.pdf")) #另外一种调用方式

# RAG在线链路-ElasticSearch引擎
# 环境需求：
# Python版本：3.7及以上
# ES集群版本：8.9以上（如果是阿里云ES 需要提前开通并设置访问ip白名单 https://help.aliyun.com/zh/es/user-guide/configure-a-public-or-private-ip-address-whitelist-for-an-elasticsearch-cluster）
# 包需求：
# pip install alibabacloud_searchplat20240529
# pip install elasticsearch
# OpenSearch搜索开发工作台配置
api_key = "OS-xxx"
aisearch_endpoint = "xxx.platform-cn-shanghai.opensearch.aliyuncs.com"
workspace_name = "default"
service_id_config = {
    "rank": "ops-bge-reranker-larger",
    "text_embedding": "ops-text-embedding-001",
    "text_sparse_embedding": "ops-text-sparse-embedding-001",
    "llm": "ops-qwen-turbo",
    "query_analyze": "ops-query-analyze-001"
}
# ES配置
es_host = 'http://es-cn-xxx.public.elasticsearch.aliyuncs.com:9200'
es_auth = ('elastic', 'xxx')
# 用户query:
user_query = "OpenSearch搜索开发平台可以做什么？"
import asyncio
from elasticsearch import AsyncElasticsearch
from alibabacloud_tea_openapi.models import Config
from alibabacloud_searchplat20240529.client import Client
from alibabacloud_searchplat20240529.models import GetTextEmbeddingRequest,  \
    GetDocumentRankRequest, GetTextGenerationRequest, GetTextGenerationRequestMessages, \
    GetQueryAnalysisRequest
# 生成opensearch开发平台client
config = Config(bearer_token=api_key, endpoint=aisearch_endpoint, protocol="http")
ops_client = Client(config=config)
async def es_retrieve(query):
    es = AsyncElasticsearch(
        [es_host],
        basic_auth=es_auth,
        verify_certs=False,
        request_timeout=30,
        max_retries=10,
        retry_on_timeout=True
    )
    index_name = 'dense_vertex_index'
    # query向量化
    query_emb_result = await ops_client.get_text_embedding_async(workspace_name, service_id_config["text_embedding"],
                                                                 GetTextEmbeddingRequest(input=[query],
                                                                                         input_type="query"))
    query_emb = query_emb_result.body.result.embeddings[0].embedding
    query = {
        "field": "emb",
        "query_vector": query_emb,
        "k": 5,  # 返回文档切片数量
        "num_candidates": 100  # HNSW搜索参数efsearch
    }
    res = await es.search(index=index_name, knn=query)
    search_results = [item['_source']['content'] for item in res['hits']['hits']]
    await es.close()
    return search_results
# 在线问答流水线，输入是用户问题
async def query_pipeline_execute():
    # Step 1: query分析
    query_analyze_response = ops_client.get_query_analysis(workspace_name, service_id_config['query_analyze'],
                                                           GetQueryAnalysisRequest(query=user_query))
    print("query analysis rewrite result:" + query_analyze_response.body.result.query)
    # Step 2: 召回文档
    all_query_results = []
    user_query_results = await es_retrieve(user_query)
    all_query_results.extend(user_query_results)
    rewrite_query_results = await es_retrieve(query_analyze_response.body.result.query)
    all_query_results.extend(rewrite_query_results)
    for extend_query in query_analyze_response.body.result.queries:
        extend_query_result = await es_retrieve(extend_query)
        all_query_results.extend(extend_query_result)
    # 对所有召回结果进行去重
    remove_duplicate_results = list(set(all_query_results))
    # Step 3: 对召回文档进行重排序
    rerank_top_k = 8
    score_results = await ops_client.get_document_rank_async(workspace_name, service_id_config["rank"],GetDocumentRankRequest(remove_duplicate_results, user_query))
    rerank_results = [remove_duplicate_results[item.index] for item in score_results.body.result.scores[:rerank_top_k]]
    # Step 4: 调用大模型回答问题
    docs = '\n'.join([f"<article>{s}</article>" for s in rerank_results])
    messages = [
        GetTextGenerationRequestMessages(role="system", content="You are a helpful assistant."),
        GetTextGenerationRequestMessages(role="user",
                                         content=f"""已知信息包含多个独立文档，每个文档在<article>和</article>之间，已知信息如下：\n'''{docs}'''
                                         \n\n根据上述已知信息，详细且有条理地回答用户的问题。确保答案充分回答了问题并且正确使用了已知信息。如果信息不足以回答问题，请说“根据已知信息无法回答该问题”。
                                         不要使用不在已知信息中的内容生成答案，确保答案中每一个陈述在上述已知信息中有相应内容支撑。答案请使用中文。
                                         \n问题是：'''{user_query}'''""""")
    ]
    response = await ops_client.get_text_generation_async(workspace_name, service_id_config["llm"],
                                                          GetTextGenerationRequest(messages=messages))
    print("大模型最终回答： ", response.body.result.text)
if __name__ == "__main__":
    #    import nest_asyncio # 如果在Jupyter notebook中运行，反注释这两行
    #    nest_asyncio.apply() # 如果在Jupyter notebook中运行，反注释这两行
    asyncio.run(query_pipeline_execute())