From 3519b718ea989cdea9806137e1cc2cf45d652667 Mon Sep 17 00:00:00 2001 From: Roland Minrui <114476598+RolandMinrui@users.noreply.github.com> Date: Fri, 14 Mar 2025 16:08:29 +0800 Subject: [PATCH] feat: add constraint labels for semantic search (#680) * add constraint labels for semantic search * reformat * reformat * reformat --------- Co-authored-by: Xu --- .../components/knowledge_management/graph.py | 22 ++++++++---- .../knowledge_management/vector_base.py | 35 ++++++++++++++----- 2 files changed, 43 insertions(+), 14 deletions(-) diff --git a/rdagent/components/knowledge_management/graph.py b/rdagent/components/knowledge_management/graph.py index 9ea5b7a5..c6ac5661 100644 --- a/rdagent/components/knowledge_management/graph.py +++ b/rdagent/components/knowledge_management/graph.py @@ -271,27 +271,37 @@ class UndirectedGraph(Graph): node: UndirectedNode | str, similarity_threshold: float = 0.0, topk_k: int = 5, + constraint_labels: list[str] | None = None, ) -> list[UndirectedNode]: """ - semantic search by node's embedding + Semantic search by node's embedding. Parameters ---------- - topk_k - node - similarity_threshold: Returns nodes whose distance score from the input - node is greater than similarity_threshold + node : UndirectedNode | str + The node to search for. + similarity_threshold : float, optional + The minimum similarity score for a node to be included in the results. + Nodes with a similarity score less than or equal to this threshold will be excluded. + topk_k : int, optional + The maximum number of similar nodes to return. + constraint_labels : list[str], optional + If provided, only nodes with matching labels will be considered. Returns ------- - + list[UndirectedNode] + A list of `topk_k` nodes that are semantically similar to the input node, sorted by similarity score. + All nodes shall meet the `similarity_threshold` and `constraint_labels` criteria. """ + # Question: why do we need to convert to Node object first? if isinstance(node, str): node = UndirectedNode(content=node) docs, scores = self.vector_base.search( content=node.content, topk_k=topk_k, similarity_threshold=similarity_threshold, + constraint_labels=constraint_labels, ) return [self.get_node(doc.id) for doc in docs] diff --git a/rdagent/components/knowledge_management/vector_base.py b/rdagent/components/knowledge_management/vector_base.py index bb98caa2..2f3d00a6 100644 --- a/rdagent/components/knowledge_management/vector_base.py +++ b/rdagent/components/knowledge_management/vector_base.py @@ -155,29 +155,48 @@ class PDVectorBase(VectorBase): for doc in document: self.add(document=doc) - def search(self, content: str, topk_k: int = 5, similarity_threshold: float = 0) -> Tuple[List[Document], List]: + def search( + self, content: str, topk_k: int = 5, similarity_threshold: float = 0, constraint_labels: list[str] | None = None + ) -> Tuple[List[Document], List]: """ - search vector by node + Search vector by node's embedding. + Parameters ---------- - similarity_threshold - content - topk_k: return topk_k nearest vector + content : str + The content to search for. + topk_k : int, optional + The number of nearest vectors to return. + similarity_threshold : float, optional + The minimum similarity score for a vector to be considered. + constraint_labels : List[str], optional + If provided, only nodes with matching labels will be considered. Returns ------- - + Tuple[List[Document], List] + A list of `topk_k` nodes that are semantically similar to the input node, sorted by similarity score. + All nodes shall meet the `similarity_threshold` and `constraint_labels` criteria. """ if not self.vector_df.shape[0]: return [], [] + document = Document(content=content) document.create_embedding() - similarities = self.vector_df["embedding"].apply( + + filtered_df = self.vector_df + if constraint_labels is not None: + filtered_df = self.vector_df[self.vector_df["label"].isin(constraint_labels)] + + similarities = filtered_df["embedding"].apply( lambda x: 1 - cosine(x, document.embedding) ) # cosine is cosine distance, 1-similarity + searched_similarities = similarities[similarities > similarity_threshold].nlargest(topk_k) - most_similar_docs = self.vector_df.loc[searched_similarities.index] + most_similar_docs = filtered_df.loc[searched_similarities.index] + docs = [] for _, similar_docs in most_similar_docs.iterrows(): docs.append(Document().from_dict(similar_docs.to_dict())) + return docs, searched_similarities.to_list()