upstream-archive/qdrant-client
1
0
Fork 0
mirror of https://github.com/qdrant/qdrant-client.git synced 2026-08-06 01:50:58 -05:00
Code Issues Packages Projects Releases Wiki Activity

fix mypy

Browse Source
This commit is contained in:
Andrey Vasnetsov
2023-03-26 14:52:30 +02:00
parent be8c6c2c49
commit f9e00f3bde
9 changed files with 86 additions and 63 deletions
Download Patch File Download Diff File Expand all files Collapse all files

11
qdrant_client/client_base.py
View File

@@ -1,6 +1,7 @@
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
from qdrant_client.conversions import common_types as types
from qdrant_client.http import models
class QdrantBase:
@@ -27,11 +28,11 @@ class QdrantBase:
query_vector: Union[
types.NumpyArray, Sequence[float], Tuple[str, List[float]], types.NamedVector
],
query_filter: Optional[types.Filter] = None,
search_params: Optional[types.SearchParams] = None,
query_filter: Optional[models.Filter] = None,
search_params: Optional[models.SearchParams] = None,
limit: int = 10,
offset: int = 0,
with_payload: Union[bool, Sequence[str], types.PayloadSelector] = True,
with_payload: Union[bool, Sequence[str], models.PayloadSelector] = True,
with_vectors: Union[bool, Sequence[str]] = False,
score_threshold: Optional[float] = None,
**kwargs: Any,
@@ -515,7 +516,6 @@ class QdrantBase:
self,
collection_name: str,
vectors_config: Union[types.VectorParams, Mapping[str, types.VectorParams]],
init_from: Optional[types.InitFrom] = None,
**kwargs: Any,
) -> bool:
"""Create empty collection with given parameters
@@ -526,7 +526,6 @@ class QdrantBase:
Configuration of the vector storage. Vector params contains size and distance for the vector storage.
If dict is passed, service will create a vector storage for each key in the dict.
If single VectorParams is passed, service will create a single anonymous vector storage.
init_from: Use data stored in another collection to initialize this collection
Returns:
Operation result
@@ -537,7 +536,6 @@ class QdrantBase:
self,
collection_name: str,
vectors_config: Union[types.VectorParams, Mapping[str, types.VectorParams]],
init_from: Optional[types.InitFrom] = None,
**kwargs: Any,
) -> bool:
"""Delete and create empty collection with given parameters
@@ -548,7 +546,6 @@ class QdrantBase:
Configuration of the vector storage. Vector params contains size and distance for the vector storage.
If dict is passed, service will create a vector storage for each key in the dict.
If single VectorParams is passed, service will create a single anonymous vector storage.
init_from: Use data stored in another collection to initialize this collection
Returns:
Operation result

14
qdrant_client/local/distances.py
View File

@@ -2,7 +2,7 @@ from enum import Enum
import numpy as np
from qdrant_client import models
from qdrant_client.http import models
class DistanceOrder(str, Enum):
@@ -24,7 +24,7 @@ def distance_to_order(distance: models.Distance) -> DistanceOrder:
return DistanceOrder.BIGGER_IS_BETTER
def cosine_similarity(query: np.ndarray, vectors: np.ndarray):
def cosine_similarity(query: np.ndarray, vectors: np.ndarray) -> np.ndarray:
"""
Calculate cosine distance between query and vectors
Args:
@@ -38,7 +38,7 @@ def cosine_similarity(query: np.ndarray, vectors: np.ndarray):
return np.dot(vectors, query)
def dot_product(query: np.ndarray, vectors: np.ndarray):
def dot_product(query: np.ndarray, vectors: np.ndarray) -> np.ndarray:
"""
Calculate dot product between query and vectors
Args:
@@ -50,7 +50,7 @@ def dot_product(query: np.ndarray, vectors: np.ndarray):
return np.dot(vectors, query)
def euclidean_distance(query: np.ndarray, vectors: np.ndarray):
def euclidean_distance(query: np.ndarray, vectors: np.ndarray) -> np.ndarray:
"""
Calculate euclidean distance between query and vectors
Args:
@@ -62,7 +62,9 @@ def euclidean_distance(query: np.ndarray, vectors: np.ndarray):
return np.linalg.norm(vectors - query, axis=1)
def calculate_distance(query: np.ndarray, vectors: np.ndarray, distance_type: models.Distance):
def calculate_distance(
query: np.ndarray, vectors: np.ndarray, distance_type: models.Distance
) -> np.ndarray:
if distance_type == models.Distance.COSINE:
return cosine_similarity(query, vectors)
elif distance_type == models.Distance.DOT:
@@ -73,7 +75,7 @@ def calculate_distance(query: np.ndarray, vectors: np.ndarray, distance_type: mo
raise ValueError(f"Unknown distance type {distance_type}")
def test_distances():
def test_distances() -> None:
query = np.array([1.0, 2.0, 3.0])
vectors = np.array([[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]])
assert np.allclose(calculate_distance(query, vectors, models.Distance.COSINE), [1.0, 1.0])

2
qdrant_client/local/geo.py
View File

@@ -27,7 +27,7 @@ def geo_distance(lon1: float, lat1: float, lon2: float, lat2: float) -> float:
return km * 1000
def test_geo_distance():
def test_geo_distance() -> None:
moscow = {"lon": 37.6173, "lat": 55.7558}
london = {"lon": -0.1278, "lat": 51.5074}
berlin = {"lon": 13.4050, "lat": 52.5200}

54
qdrant_client/local/local_collection.py
View File

@@ -1,10 +1,10 @@
from collections import defaultdict
from typing import Dict, List, Optional, Sequence, Tuple, Union
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
import numpy as np
from qdrant_client import models
from qdrant_client.conversions import common_types as types
from qdrant_client.http import models
from qdrant_client.local.distances import (
DistanceOrder,
calculate_distance,
@@ -27,19 +27,19 @@ class LocalCollection:
Args:
location: path to the collection directory. If None, the collection will be created in memory.
"""
self.vectors = {}
self.payload: List[Dict[str, any]] = []
self.deleted = None
self.ids = {} # Mapping from external id to internal id
self.ids_inv = [] # Mapping from internal id to external id
self.vectors: Dict[str, np.ndarray] = {}
self.payload: List[models.Payload] = []
self.deleted = np.zeros(0, dtype=bool)
self.ids: Dict[models.ExtendedPointId, int] = {} # Mapping from external id to internal id
self.ids_inv: List[models.ExtendedPointId] = [] # Mapping from internal id to external id
self.persistent = location is not None
self.storage = None
self.config = config
if self.persistent:
if location is not None:
self.storage = CollectionPersistence(location)
self.load()
def load(self):
def load(self) -> None:
if self.storage is not None:
vectors = defaultdict(list)
for idx, point in enumerate(self.storage.load()):
@@ -161,7 +161,7 @@ class LocalCollection:
)
name, vector = self._resolve_vector_name(query_vector)
result = []
result: List[models.ScoredPoint] = []
if name not in self.vectors:
raise ValueError(f"Vector {name} is not found in the collection")
@@ -213,7 +213,7 @@ class LocalCollection:
ids: Sequence[types.PointId],
with_payload: Union[bool, Sequence[str], types.PayloadSelector] = True,
with_vectors: Union[bool, Sequence[str]] = False,
):
) -> List[models.Record]:
result = []
for point_id in ids:
@@ -244,7 +244,7 @@ class LocalCollection:
using: Optional[str] = None,
lookup_from_collection: Optional["LocalCollection"] = None,
lookup_from_vector_name: Optional[str] = None,
):
) -> List[models.ScoredPoint]:
collection = self if lookup_from_collection is None else lookup_from_collection
search_in_vector_name = using if using is not None else DEFAULT_VECTOR_NAME
vector_name = (
@@ -306,7 +306,7 @@ class LocalCollection:
sorted_ids = sorted(self.ids.items(), key=lambda x: x[0])
result = []
result: List[types.Record] = []
payload_mask = calculate_payload_mask(
payloads=self.payload,
@@ -348,7 +348,7 @@ class LocalCollection:
mask = payload_mask & ~self.deleted
return models.CountResult(count=np.count_nonzero(mask))
def _update_point(self, point: models.PointStruct):
def _update_point(self, point: models.PointStruct) -> None:
idx = self.ids[point.id]
self.payload[idx] = point.payload
@@ -366,12 +366,12 @@ class LocalCollection:
self.deleted[idx] = 0
def _add_point(self, point: models.PointStruct):
def _add_point(self, point: models.PointStruct) -> None:
idx = len(self.ids)
self.ids[point.id] = idx
self.ids_inv.append(point.id)
self.payload.append(point.payload)
self.deleted.append(0)
np.append(self.deleted, 0)
if isinstance(point.vector, list):
vectors = {DEFAULT_VECTOR_NAME: point.vector}
@@ -383,18 +383,18 @@ class LocalCollection:
), f"Expected all vectors to be present: {vectors.keys()} != {self.vectors.keys()}"
for vector_name, vector in vectors.items():
self.vectors[vector_name].append(vector)
np.append(self.vectors[vector_name], vector)
def _upsert_point(self, point: models.PointStruct):
def _upsert_point(self, point: models.PointStruct) -> None:
if point.id in self.ids:
self._update_point(point)
else:
self._add_point(point)
if self.persistent:
if self.storage is not None:
self.storage.persist(point)
def upsert(self, points: Union[List[models.PointStruct], models.Batch]):
def upsert(self, points: Union[List[models.PointStruct], models.Batch]) -> None:
if isinstance(points, list):
for point in points:
self._upsert_point(point)
@@ -422,12 +422,12 @@ class LocalCollection:
else:
raise ValueError(f"Unsupported type: {type(points)}")
def _delete_ids(self, ids: List[types.PointId]):
def _delete_ids(self, ids: List[types.PointId]) -> None:
for point_id in ids:
idx = self.ids[point_id]
self.deleted[idx] = 1
if self.persistent:
if self.storage is not None:
for point_id in ids:
self.storage.delete(point_id)
@@ -463,7 +463,7 @@ class LocalCollection:
selector: Union[
models.Filter, List[models.ExtendedPointId], models.FilterSelector, models.PointIdsList
],
):
) -> None:
ids = self._selector_to_ids(selector)
self._delete_ids(ids)
@@ -473,7 +473,7 @@ class LocalCollection:
selector: Union[
models.Filter, List[models.ExtendedPointId], models.FilterSelector, models.PointIdsList
],
):
) -> None:
ids = self._selector_to_ids(selector)
for point_id in ids:
idx = self.ids[point_id]
@@ -488,7 +488,7 @@ class LocalCollection:
selector: Union[
models.Filter, List[models.ExtendedPointId], models.FilterSelector, models.PointIdsList
],
):
) -> None:
ids = self._selector_to_ids(selector)
for point_id in ids:
idx = self.ids[point_id]
@@ -500,7 +500,7 @@ class LocalCollection:
selector: Union[
models.Filter, List[models.ExtendedPointId], models.FilterSelector, models.PointIdsList
],
):
) -> None:
ids = self._selector_to_ids(selector)
for point_id in ids:
idx = self.ids[point_id]
@@ -513,7 +513,7 @@ class LocalCollection:
selector: Union[
models.Filter, List[models.ExtendedPointId], models.FilterSelector, models.PointIdsList
],
):
) -> None:
ids = self._selector_to_ids(selector)
for point_id in ids:
idx = self.ids[point_id]

12
qdrant_client/local/payload_filters.py
View File

@@ -2,7 +2,7 @@ from typing import Any, List, Optional
import numpy as np
from qdrant_client import models
from qdrant_client.http import models
from qdrant_client.local.geo import geo_distance
from qdrant_client.local.payload_value_extractor import value_by_key
@@ -81,10 +81,10 @@ def check_condition(
values = value_by_key(payload, condition.is_empty.key)
if values is None or len(values) == 0:
return True
if isinstance(condition, models.HasIdCondition):
elif isinstance(condition, models.HasIdCondition):
if point_id in condition.has_id:
return True
if isinstance(condition, models.FieldCondition):
elif isinstance(condition, models.FieldCondition):
values = value_by_key(payload, condition.key)
if condition.match is not None:
if values is None:
@@ -104,9 +104,11 @@ def check_condition(
return any(check_geo_radius(condition.geo_radius, v) for v in values)
if condition.values_count is not None:
return check_values_count(condition.values_count, values)
if isinstance(condition, models.Filter):
elif isinstance(condition, models.Filter):
return check_filter(condition, payload, point_id)
else:
raise ValueError(f"Unknown condition: {condition}")
return False
def check_must(

4
qdrant_client/local/payload_value_extractor.py
View File

@@ -19,7 +19,7 @@ def value_by_key(payload: dict, key: str) -> Optional[List[Any]]:
keys = key.split(".")
result = []
def _get_value(data, k_list):
def _get_value(data: Any, k_list: List[str]) -> None:
if not k_list:
return
@@ -59,7 +59,7 @@ def value_by_key(payload: dict, key: str) -> Optional[List[Any]]:
return result if result else None
def test_value_by_key():
def test_value_by_key() -> None:
payload = {
"name": "John",
"counts": [1, 2, 3],

10
qdrant_client/local/persistence.py
View File

@@ -3,7 +3,7 @@ import pickle
from pathlib import Path
from typing import Iterable
from qdrant_client import models
from qdrant_client.http import models
STORAGE_FILE_NAME = "storage.dbm"
@@ -19,10 +19,10 @@ class CollectionPersistence:
self.location = Path(location) / STORAGE_FILE_NAME
self.storage = dbm.open(str(self.location), "c")
def __del__(self):
def __del__(self) -> None:
self.storage.close()
def persist(self, point: models.PointStruct):
def persist(self, point: models.PointStruct) -> None:
"""
Persist a point in the local storage.
Args:
@@ -32,7 +32,7 @@ class CollectionPersistence:
value = pickle.dumps(point)
self.storage[key] = value
def delete(self, point_id: models.ExtendedPointId):
def delete(self, point_id: models.ExtendedPointId) -> None:
"""
Delete a point from the local storage.
Args:
@@ -53,7 +53,7 @@ class CollectionPersistence:
yield pickle.loads(value)
def test_persistence():
def test_persistence() -> None:
import tempfile
with tempfile.TemporaryDirectory() as tmpdir:

4
qdrant_client/local/qdrant_local.py
View File

@@ -4,9 +4,9 @@ import shutil
from itertools import zip_longest
from typing import Any, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Union
from qdrant_client import models as rest_models
from qdrant_client.client_base import QdrantBase
from qdrant_client.conversions import common_types as types
from qdrant_client.http import models as rest_models
from qdrant_client.local.local_collection import LocalCollection
@@ -383,7 +383,7 @@ class QdrantLocal(QdrantBase):
payload=payload or {},
)
for idx, (point_id, vector, payload) in enumerate(
zip_longest(ids, vectors, payload)
zip_longest(ids or [], vectors, payload or [])
)
]
)

38
qdrant_client/qdrant_client.py
View File

@@ -240,6 +240,7 @@ class QdrantClient(QdrantBase):
collection_name: str,
requests: Sequence[types.SearchRequest],
consistency: Optional[types.ReadConsistency] = None,
**kwargs: Any,
) -> List[List[types.ScoredPoint]]:
"""Search for points in multiple collections
@@ -309,6 +310,7 @@ class QdrantClient(QdrantBase):
score_threshold: Optional[float] = None,
append_payload: bool = True,
consistency: Optional[types.ReadConsistency] = None,
**kwargs: Any,
) -> List[types.ScoredPoint]:
"""Search for closest vectors in collection taking into account filtering conditions
@@ -480,6 +482,7 @@ class QdrantClient(QdrantBase):
collection_name: str,
requests: Sequence[types.RecommendRequest],
consistency: Optional[types.ReadConsistency] = None,
**kwargs: Any,
) -> List[List[types.ScoredPoint]]:
"""Perform multiple recommend requests in batch mode
@@ -551,6 +554,7 @@ class QdrantClient(QdrantBase):
using: Optional[str] = None,
lookup_from: Optional[types.LookupLocation] = None,
consistency: Optional[types.ReadConsistency] = None,
**kwargs: Any,
) -> List[types.ScoredPoint]:
"""Recommend points: search for similar points based on already stored in Qdrant examples.
@@ -737,6 +741,7 @@ class QdrantClient(QdrantBase):
with_payload: Union[bool, Sequence[str], types.PayloadSelector] = True,
with_vectors: Union[bool, Sequence[str]] = False,
consistency: Optional[types.ReadConsistency] = None,
**kwargs: Any,
) -> Tuple[List[types.Record], Optional[types.PointId]]:
"""Scroll over all (matching) points in the collection.
@@ -883,6 +888,7 @@ class QdrantClient(QdrantBase):
points: types.Points,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Update or insert a new point into the collection.
@@ -971,6 +977,7 @@ class QdrantClient(QdrantBase):
with_payload: Union[bool, Sequence[str], types.PayloadSelector] = True,
with_vectors: Union[bool, Sequence[str]] = False,
consistency: Optional[types.ReadConsistency] = None,
**kwargs: Any,
) -> List[types.Record]:
"""Retrieve stored points by IDs
@@ -1153,6 +1160,7 @@ class QdrantClient(QdrantBase):
points_selector: types.PointsSelector,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Deletes selected points from collection
@@ -1212,6 +1220,7 @@ class QdrantClient(QdrantBase):
points: types.PointsSelector,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Modifies payload of the specified points
@@ -1292,6 +1301,7 @@ class QdrantClient(QdrantBase):
points: types.PointsSelector,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Overwrites payload of the specified points
After this operation is applied, only the specified payload will be present in the point.
@@ -1376,6 +1386,7 @@ class QdrantClient(QdrantBase):
points: types.PointsSelector,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Remove values from point's payload
@@ -1438,6 +1449,7 @@ class QdrantClient(QdrantBase):
points_selector: types.PointsSelector,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Delete all payload for selected points
@@ -1494,6 +1506,7 @@ class QdrantClient(QdrantBase):
self,
change_aliases_operations: Sequence[types.AliasOperations],
timeout: Optional[int] = None,
**kwargs: Any,
) -> bool:
"""Operation for performing changes of collection aliases.
@@ -1542,7 +1555,7 @@ class QdrantClient(QdrantBase):
assert result is not None, "Get collection aliases returned None"
return result
def get_aliases(self) -> types.CollectionsAliasesResponse:
def get_aliases(self, **kwargs: Any) -> types.CollectionsAliasesResponse:
"""Get all aliases
Returns:
@@ -1554,7 +1567,7 @@ class QdrantClient(QdrantBase):
assert result is not None, "Get aliases returned None"
return result
def get_collections(self) -> types.CollectionsResponse:
def get_collections(self, **kwargs: Any) -> types.CollectionsResponse:
"""Get list name of all existing collections
Returns:
@@ -1605,6 +1618,7 @@ class QdrantClient(QdrantBase):
optimizer_config: Optional[types.OptimizersConfigDiff] = None,
collection_params: Optional[types.CollectionParamsDiff] = None,
timeout: Optional[int] = None,
**kwargs: Any,
) -> bool:
"""Update parameters of the collection
@@ -1635,7 +1649,9 @@ class QdrantClient(QdrantBase):
assert result is not None, "Update collection returned None"
return result
def delete_collection(self, collection_name: str, timeout: Optional[int] = None) -> bool:
def delete_collection(
self, collection_name: str, timeout: Optional[int] = None, **kwargs: Any
) -> bool:
"""Removes collection and all it's data
Args:
@@ -1667,6 +1683,7 @@ class QdrantClient(QdrantBase):
quantization_config: Optional[types.QuantizationConfig] = None,
init_from: Optional[types.InitFrom] = None,
timeout: Optional[int] = None,
**kwargs: Any,
) -> bool:
"""Create empty collection with given parameters
@@ -1753,6 +1770,7 @@ class QdrantClient(QdrantBase):
quantization_config: Optional[types.QuantizationConfig] = None,
init_from: Optional[types.InitFrom] = None,
timeout: Optional[int] = None,
**kwargs: Any,
) -> bool:
"""Delete and create empty collection with given parameters
@@ -1805,7 +1823,6 @@ class QdrantClient(QdrantBase):
optimizers_config=optimizers_config,
wal_config=wal_config,
quantization_config=quantization_config,
init_from=init_from,
timeout=timeout,
)
@@ -1868,6 +1885,7 @@ class QdrantClient(QdrantBase):
parallel: int = 1,
method: Optional[str] = None,
max_retries: int = 3,
**kwargs: Any,
) -> None:
"""Upload records to the collection
@@ -1899,6 +1917,7 @@ class QdrantClient(QdrantBase):
parallel: int = 1,
method: Optional[str] = None,
max_retries: int = 3,
**kwargs: Any,
) -> None:
"""Upload vectors and payload to the collection.
This method will perform automatic batching of the data.
@@ -1929,6 +1948,7 @@ class QdrantClient(QdrantBase):
field_type: Optional[types.PayloadSchemaType] = None,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Creates index for a given payload field.
Indexed fields allow to perform filtered search operations faster.
@@ -1979,6 +1999,7 @@ class QdrantClient(QdrantBase):
field_name: str,
wait: bool = True,
ordering: Optional[types.WriteOrdering] = None,
**kwargs: Any,
) -> types.UpdateResult:
"""Removes index for a given payload field.
@@ -2060,7 +2081,7 @@ class QdrantClient(QdrantBase):
assert result is not None, "Delete snapshot API returned None"
return result
def list_full_snapshots(self) -> List[types.SnapshotDescription]:
def list_full_snapshots(self, **kwargs: Any) -> List[types.SnapshotDescription]:
"""List all snapshots for a whole storage
Returns:
@@ -2070,7 +2091,7 @@ class QdrantClient(QdrantBase):
assert snapshots is not None, "List full snapshots API returned None result"
return snapshots
def create_full_snapshot(self) -> types.SnapshotDescription:
def create_full_snapshot(self, **kwargs: Any) -> types.SnapshotDescription:
"""Create snapshot for a whole storage.
Returns:
@@ -2100,6 +2121,7 @@ class QdrantClient(QdrantBase):
collection_name: str,
location: str,
priority: Optional[types.SnapshotPriority] = None,
**kwargs: Any,
) -> bool:
"""Recover collection from snapshot.
@@ -2132,7 +2154,7 @@ class QdrantClient(QdrantBase):
assert result is not None, "Lock storage returned None"
return result
def unlock_storage(self) -> types.LocksOption:
def unlock_storage(self, **kwargs: Any) -> types.LocksOption:
"""Unlock storage for writing."""
result: Optional[types.LocksOption] = self.openapi_client.service_api.post_locks(
rest_models.LocksOption(write=False)
@@ -2140,7 +2162,7 @@ class QdrantClient(QdrantBase):
assert result is not None, "Post locks returned None"
return result
def get_locks(self) -> types.LocksOption:
def get_locks(self, **kwargs: Any) -> types.LocksOption:
"""Get current locks state."""
result: Optional[types.LocksOption] = self.openapi_client.service_api.get_locks().result
assert result is not None, "Get locks returned None"