はじめに
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)をまでを整理する予定です)
コードの解説
前提として、事前にVector Indexは生成されているものとしています。
検索のみのコードサンプル
VectorStoreIndexからのデータ検索の一般的な実装は下記の通りです。
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/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
に指定された保存データからdocstore
、index_store
等が順次復元されていきます(111行目以降)
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 )
・・・
- 上記114行目の
index_store
の具体的に復元を見ていきます - ここで
SimpleIndexStoreクラスのfrom_persist_dir()
が呼ばれます(下記、33行目)。 - 次に `from_persist_path()を呼び出します(43行目)
- index_storeをkey-value形式で復元し
SimpleIndexStore
のインスタンスとして返却します(53-54行目) - これにより
persist_dir
に保存された情報からSimpleIndexStore
を生成(復元)します
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)
- 以降、順次復元がされていきます(117行目以降)
- VectorStoreが復元されます(134行目)
- 最後に各復元されたオブジェクトから
StorageContext
インスタンスを生成します(141行目)
・・・
# ⭐️オブジェクトを順次復元(続き)
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 )
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行目)
・・・
# ⭐️ 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
3.indexを検索モジュール(retriever)として生成
retriever = index.as_retriever.retrieve(similarity_top_k=5)
生成したindexをVectorIndex Retriver(検索モジュール)として生成します。
- as_retriever()が呼び出される(113行目)
- ノードのID情報等を与えてVectorIndexRetrieverインスタンスを生成(119行目)
・・・
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 )
・・・
4. 検索モジュールからクエリを与えてデータを取得する
retrieved_docs = retriever.retrieve("ポケモンについて教えて")
先ほど作成した検索モジュールを使って、Queryと類似度が高いデータをindexを検索して返却します。
-
BaseRetriever
クラスのretrieve()
が呼ばれます(222行目) - 類似度が高いデータ(ノード)を取得するため
_retrive()
を呼び出します(245行目)
・・・
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
-
_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行目)
・・・
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)
上記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行目)
・・・
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)
・・・
-
get_top_k_embeddings()
が呼び出されます(11行目) - 類似度比較関数にdefalutの
base/embeddings/base/simirarity
が設定されます(23行目) - 類似度の検証を行います(30行目)
- heapqを利用し
top_k
数を超えたらキューの中から類似度が最も低いものを削除します(34行目)
・・・
# ⭐️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
・・・
similarityは、defaultではコサイン類似度で類似度が測定されます。
(indexにある全ノードと比較されるので件数が極端に多いときは処理性能には注意が必要と思われます)
・・・
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
・・・
以上がlamaindexの検索のみ時(index.as_retriever)のコードの流れになります。
おわりに
コードを見ると意外にシンプルだと思います。少しでもRAGの理解のお役に立てればと思います。
(記載に間違えがあればご指摘ください。反映してきます)
次回は、LLMを検索して回答を得るQueryEngineのコードを解説します。