gavindiaz/NexQuant
1
0
Fork 0
mirror of https://github.com/NicolasBohn/NexQuant.git synced 2026-07-29 00:17:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add constraint labels for semantic search (#680)

Browse Source 
* add constraint labels for semantic search

* reformat

* reformat

* reformat

---------

Co-authored-by: Xu <v-xuminrui@microsoft.com>
This commit is contained in:
Roland Minrui
2025-03-14 16:08:29 +08:00
committed by GitHub
parent cde173688b
commit 3519b718ea
2 changed files with 43 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rdagent/components/knowledge_management/graph.py
+16 -6
View File
@@ -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]
rdagent/components/knowledge_management/vector_base.py
+27 -8
View File
@@ -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()