refactor(tree): extract _embed_l2_subtree to reduce embed_all complexity

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-07-07 01:37:09 -04:00
parent 22ad014973
commit c136de076d
+21 -7
View File
@@ -481,6 +481,27 @@ class TreeIndex:
if l1.embedding is None: if l1.embedding is None:
l1.embedding = embed_fn(l1.summary)[0].astype(np.float32) l1.embedding = embed_fn(l1.summary)[0].astype(np.float32)
for l2 in l1.children: for l2 in l1.children:
self._embed_l2_subtree(l2, embed_fn)
self.metadata.embed_model = model_name
self.metadata.embed_dim = embed_dim
logger.info(
"embed_all 完成",
model=model_name,
embed_dim=embed_dim,
)
def _embed_l2_subtree(
self,
l2: L2Node,
embed_fn: Callable[[str | list[str]], np.ndarray],
) -> None:
"""对单个 L2 节点及其 L3 子节点执行 embedding(仅处理 embedding 为 None 的节点)。
参数:
l2: 待 embed 的 L2 节点。
embed_fn: EmbeddingModel.embed 方法,接受 str 或 List[str]
返回 [N, D] ndarray。
"""
if l2.embedding is None: if l2.embedding is None:
l2.embedding = embed_fn(l2.description)[0].astype(np.float32) l2.embedding = embed_fn(l2.description)[0].astype(np.float32)
# L3 批量 embed # L3 批量 embed
@@ -490,13 +511,6 @@ class TreeIndex:
embs = embed_fn(texts).astype(np.float32) # [N, D] embs = embed_fn(texts).astype(np.float32) # [N, D]
for l3, emb in zip(need_embed, embs, strict=True): for l3, emb in zip(need_embed, embs, strict=True):
l3.embedding = emb l3.embedding = emb
self.metadata.embed_model = model_name
self.metadata.embed_dim = embed_dim
logger.info(
"embed_all 完成",
model=model_name,
embed_dim=embed_dim,
)
# ------------------------------------------------------------------ # # ------------------------------------------------------------------ #
# 嵌入矩阵提取 # 嵌入矩阵提取