From 4cdd976399e44f1bfc5c8a38abedac6dad446447 Mon Sep 17 00:00:00 2001 From: Alberto Lanaro Date: Tue, 4 Jun 2024 15:45:56 +0200 Subject: [PATCH] Fix qdrant `remove_training_data` api Return `True` for successful deletion --- src/vanna/qdrant/qdrant.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/vanna/qdrant/qdrant.py b/src/vanna/qdrant/qdrant.py index 687de92b..3a70156c 100644 --- a/src/vanna/qdrant/qdrant.py +++ b/src/vanna/qdrant/qdrant.py @@ -3,6 +3,7 @@ import pandas as pd from qdrant_client import QdrantClient, grpc, models +from qdrant_client.http.models.models import UpdateStatus from ..base import VannaBase from ..utils import deterministic_uuid @@ -210,7 +211,8 @@ def get_training_data(self, **kwargs) -> pd.DataFrame: def remove_training_data(self, id: str, **kwargs) -> bool: try: id, collection_name = self._parse_point_id(id) - self._client.delete(collection_name, points_selector=[id]) + res = self._client.delete(collection_name, points_selector=[id]) + res == UpdateStatus.COMPLETED except ValueError: return False