Why not login to Qiita and try out its useful features?

We'll deliver articles that match you.

You can read useful information later.

18
17

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?

RAG(LlamaIndex)をDeepに理解しよう!コード解説(as_retriver)

Last updated at Posted at 2025-03-18

はじめに

OpenAI が AI Agents SDK(しかも OSS!(^^))を発表し、いよいよ AI エージェントの実サービスへの活用が本格化しそうな気配が感じられます。
その一方で、これからの AI エージェントや RAG(Retrieval-Augmented Generation)をうまく使いこなすためには、RAG の基本をしっかりと整理し、RAG や AI エージェントで「できること」「組み合わせて可能になること」をあらためて考えることが重要だと感じています。

OpenAI Agents SDK

今回は、自分自身の理解を整理する目的も兼ねて、LlamaIndex のコードを読み解きながら、基本的な使い方やその裏側の実装がどうなっているのか をまとめてみました。

対象とする LlamaIndex のバージョンは、2025年3月7日時点の最新バージョン「0.12.23」です。

第一弾として、最も基本的なVectorStoreIndexから検索(主にindex.as_retriever()の範囲)する処理を解説します。
(第二弾として、QueryEngineによる生成AI検索(主にindex.as_query_engine)をまでを整理する予定です)

llamaindex_構造.png

コードの解説

前提として、事前にVector Indexは生成されているものとしています。

検索のみのコードサンプル

VectorStoreIndexからのデータ検索の一般的な実装は下記の通りです。

test.py
1  # ✅ 1.保存したVectorStoreIndexからStorageContextを復元
2  storage_context = StorageContext.from_defaults(persist_dir=LLAMAINDEX_PERSIST_FOLDER)  #セーブしたフォルダ名
3  # ✅ 2.StorageContextからindex(VectorStoreIndex)を復元
4  index = load_index_from_storage(storage_context)
5  # ✅ 3.indexを検索モジュール(retriever)として生成。この時、クエリに対して類似度の高い上位5件のデータを取得するよう設定
6  retriever = index.as_retriever(similarity_top_k=5)
7  # ✅ 4.検索モジュールからクエリを与えてデータを(5つ)取得する
8  retrieved_docs = retriever.retrieve("ポケモンについて教えて")
9
10 # データを表示。retriever.retrieve()の戻り値は List[NodeWithScore]
11 # 各データは node.text(テキスト)や node.metadata(メタデータ、設定されている場合)から取得する
12 response_text = ""
13 for doc in retrieved_docs:
14     response_text += f"{doc.node.text}\n"
15     response_text += f"{doc.node.metadata}\n"

以降、それぞれの処理でどのような処理が実際に行われているのか、順次コードを追って見ていきます。

1. 保存したVectorStoreIndexからStorageContextを復元

storage_context = StorageContext.from_defaults()

コードを見る前に、まずはセーブされたVector Store Indexのディレクトリ構造を見てみましょう。データはJSON形式で保存されています。それぞれのファイル名と概要は以下のとおりです(データ構造の詳細はドキュメント等には開示されていません)。

ファイル名 概要
default_vector_store.json 埋め込みデータ(主にクエリーとデータとの類似度の比較に 使用
docstore.json ドキュメントストア(データやメタデータ)
graph_store.json グラフストア(基本未使用)
image_vector_store.json 画像ベクトルストア(画像データを使わなければ未使用)
index_store.json ノード(文書やデータの単位)の情報

index_storeには、下記のようにindexストアの情報も__type__で格納されています(下記の例の場合は、VectorStoreIndex)。

index_store.json
{"index_store/data": {"a0e20a7c-f561-4f66-b23a-xxxxxx": {"__type__": "vector_store", 
"__data__": "{\"index_id\": \"a0e20a7c-f561-4f66-b23a-xxxxxx\",・・・・

StorageContext.from_defaults() では、この保存したJSONファイルを読み込んでStorageContextインスタンスを生成(復元)します。

以下、内部でどのような処理が行われているのか解説していきます。

  • StorageContext クラスの from_defaults() が呼ばれます(73行目)
  • persist_dir に指定された保存データから docstoreindex_store等が順次復元されていきます(111行目以降)
storage/storage_context.py
 52 @dataclass
 53 class StorageContext:
 ・・・
 72 @classmethod
    # ⭐️from_defaultsが呼び出される
 73   def from_defaults(
 74       cls,
 75       docstore: Optional[BaseDocumentStore] = None,
 76       index_store: Optional[BaseIndexStore] = None,
 77       vector_store: Optional[BasePydanticVectorStore] = None,
 78       image_store: Optional[BasePydanticVectorStore] = None,
 79       vector_stores: Optional[Dict[str, BasePydanticVectorStore]] = None,
 80       graph_store: Optional[GraphStore] = None,
 81       property_graph_store: Optional[PropertyGraphStore] = None,
 82       persist_dir: Optional[str] = None,
 83       fs: Optional[fsspec.AbstractFileSystem] = None,
 84   ) -> "StorageContext":
 85       """Create a StorageContext from defaults.
 86
 87       Args:
 88           docstore (Optional[BaseDocumentStore]): document store
 89           index_store (Optional[BaseIndexStore]): index store
 90           vector_store (Optional[BasePydanticVectorStore]): vector store
 91           graph_store (Optional[GraphStore]): graph store
 92           image_store (Optional[BasePydanticVectorStore]): image store
 93
 94       """
        
        print("#### Aelurus: storage/storage_context.py - 001")
        
 95       if persist_dir is None:
 96           docstore = docstore or SimpleDocumentStore()
 97           index_store = index_store or SimpleIndexStore()
 98           graph_store = graph_store or SimpleGraphStore()
 99           image_store = image_store or SimpleVectorStore()
100
101           if vector_store:
102               vector_stores = {DEFAULT_VECTOR_STORE: vector_store}
103           else:
104               vector_stores = vector_stores or {
105               DEFAULT_VECTOR_STORE: SimpleVectorStore()
106               }
107            if image_store:
108                # append image store to vector stores
109                vector_stores[IMAGE_VECTOR_STORE_NAMESPACE] = image_store
110       else:
               # ⭐️オブジェクトを順次復元
111            docstore = docstore or SimpleDocumentStore.from_persist_dir(
112                persist_dir, fs=fs
113            )
            
            print("#### Aelurus: storage/storage_context.py - 002")
               # ⭐️index_soreを復元
114            index_store = index_store or SimpleIndexStore.from_persist_dir(
115                persist_dir, fs=fs
116            )
・・・            

storage_context.pyの全文はこちら

  • 上記114行目のindex_store の具体的に復元を見ていきます
  • ここで SimpleIndexStoreクラスのfrom_persist_dir() が呼ばれます(下記、33行目)。
  • 次に `from_persist_path()を呼び出します(43行目)
  • index_storeをkey-value形式で復元し SimpleIndexStore のインスタンスとして返却します(53-54行目)
  • これにより persist_dir に保存された情報から SimpleIndexStoreを生成(復元)します
storage/index_store/simple_index_store.py
 16 class SimpleIndexStore(KVIndexStore):
 17     """Simple in-memory Index store.
 18
 19     Args:
 20         simple_kvstore (SimpleKVStore): simple key-value store
 21
 22     """
 23
 24     def __init__(
 25         self,
 26         simple_kvstore: Optional[SimpleKVStore] = None,
 27     ) -> None:
 28         """Init a SimpleIndexStore."""
 29         simple_kvstore = simple_kvstore or SimpleKVStore()
 30         super().__init__(simple_kvstore)
 31
        # ⭐️from_persist_dirが呼び出される
 32     @classmethod
 33     def from_persist_dir(
 34         cls,
 35         persist_dir: str = DEFAULT_PERSIST_DIR,
 36         fs: Optional[fsspec.AbstractFileSystem] = None,
 37     ) -> "SimpleIndexStore":
 38         """Create a SimpleIndexStore from a persist directory."""     
 39         if fs is not None:
 40             persist_path = concat_dirs(persist_dir, DEFAULT_PERSIST_FNAME)
 41         else:
 42             persist_path = os.path.join(persist_dir, DEFAULT_PERSIST_FNAME)
           # ⭐️from_persist_pathを呼び出しリターン
 43         return cls.from_persist_path(persist_path, fs=fs)
 44
 45     @classmethod
 46     def from_persist_path(
 47         cls,
 48         persist_path: str,
 49         fs: Optional[fsspec.AbstractFileSystem] = None,
 50     ) -> "SimpleIndexStore":
 51         """Create a SimpleIndexStore from a persist path."""
        # ⭐️index_storeをkey-value形式で復元しSimpleIndexStoreのインスタンスとして返却
 52         fs = fs or fsspec.filesystem("file")
 53         simple_kvstore = SimpleKVStore.from_persist_path(persist_path, fs=fs)
 54         return cls(simple_kvstore)

simple_index_store.pyの全文はこちら

  • 以降、順次復元がされていきます(117行目以降)
  • VectorStoreが復元されます(134行目)
  • 最後に各復元されたオブジェクトから StorageContext インスタンスを生成します(141行目)
storage/storage_context.py
・・・
         # ⭐️オブジェクトを順次復元(続き)
117            graph_store = graph_store or SimpleGraphStore.from_persist_dir(
118                persist_dir, fs=fs
119            )
120
121            try:
122                property_graph_store = (
123                    property_graph_store
124                    or SimplePropertyGraphStore.from_persist_dir(persist_dir, fs=fs)
125                )
126            except FileNotFoundError:
127                property_graph_store = None
128
129            if vector_store:               
130                vector_stores = {DEFAULT_VECTOR_STORE: vector_store}
131            elif vector_stores:
132                vector_stores = vector_stores
133            else:
                   # ⭐️vector_storeを復元(続き)
134                vector_stores = SimpleVectorStore.from_namespaced_persist_dir(
135                    persist_dir, fs=fs
136                )
137            if image_store:
138                # append image store to vector stores
139                vector_stores[IMAGE_VECTOR_STORE_NAMESPACE] = image_store  # type: ignore
140
           # ⭐️生成オブジェクトからStorageContextインスタンスを生成し復元
141        return cls(
142            docstore=docstore,
143            index_store=index_store,
144            vector_stores=vector_stores,  # type: ignore
145            graph_store=graph_store,
146            property_graph_store=property_graph_store,
147        )

storage_context.pyの全文はこちら

2. StorageContextからindex(VectorStoreIndex)を復元

index = load_index_from_storage(storage_context)

ここでは、生成(復元)したStorageContextからindexを生成します。

  • load_index_from_storage() が呼ばれます(12行目)
  • index_idは指定されていないので、Indexの情報を取得します(基本1つしかないはず)(64行目)
  • 続いて、このindexのタイプを取得しますが、これは先ほどのVecotStoreIndexのディレクトリ構造で説明した、inedex_store内のJSONデータの __type__ を取得しますので、"vector_sore"、クラス名=VectorStoreIndexが返却されます(76,77行目)
  • VectorStoreIndexのインスタンスが生成され返却されます(78行目)
indices/loading.py
・・・
    # ⭐️ load_index_from_storageの呼び出し
 12 def load_index_from_storage(
 13     storage_context: StorageContext,
 14     index_id: Optional[str] = None,
 15     **kwargs: Any,
 16 ) -> BaseIndex:
 17     """Load index from storage context.
 18
 19     Args:
 20         storage_context (StorageContext): storage context containing
 21             docstore, index store and vector store.
 22         index_id (Optional[str]): ID of the index to load.
 23             Defaults to None, which assumes there's only a single index
 24             in the index store and load it.
 25         **kwargs: Additional keyword args to pass to the index constructors.
 26     """
 27     index_ids: Optional[Sequence[str]]
 28     if index_id is None:
 29         index_ids = None
 30     else:
 31         index_ids = [index_id]
 32
        # ⭐️Storegeコンテキストからindexの読み込み 
 33     indices = load_indices_from_storage(storage_context, index_ids=index_ids, **kwargs)
 34
 35     if len(indices) == 0:
 36         raise ValueError(
 37             "No index in storage context, check if you specified the right persist_dir."
 38         )
 39     elif len(indices) > 1:
 40         raise ValueError(
 41             f"Expected to load a single index, but got {len(indices)} instead. "
 42             "Please specify index_id."
 43         )
 44
 45     return indices[0]
 46
 47
 48 def load_indices_from_storage(
 49     storage_context: StorageContext,
 50     index_ids: Optional[Sequence[str]] = None,
 51     **kwargs: Any,
 52 ) -> List[BaseIndex]:
 53     """Load multiple indices from storage context.
 54
 55     Args:
 56         storage_context (StorageContext): storage context containing
 57             docstore, index store and vector store.
 58         index_id (Optional[Sequence[str]]): IDs of the indices to load.
 59             Defaults to None, which loads all indices in the index store.
 60         **kwargs: Additional keyword args to pass to the index constructors.
 61     """
 62     if index_ids is None:
 63         logger.info("Loading all indices.")
            # ⭐️Storageコンテキストのindexを取得する(基本は1つ) 
 64         index_structs = storage_context.index_store.index_structs()
 65     else:
 66         logger.info(f"Loading indices with ids: {index_ids}")
 67         index_structs = []
 68         for index_id in index_ids:
 69             index_struct = storage_context.index_store.get_index_struct(index_id)
 70             if index_struct is None:
 71                 raise ValueError(f"Failed to load index with ID {index_id}")
 72             index_structs.append(index_struct)
 73
 74     indices = []
 75     for index_struct in index_structs:
 76         type_ = index_struct.get_type()
            # ⭐️type_は「IndexStructType.VECTOR_STORE(VectorStoreIndex」が返却
 77         index_cls = INDEX_STRUCT_TYPE_TO_INDEX_CLASS[type_]
            # ⭐️ VectorStoreIndexクラスが生成される
 78         index = index_cls(
 79             index_struct=index_struct, storage_context=storage_context, **kwargs
 80         )
 81         indices.append(index)
 82     return indices

loading.pyの全文はこちら

3.indexを検索モジュール(retriever)として生成

retriever = index.as_retriever.retrieve(similarity_top_k=5)

生成したindexをVectorIndex Retriver(検索モジュール)として生成します。

  • as_retriever()が呼び出される(113行目)
  • ノードのID情報等を与えてVectorIndexRetrieverインスタンスを生成(119行目)
indices/vector_store/base.py
・・・
 36 class VectorStoreIndex(BaseIndex[IndexDict]):
 37     """
 38     Vector Store Index.
 39  
 40     Args:
 41         use_async (bool): Whether to use asynchronous calls. Defaults to False.
 42         show_progress (bool): Whether to show tqdm progress bars. Defaults to False.
 43         store_nodes_override (bool): set to True to always store Node objects in index
 44             store and document store even if vector store keeps text. Defaults to False
 45   """
・・・
      # ⭐️as_retriever()が呼ばれる
 113 def as_retriever(self, **kwargs: Any) -> BaseRetriever:
 114     # NOTE: lazy import
 115     from llama_index.core.indices.vector_store.retrievers import (
 116         VectorIndexRetriever,
 117     )
 118
         # ⭐️VectorIndexRetrieverの生成
 119     return VectorIndexRetriever(
 120         self,
 121         node_ids=list(self.index_struct.nodes_dict.values()),
 122         callback_manager=self._callback_manager,
 123         object_map=self._object_map,
 124         **kwargs,
 125     )
 ・・・

vector_store/base.pyの全文はこちら

4. 検索モジュールからクエリを与えてデータを取得する

retrieved_docs = retriever.retrieve("ポケモンについて教えて")

先ほど作成した検索モジュールを使って、Queryと類似度が高いデータをindexを検索して返却します。

  • BaseRetrieverクラスの retrieve() が呼ばれます(222行目)
  • 類似度が高いデータ(ノード)を取得するため _retrive()を呼び出します(245行目)
base/base_retriever.py
・・・
 42 class BaseRetriever(ChainableMixin, PromptMixin, DispatcherSpanMixin):
 43     """Base retriever."""
 44
 45     def __init__(
 46         self,
 47         callback_manager: Optional[CallbackManager] = None,
 48         object_map: Optional[Dict] = None,
 49         objects: Optional[List[IndexNode]] = None,
 50         verbose: bool = False,
 51     ) -> None:
・・・
    # ⭐️retrieve()が呼び出される
222 def retrieve(self, str_or_query_bundle: QueryType) -> List[NodeWithScore]:
223     """Retrieve nodes given query.
224
225     Args:
226          str_or_query_bundle (QueryType): Either a query string or
227              a QueryBundle object.
228
229      """
230     self._check_callback_manager()
231     dispatcher.event(
232         RetrievalStartEvent(
233             str_or_query_bundle=str_or_query_bundle,
234         )
235     )
236     if isinstance(str_or_query_bundle, str):
237         query_bundle = QueryBundle(str_or_query_bundle)
238     else:
239         query_bundle = str_or_query_bundle
240     with self.callback_manager.as_trace("query"):
241         with self.callback_manager.event(
242             CBEventType.RETRIEVE,
243             payload={EventPayload.QUERY_STR: query_bundle.query_str},
244         ) as retrieve_event:
                # ⭐️ノードを取得 
245             nodes = self._retrieve(query_bundle)
246             nodes = self._handle_recursive_retrieval(query_bundle, nodes)
247             retrieve_event.on_end(
248                 payload={EventPayload.NODES: nodes},
249             )
250     dispatcher.event(
251         RetrievalEndEvent(
252             str_or_query_bundle=str_or_query_bundle,
253             nodes=nodes,
254         )
255     )
256     return nodes

base_retriever.pyの全文はこちら

  • _retrieve() は、VectorIndexRetrieverクラスの _retrieve()が呼ばれます(92行目)
  • _get_nodes_with_embeddings() を呼び出してリターンします(103行目)
  • _get_nodes_with_embeddings() では、_build_vector_store_query()を呼び出し、VectorStoreQuery オブジェクトを生成します(179行目)
  • さらに、VectorStore(SimpleVectorStore)query()しノード(データ)を取得します(180行目)
  • _build_node_list_from_query_result()を呼び出し、検索されたノードindexからノードIDリストを作り、ドキュメントストアから対応するノード(データ)を取得します
  • 検索されたデータ(Node)とScoreで、NodeWithScoreオブジェクトを生成し、リターンします(172-173行目)
indices/vector_store/retrivers/retriever.py
・・・
 24 class VectorIndexRetriever(BaseRetriever):
 25     """Vector index retriever.
 26
 27     Args:
 28         index (VectorStoreIndex): vector store index.
 29         similarity_top_k (int): number of top k results to return.
 30         vector_store_query_mode (str): vector store query mode
 31             See reference for VectorStoreQueryMode for full list of supported modes.
 32         filters (Optional[MetadataFilters]): metadata filters, defaults to None
 33         alpha (float): weight for sparse/dense retrieval, only used for
 34             hybrid query mode.
 35         doc_ids (Optional[List[str]]): list of documents to constrain search.
 36         vector_store_kwargs (dict): Additional vector store specific kwargs to pass
 37             through to the vector store at query time.
 38
 39     """
 40
 41     def __init__(
 42         self,
 43         index: VectorStoreIndex,
 44         similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
 45         vector_store_query_mode: VectorStoreQueryMode = VectorStoreQueryMode.DEFAULT,
 46         filters: Optional[MetadataFilters] = None,
 47         alpha: Optional[float] = None,
 48         node_ids: Optional[List[str]] = None,
 49         doc_ids: Optional[List[str]] = None,
 50         sparse_top_k: Optional[int] = None,
 51         hybrid_top_k: Optional[int] = None,
 52         callback_manager: Optional[CallbackManager] = None,
 53         object_map: Optional[dict] = None,
 54         embed_model: Optional[BaseEmbedding] = None,
 55         verbose: bool = False,
 56         **kwargs: Any,
 57     ) -> None:
・・・
        # ⭐️VectorIndexのノード検索 
 92     def _retrieve(
 93         self,
 94         query_bundle: QueryBundle,
 95     ) -> List[NodeWithScore]:
 96         if self._vector_store.is_embedding_query:
 97             if query_bundle.embedding is None and len(query_bundle.embedding_strs) > 0:
 98                 query_bundle.embedding = (
 99                     self._embed_model.get_agg_embedding_from_queries(
100                         query_bundle.embedding_strs
101                     )
102                 )
            # ⭐️ _get_nodes_with_embeddings()を呼び出す
103         return self._get_nodes_with_embeddings(query_bundle)
・・・
134     def _build_node_list_from_query_result(
135         self, query_result: VectorStoreQueryResult
136     ) -> List[NodeWithScore]:
137         if query_result.nodes is None:
                # ⭐️refineでなければこの処理が行われる  
138             # NOTE: vector store does not keep text and returns node indices.
139             # Need to recover all nodes from docstore
140             if query_result.ids is None:
141                 raise ValueError(
142                     "Vector store query result should return at "
143                     "least one of nodes or ids."
144                 )
145             assert isinstance(self._index.index_struct, IndexDict)
                # ⭐️検索されたノードindexからノードIDリストを作り、ドキュメントストアから対応するノード(データ)を取得する   
146             node_ids = [
147                 self._index.index_struct.nodes_dict[idx] for idx in query_result.ids
148             ]
149             nodes = self._docstore.get_nodes(node_ids)
150             query_result.nodes = nodes
151         else:
・・・
164
165         log_vector_store_query_result(query_result)
166
167         node_with_scores: List[NodeWithScore] = []
168         for ind, node in enumerate(query_result.nodes):
169             score: Optional[float] = None
170             if query_result.similarities is not None:
171                 score = query_result.similarities[ind]
                # ⭐️検索されたデータ(Node)とScoreでNodeWithScoreオブジェクトを生成し追加する
172             node_with_scores.append(NodeWithScore(node=node, score=score))
173         return node_with_scores
・・・
176     def _get_nodes_with_embeddings(
177         self, query_bundle_with_embeddings: QueryBundle
178     ) -> List[NodeWithScore]:
            # ⭐️ _build_vector_store_query()を呼び出し、VectorStoreQeryオブジェクトを生成
179         query = self._build_vector_store_query(query_bundle_with_embeddings)
            #  ⭐️ VectorStore(vector_stores.simple.SimpleVectorStore)をquery()し結果を取得
180         query_result = self._vector_store.query(query, **self._kwargs)
        return self._build_node_list_from_query_result(query_result)

retriever.pyの全文はこちら

上記180行目の self._vector_store.query()メソッドは、先ほどのStorageContexの生成で出てきた、storage/storage_context.pyの134行目 vector_stores = SimpleVectorStore.from_namespaced_persist_dir()で生成した SimpleVectorStoreクラスが使われます。

  • query()を呼び出し、類似度が高いVestorストア上のノードを検索します(317行目)
  • デフォルトのqueryモードで実行され、get_top_k_embeddings() を呼び出す(376行目)
vector_store/simple.py
・・・
139 class SimpleVectorStore(BasePydanticVectorStore):
140     """Simple Vector Store.
141
142     In this vector store, embeddings are stored within a simple, in-memory dictionary.
143
144     Args:
145         simple_vector_store_data_dict (Optional[dict]): data dict
146             containing the embeddings and doc_ids. See SimpleVectorStoreData
147             for more details.
148     """
・・・
        # ⭐️queryと類似度が高いVestorストア上のノードを検索
317     def query(
318         self,
319         query: VectorStoreQuery,
320         **kwargs: Any,
321     ) -> VectorStoreQueryResult:
322         """Get nodes for response."""
・・・
337
338         if query.node_ids is not None:
339             available_ids = set(query.node_ids)
340
341             def node_filter_fn(node_id: str) -> bool:
342                 return node_id in available_ids
343
344         else:
・・・
349         node_ids = []
350         embeddings = []
351         # TODO: consolidate with get_query_text_embedding_similarities
352         for node_id, embedding in self.data.embedding_dict.items():
353             if node_filter_fn(node_id) and query_filter_fn(node_id):
354                 node_ids.append(node_id)
355                 embeddings.append(embedding)
356
357         query_embedding = cast(List[float], query.query_embedding)
358
359         if query.mode in LEARNER_MODES:
・・・
366         elif query.mode == MMR_MODE:
・・・
375         elif query.mode == VectorStoreQueryMode.DEFAULT:
                # ⭐️デフォルトのqueryモードが実行される
376             top_similarities, top_ids = get_top_k_embeddings(
377                 query_embedding,
378                 embeddings,
379                 similarity_top_k=query.similarity_top_k,
380                 embedding_ids=node_ids,
381             )
382         else:
383             raise ValueError(f"Invalid query mode: {query.mode}")
384
385         return VectorStoreQueryResult(similarities=top_similarities, ids=top_ids)
・・・

vector_store/simple.pyの全文はこちら

  • get_top_k_embeddings() が呼び出されます(11行目)
  • 類似度比較関数にdefalutの base/embeddings/base/simirarityが設定されます(23行目)
  • 類似度の検証を行います(30行目)
  • heapqを利用しtop_k 数を超えたらキューの中から類似度が最も低いものを削除します(34行目)
indices/query/embedding_utils.py
・・・
    # ⭐️get_top_k_embeddings()の呼び出し
 11 def get_top_k_embeddings(
 12     query_embedding: List[float],
 13     embeddings: List[List[float]],
 14     similarity_fn: Optional[Callable[..., float]] = None,
 15     similarity_top_k: Optional[int] = None,
 16     embedding_ids: Optional[List] = None,
 17     similarity_cutoff: Optional[float] = None,
 18 ) -> Tuple[List[float], List]:
 19     """Get top nodes by similarity to the query."""
 20     if embedding_ids is None:
 21         embedding_ids = list(range(len(embeddings)))
 22
        # ⭐️defalutの類似度比較関数であるbase/embeddings/base/simirarityが呼ばれる
 23     similarity_fn = similarity_fn or default_similarity_fn
 24
 25     embeddings_np = np.array(embeddings)
 26     query_embedding_np = np.array(query_embedding)
 27
 28     similarity_heap: List[Tuple[float, Any]] = []
 29     for i, emb in enumerate(embeddings_np):
            # ⭐️類似度の検証
 30         similarity = similarity_fn(query_embedding_np, emb)  # type: ignore[arg-type]
 31         if similarity_cutoff is None or similarity > similarity_cutoff:
 32             heapq.heappush(similarity_heap, (similarity, embedding_ids[i]))
 33             if similarity_top_k and len(similarity_heap) > similarity_top_k:
                    # ⭐️top_k数を超えたらheapqで類似度が最も低いものを削除    
 34                 heapq.heappop(similarity_heap)
 35     result_tups = sorted(similarity_heap, key=lambda x: x[0], reverse=True)
 36
 37     result_similarities = [s for s, _ in result_tups]
 38     result_ids = [n for _, n in result_tups]
 39
 40     return result_similarities, result_ids
・・・

vector_store/simple.pyの全文はこちら

similarityは、defaultではコサイン類似度で類似度が測定されます。
(indexにある全ノードと比較されるので件数が極端に多いときは処理性能には注意が必要と思われます)

base/embeddings/base.py
・・・
 50 def similarity(
 51     embedding1: Embedding,
 52     embedding2: Embedding,
 53     mode: SimilarityMode = SimilarityMode.DEFAULT,
 54 ) -> float:
 55     """Get embedding similarity."""
 56     if mode == SimilarityMode.EUCLIDEAN:
 57         # Using -euclidean distance as similarity to achieve same ranking order
 58         return -float(np.linalg.norm(np.array(embedding1) - np.array(embedding2)))
 59     elif mode == SimilarityMode.DOT_PRODUCT:
 60         return np.dot(embedding1, embedding2)
 61     else:
 62         # ⭐️COS類似度の算出
 63         product = np.dot(embedding1, embedding2)
 64         norm = np.linalg.norm(embedding1) * np.linalg.norm(embedding2)
 65         return product / norm
・・・

embeddings/base.pyの全文はこちら

以上がlamaindexの検索のみ時(index.as_retriever)のコードの流れになります。

おわりに

コードを見ると意外にシンプルだと思います。少しでもRAGの理解のお役に立てればと思います。
(記載に間違えがあればご指摘ください。反映してきます)

次回は、LLMを検索して回答を得るQueryEngineのコードを解説します。

18
17
0

Register as a new user and use Qiita more conveniently

  1. You get articles that match your needs
  2. You can efficiently read back useful information
  3. You can use dark theme
What you can do with signing up

Qiita Conference 2025 will be held!: 4/23(wed) - 4/25(Fri)

Qiita Conference is the largest tech conference in Qiita!

Keynote Speaker

ymrl、Masanobu Naruse, Takeshi Kano, Junichi Ito, uhyo, Hiroshi Tokumaru, MinoDriven, Minorun, Hiroyuki Sakuraba, tenntenn, drken, konifar

View event details
18
17

Delete article

Deleted articles cannot be recovered.

Draft of this article would be also deleted.

Are you sure you want to delete this article?