#!/usr/bin/env python
# coding: utf-8

# <a href="https://colab.research.google.com/github/jerryjliu/llama_index/blob/main/docs/docs/examples/low_level/retrieval.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="在 Colab 中打开"/></a>
# 

# # 从零开始构建检索系统
# 
# 在本教程中，我们将向您展示如何针对向量数据库构建一个标准的检索器，该检索器将通过前k个相似度来获取节点。
# 
# 我们将使用Pinecone作为向量数据库。我们将使用我们的高级摄入抽象来加载节点（要了解如何从头开始构建这个过程，请参阅我们之前的教程！）。
# 
# 我们将展示如何完成以下操作：
# 1. 如何生成查询嵌入
# 2. 如何使用不同的搜索模式（密集、稀疏、混合）查询向量数据库
# 3. 如何将结果解析为一组节点
# 4. 如何将其放入自定义的检索器
# 

# ## 设置
# 
# 我们构建一个空的Pinecone索引，并定义必要的LlamaIndex包装器/抽象，以便我们可以开始将数据加载到Pinecone中。
# 

# 如果您在colab上打开这个笔记本，您可能需要安装LlamaIndex 🦙。
# 

# In[ ]:


get_ipython().run_line_magic('pip', 'install llama-index-readers-file pymupdf')
get_ipython().run_line_magic('pip', 'install llama-index-vector-stores-pinecone')
get_ipython().run_line_magic('pip', 'install llama-index-embeddings-openai')


# In[ ]:


get_ipython().system('pip install llama-index')


# ### 构建Pinecone索引
# 

# In[ ]:


import pinecone
import os

api_key = os.environ["PINECONE_API_KEY"]
pinecone.init(api_key=api_key, environment="us-west1-gcp")


# In[ ]:


# dimensions are for text-embedding-ada-002
pinecone.create_index(
    "quickstart", dimension=1536, metric="euclidean", pod_type="p1"
)


# In[ ]:


pinecone_index = pinecone.Index("quickstart")


# In[ ]:


# [可选] 删除索引中的所有内容
pinecone_index.delete(deleteAll=True)


# #### 创建PineconeVectorStore
# 
# 简单的包装抽象，用于在LlamaIndex中使用。包装在StorageContext中，以便我们可以轻松地加载节点。
# 

# In[ ]:


from llama_index.vector_stores.pinecone import PineconeVectorStore


# In[ ]:


vector_store = PineconeVectorStore(pinecone_index=pinecone_index)


# #### 加载文档
# 

# In[ ]:


get_ipython().system('mkdir data')
get_ipython().system('wget --user-agent "Mozilla" "https://arxiv.org/pdf/2307.09288.pdf" -O "data/llama2.pdf"')


# In[ ]:


from pathlib import Path
from llama_index.readers.file import PyMuPDFReader


# In[ ]:


loader = PyMuPDFReader()
documents = loader.load(file_path="./data/llama2.pdf")


# #### 加载到向量存储
# 
# 将文档加载到PineconeVectorStore中。
# 
# **注意**：在这里我们使用高级的摄取抽象，使用`VectorStoreIndex.from_documents`。在本教程的其余部分中，我们将避免使用`VectorStoreIndex`。
# 

# In[ ]:


from llama_index.core import VectorStoreIndex
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import StorageContext


# In[ ]:


splitter = SentenceSplitter(chunk_size=1024)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents, transformations=[splitter], storage_context=storage_context
)


# ## 定义向量检索器
# 
# 现在我们准备定义针对这个向量存储的检索器，以检索一组节点。
# 
# 我们将逐步展示这些过程，然后将其封装成一个函数。
# 

# In[ ]:


query_str = "Can you tell me about the key concepts for safety finetuning"


# ### 1. 生成查询嵌入
# 
# 我们将使用`sentence-transformers`库来生成查询嵌入。首先，我们需要安装这个库，然后加载预训练的模型。接下来，我们将使用加载的模型来生成查询的嵌入向量。
# 

# In[ ]:


from llama_index.embeddings.openai import OpenAIEmbedding

embed_model = OpenAIEmbedding()


# In[ ]:


query_embedding = embed_model.get_query_embedding(query_str)


# ### 2. 查询向量数据库
# 
# 我们将展示如何使用不同的模式查询向量数据库：默认模式、稀疏模式和混合模式。
# 
# 我们首先构建一个 `VectorStoreQuery`，然后查询向量数据库。
# 

# In[ ]:


# 构建向量存储查询
from llama_index.core.vector_stores import VectorStoreQuery

query_mode = "default"
# query_mode = "sparse"
# query_mode = "hybrid"

vector_store_query = VectorStoreQuery(
    query_embedding=query_embedding, similarity_top_k=2, mode=query_mode
)


# In[ ]:


# 返回一个VectorStoreQueryResult
query_result = vector_store.query(vector_store_query)
query_result


# ### 3. 将结果解析为一组节点
# 
# `VectorStoreQueryResult` 返回一组节点和相似度。我们使用这些信息构造一个 `NodeWithScore` 对象。
# 

# In[ ]:


from llama_index.core.schema import NodeWithScore
from typing import Optional

nodes_with_scores = []
for index, node in enumerate(query_result.nodes):
    score: Optional[float] = None
    if query_result.similarities is not None:
        score = query_result.similarities[index]
    nodes_with_scores.append(NodeWithScore(node=node, score=score))


# In[ ]:


from llama_index.core.response.notebook_utils import display_source_node

for node in nodes_with_scores:
    display_source_node(node, source_length=1000)


# ### 4. 将这些内容放入Retriever中
# 
# 让我们将这些内容放入一个Retriever子类中，这样就可以将其插入到LlamaIndex工作流的其余部分中！
# 

# In[ ]:


from llama_index.core import QueryBundle
from llama_index.core.retrievers import BaseRetriever
from typing import Any, List


class PineconeRetriever(BaseRetriever):
    """基于Pinecone向量存储的检索器。"""

    def __init__(
        self,
        vector_store: PineconeVectorStore,
        embed_model: Any,
        query_mode: str = "default",
        similarity_top_k: int = 2,
    ) -> None:
        """初始化参数。"""
        self._vector_store = vector_store
        self._embed_model = embed_model
        self._query_mode = query_mode
        self._similarity_top_k = similarity_top_k
        super().__init__()

    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
        """检索。"""
        query_embedding = embed_model.get_query_embedding(query_str)
        vector_store_query = VectorStoreQuery(
            query_embedding=query_embedding,
            similarity_top_k=self._similarity_top_k,
            mode=self._query_mode,
        )
        query_result = vector_store.query(vector_store_query)

        nodes_with_scores = []
        for index, node in enumerate(query_result.nodes):
            score: Optional[float] = None
            if query_result.similarities is not None:
                score = query_result.similarities[index]
            nodes_with_scores.append(NodeWithScore(node=node, score=score))

        return nodes_with_scores


# In[ ]:


retriever = PineconeRetriever(
    vector_store, embed_model, query_mode="default", similarity_top_k=2
)


# In[ ]:


retrieved_nodes = retriever.retrieve(query_str)
for node in retrieved_nodes:
    display_source_node(node, source_length=1000)


# ## 将这个内容插入到我们的RetrieverQueryEngine中，以合成一个响应
# 
# **注意**：在未来的教程中，我们将更多地介绍如何从头开始构建响应合成的内容！
# 

# In[ ]:


from llama_index.core.query_engine import RetrieverQueryEngine

query_engine = RetrieverQueryEngine.from_args(retriever)


# In[ ]:


response = query_engine.query(query_str)


# In[ ]:


print(str(response))