Skip to content

Commit

Permalink
update embedding retrieval to use v4 API
Browse files Browse the repository at this point in the history
  • Loading branch information
hsm207 committed Feb 27, 2024
1 parent 995287d commit 756b9ef
Showing 1 changed file with 35 additions and 24 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -494,30 +494,41 @@ def _embedding_retrieval(
msg = "Can't use 'distance' and 'certainty' parameters together"
raise ValueError(msg)

collection_name = self._collection_settings["class"]
properties = self._client.schema.get(self._collection_settings["class"]).get("properties", [])
properties = [prop["name"] for prop in properties]

near_vector: Dict[str, Union[float, List[float]]] = {
"vector": query_embedding,
}
if distance is not None:
near_vector["distance"] = distance

if certainty is not None:
near_vector["certainty"] = certainty

query_builder = (
self._client.query.get(collection_name, properties=properties)
.with_near_vector(near_vector)
.with_additional(["vector"])
)
# collection_name = self._collection_settings["class"]
# properties = self._client.schema.get(self._collection_settings["class"]).get("properties", [])
# properties = [prop["name"] for prop in properties]

if filters:
query_builder = query_builder.with_where(convert_filters(filters))
# near_vector: Dict[str, Union[float, List[float]]] = {
# "vector": query_embedding,
# }
# if distance is not None:
# near_vector["distance"] = distance

# if certainty is not None:
# near_vector["certainty"] = certainty

if top_k:
query_builder = query_builder.with_limit(top_k)
# query_builder = (
# self._client.query.get(collection_name, properties=properties)
# .with_near_vector(near_vector)
# .with_additional(["vector"])
# )

# if filters:
# query_builder = query_builder.with_where(convert_filters(filters))

result = query_builder.do()
return [self._to_document(doc) for doc in result["data"]["Get"][collection_name]]
# if top_k:
# query_builder = query_builder.with_limit(top_k)

# result = query_builder.do()
# return [self._to_document(doc) for doc in result["data"]["Get"][collection_name]]

result = self._collection.query.near_vector(
near_vector=query_embedding,
distance=distance,
certainty=certainty,
include_vector=True,
filters=convert_filters(filters) if filters else None,
limit=top_k,
)

return [self._to_document(self._convert_weaviate_v4_object_to_v3_object(doc)) for doc in result.objects]

0 comments on commit 756b9ef

Please sign in to comment.