MCPcopy Create free account
hub / github.com/dgraph-io/dgraph-benchmarks / vector_search

Method vector_search

vector/beir/dgraph_client.py:118–234  ·  view source on GitHub ↗

Perform vector similarity search with computed similarity scores. Args: query_embedding: Query embedding vector top_k: Number of results to return ef: Override efSearch for this query (PR#9514, requires use_new_syntax=True) di

(
        self, 
        query_embedding: np.ndarray, 
        top_k: int = 10,
        ef: int = 0,
        distance_threshold: float = 0.0,
        use_new_syntax: bool = False
    )

Source from the content-addressed store, hash-verified

116 print(f"Inserted {total_docs} documents")
117
118 def vector_search(
119 self,
120 query_embedding: np.ndarray,
121 top_k: int = 10,
122 ef: int = 0,
123 distance_threshold: float = 0.0,
124 use_new_syntax: bool = False
125 ) -> List[Tuple[str, float]]:
126 """
127 Perform vector similarity search with computed similarity scores.
128
129 Args:
130 query_embedding: Query embedding vector
131 top_k: Number of results to return
132 ef: Override efSearch for this query (PR#9514, requires use_new_syntax=True)
133 distance_threshold: Filter results by distance threshold (PR#9514, requires use_new_syntax=True)
134 use_new_syntax: Use new similar_to syntax from PR#9514
135
136 Returns:
137 List of tuples (doc_id, score) sorted by similarity (higher is better)
138 """
139 query_vector = query_embedding.tolist()
140 vector_str = json.dumps(query_vector)
141
142 # Build parameterized query with variables
143 variables = {
144 "$v": vector_str,
145 "$topK": str(top_k),
146 }
147
148 # Build similar_to function call based on syntax version
149 if use_new_syntax:
150 # New syntax from PR#9514: similar_to(field, topK, $v, ef: $ef, distance_threshold: $threshold)
151 param_parts = []
152 var_declarations = ["$v: float32vector", "$topK: int"]
153
154 if ef > 0:
155 param_parts.append("ef: $ef")
156 var_declarations.append("$ef: int")
157 variables["$ef"] = str(ef)
158 if distance_threshold > 0.0:
159 param_parts.append("distance_threshold: $threshold")
160 var_declarations.append("$threshold: float")
161 variables["$threshold"] = str(distance_threshold)
162
163 var_decl_str = ", ".join(var_declarations)
164 if param_parts:
165 params_str = ", ".join(param_parts)
166 similar_to_func = f"similar_to(embedding, $topK, $v, {params_str})"
167 else:
168 similar_to_func = "similar_to(embedding, $topK, $v)"
169
170 query = f"""
171 query search({var_decl_str}) {{
172 results(func: {similar_to_func}) {{
173 doc_id
174 embedding
175 }}

Callers 1

batch_vector_searchMethod · 0.95

Calls

no outgoing calls

Tested by

no test coverage detected