qdrant_fastembed.py 34 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892
  1. import uuid
  2. from itertools import tee
  3. from typing import Any, Iterable, Optional, Sequence, Union, get_args
  4. from copy import deepcopy
  5. import numpy as np
  6. from pydantic import BaseModel
  7. from qdrant_client import grpc
  8. from qdrant_client.common.client_warnings import show_warning
  9. from qdrant_client.client_base import QdrantBase
  10. from qdrant_client.embed.model_embedder import ModelEmbedder
  11. from qdrant_client.http import models
  12. from qdrant_client.conversions import common_types as types
  13. from qdrant_client.conversions.conversion import GrpcToRest
  14. from qdrant_client.embed.common import INFERENCE_OBJECT_TYPES
  15. from qdrant_client.embed.schema_parser import ModelSchemaParser
  16. from qdrant_client.hybrid.fusion import reciprocal_rank_fusion
  17. from qdrant_client.fastembed_common import FastEmbedMisc, OnnxProvider
  18. # region imports used in deprecated methods
  19. from qdrant_client.fastembed_common import (
  20. QueryResponse,
  21. TextEmbedding,
  22. SparseTextEmbedding,
  23. IDF_EMBEDDING_MODELS,
  24. )
  25. # endregion
  26. class QdrantFastembedMixin(QdrantBase):
  27. DEFAULT_EMBEDDING_MODEL = "BAAI/bge-small-en"
  28. DEFAULT_BATCH_SIZE = 8
  29. _FASTEMBED_INSTALLED: bool
  30. def __init__(self, parser: ModelSchemaParser, **kwargs: Any):
  31. self._embedding_model_name: Optional[str] = None
  32. self._sparse_embedding_model_name: Optional[str] = None
  33. self._model_embedder = ModelEmbedder(parser=parser, **kwargs)
  34. self.__class__._FASTEMBED_INSTALLED = FastEmbedMisc.is_installed()
  35. super().__init__(**kwargs)
  36. @classmethod
  37. def list_text_models(cls) -> dict[str, tuple[int, models.Distance]]:
  38. """Lists the supported dense text models.
  39. Returns:
  40. dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics.
  41. """
  42. return FastEmbedMisc.list_text_models()
  43. @classmethod
  44. def list_image_models(cls) -> dict[str, tuple[int, models.Distance]]:
  45. """Lists the supported image dense models.
  46. Returns:
  47. dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics.
  48. """
  49. return FastEmbedMisc.list_image_models()
  50. @classmethod
  51. def list_late_interaction_text_models(cls) -> dict[str, tuple[int, models.Distance]]:
  52. """Lists the supported late interaction text models.
  53. Returns:
  54. dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics.
  55. """
  56. return FastEmbedMisc.list_late_interaction_text_models()
  57. @classmethod
  58. def list_late_interaction_multimodal_models(cls) -> dict[str, tuple[int, models.Distance]]:
  59. """Lists the supported late interaction multimodal models.
  60. Returns:
  61. dict[str, tuple[int, models.Distance]]: A dict of model names, their dimensions and distance metrics.
  62. """
  63. return FastEmbedMisc.list_late_interaction_multimodal_models()
  64. @classmethod
  65. def list_sparse_models(cls) -> dict[str, dict[str, Any]]:
  66. """Lists the supported sparse text models.
  67. Returns:
  68. dict[str, dict[str, Any]]: A dict of model names and their descriptions.
  69. """
  70. return FastEmbedMisc.list_sparse_models()
  71. @property
  72. def embedding_model_name(self) -> str:
  73. if self._embedding_model_name is None:
  74. self._embedding_model_name = self.DEFAULT_EMBEDDING_MODEL
  75. return self._embedding_model_name
  76. @property
  77. def sparse_embedding_model_name(self) -> Optional[str]:
  78. return self._sparse_embedding_model_name
  79. def set_model(
  80. self,
  81. embedding_model_name: str,
  82. max_length: Optional[int] = None,
  83. cache_dir: Optional[str] = None,
  84. threads: Optional[int] = None,
  85. providers: Optional[Sequence["OnnxProvider"]] = None,
  86. cuda: bool = False,
  87. device_ids: Optional[list[int]] = None,
  88. lazy_load: bool = False,
  89. **kwargs: Any,
  90. ) -> None:
  91. """
  92. Set embedding model to use for encoding documents and queries.
  93. Args:
  94. embedding_model_name: One of the supported embedding models. See `SUPPORTED_EMBEDDING_MODELS` for details.
  95. max_length (int, optional): Deprecated. Defaults to None.
  96. cache_dir (str, optional): The path to the cache directory.
  97. Can be set using the `FASTEMBED_CACHE_PATH` env variable.
  98. Defaults to `fastembed_cache` in the system's temp directory.
  99. threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
  100. providers: The list of onnx providers (with or without options) to use. Defaults to None.
  101. Example configuration:
  102. https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options
  103. cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
  104. Defaults to False.
  105. device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
  106. workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
  107. lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
  108. Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
  109. Raises:
  110. ValueError: If embedding model is not supported.
  111. ImportError: If fastembed is not installed.
  112. Returns:
  113. None
  114. """
  115. if max_length is not None:
  116. show_warning(
  117. message="max_length parameter is deprecated and will be removed in the future. "
  118. "It's not used by fastembed models.",
  119. category=DeprecationWarning,
  120. stacklevel=3,
  121. )
  122. self._get_or_init_model(
  123. model_name=embedding_model_name,
  124. cache_dir=cache_dir,
  125. threads=threads,
  126. providers=providers,
  127. cuda=cuda,
  128. device_ids=device_ids,
  129. lazy_load=lazy_load,
  130. deprecated=True,
  131. **kwargs,
  132. )
  133. self._embedding_model_name = embedding_model_name
  134. def set_sparse_model(
  135. self,
  136. embedding_model_name: Optional[str],
  137. cache_dir: Optional[str] = None,
  138. threads: Optional[int] = None,
  139. providers: Optional[Sequence["OnnxProvider"]] = None,
  140. cuda: bool = False,
  141. device_ids: Optional[list[int]] = None,
  142. lazy_load: bool = False,
  143. **kwargs: Any,
  144. ) -> None:
  145. """
  146. Set sparse embedding model to use for hybrid search over documents in combination with dense embeddings.
  147. Args:
  148. embedding_model_name: One of the supported sparse embedding models. See `SUPPORTED_SPARSE_EMBEDDING_MODELS` for details.
  149. If None, sparse embeddings will not be used.
  150. cache_dir (str, optional): The path to the cache directory.
  151. Can be set using the `FASTEMBED_CACHE_PATH` env variable.
  152. Defaults to `fastembed_cache` in the system's temp directory.
  153. threads (int, optional): The number of threads single onnxruntime session can use. Defaults to None.
  154. providers: The list of onnx providers (with or without options) to use. Defaults to None.
  155. Example configuration:
  156. https://onnxruntime.ai/docs/execution-providers/CUDA-ExecutionProvider.html#configuration-options
  157. cuda (bool, optional): Whether to use cuda for inference. Mutually exclusive with `providers`
  158. Defaults to False.
  159. device_ids (Optional[list[int]], optional): The list of device ids to use for data parallel processing in
  160. workers. Should be used with `cuda=True`, mutually exclusive with `providers`. Defaults to None.
  161. lazy_load (bool, optional): Whether to load the model during class initialization or on demand.
  162. Should be set to True when using multiple-gpu and parallel encoding. Defaults to False.
  163. Raises:
  164. ValueError: If embedding model is not supported.
  165. ImportError: If fastembed is not installed.
  166. Returns:
  167. None
  168. """
  169. if embedding_model_name is not None:
  170. self._get_or_init_sparse_model(
  171. model_name=embedding_model_name,
  172. cache_dir=cache_dir,
  173. threads=threads,
  174. providers=providers,
  175. cuda=cuda,
  176. device_ids=device_ids,
  177. lazy_load=lazy_load,
  178. deprecated=True,
  179. **kwargs,
  180. )
  181. self._sparse_embedding_model_name = embedding_model_name
  182. @classmethod
  183. def _get_model_params(cls, model_name: str) -> tuple[int, models.Distance]:
  184. FastEmbedMisc.import_fastembed()
  185. for descriptions in (
  186. FastEmbedMisc.list_text_models(),
  187. FastEmbedMisc.list_image_models(),
  188. FastEmbedMisc.list_late_interaction_text_models(),
  189. FastEmbedMisc.list_late_interaction_multimodal_models(),
  190. ):
  191. if params := descriptions.get(model_name):
  192. return params
  193. if model_name in FastEmbedMisc.list_sparse_models():
  194. raise ValueError(
  195. "Sparse embeddings do not return fixed embedding size and distance type"
  196. )
  197. raise ValueError(f"Unsupported embedding model: {model_name}")
  198. def _get_or_init_model(
  199. self,
  200. model_name: str,
  201. cache_dir: Optional[str] = None,
  202. threads: Optional[int] = None,
  203. providers: Optional[Sequence["OnnxProvider"]] = None,
  204. deprecated: bool = False,
  205. **kwargs: Any,
  206. ) -> "TextEmbedding":
  207. FastEmbedMisc.import_fastembed()
  208. return self._model_embedder.embedder.get_or_init_model(
  209. model_name=model_name,
  210. cache_dir=cache_dir,
  211. threads=threads,
  212. providers=providers,
  213. deprecated=deprecated,
  214. **kwargs,
  215. )
  216. def _get_or_init_sparse_model(
  217. self,
  218. model_name: str,
  219. cache_dir: Optional[str] = None,
  220. threads: Optional[int] = None,
  221. providers: Optional[Sequence["OnnxProvider"]] = None,
  222. deprecated: bool = False,
  223. **kwargs: Any,
  224. ) -> "SparseTextEmbedding":
  225. FastEmbedMisc.import_fastembed()
  226. return self._model_embedder.embedder.get_or_init_sparse_model(
  227. model_name=model_name,
  228. cache_dir=cache_dir,
  229. threads=threads,
  230. providers=providers,
  231. deprecated=deprecated,
  232. **kwargs,
  233. )
  234. def _embed_documents(
  235. self,
  236. documents: Iterable[str],
  237. embedding_model_name: str = DEFAULT_EMBEDDING_MODEL,
  238. batch_size: int = 32,
  239. embed_type: str = "default",
  240. parallel: Optional[int] = None,
  241. ) -> Iterable[tuple[str, list[float]]]:
  242. embedding_model = self._get_or_init_model(model_name=embedding_model_name, deprecated=True)
  243. documents_a, documents_b = tee(documents, 2)
  244. if embed_type == "passage":
  245. vectors_iter = embedding_model.passage_embed(
  246. documents_a, batch_size=batch_size, parallel=parallel
  247. )
  248. elif embed_type == "query":
  249. vectors_iter = (
  250. list(embedding_model.query_embed(query=query))[0] for query in documents_a
  251. )
  252. elif embed_type == "default":
  253. vectors_iter = embedding_model.embed(
  254. documents_a, batch_size=batch_size, parallel=parallel
  255. )
  256. else:
  257. raise ValueError(f"Unknown embed type: {embed_type}")
  258. for vector, doc in zip(vectors_iter, documents_b):
  259. yield doc, vector.tolist()
  260. def _sparse_embed_documents(
  261. self,
  262. documents: Iterable[str],
  263. embedding_model_name: str = DEFAULT_EMBEDDING_MODEL,
  264. batch_size: int = 32,
  265. parallel: Optional[int] = None,
  266. ) -> Iterable[types.SparseVector]:
  267. sparse_embedding_model = self._get_or_init_sparse_model(
  268. model_name=embedding_model_name, deprecated=True
  269. )
  270. vectors_iter = sparse_embedding_model.embed(
  271. documents, batch_size=batch_size, parallel=parallel
  272. )
  273. for sparse_vector in vectors_iter:
  274. yield types.SparseVector(
  275. indices=sparse_vector.indices.tolist(),
  276. values=sparse_vector.values.tolist(),
  277. )
  278. def get_vector_field_name(self) -> str:
  279. """
  280. Returns name of the vector field in qdrant collection, used by current fastembed model.
  281. Returns:
  282. Name of the vector field.
  283. """
  284. model_name = self.embedding_model_name.split("/")[-1].lower()
  285. return f"fast-{model_name}"
  286. def get_sparse_vector_field_name(self) -> Optional[str]:
  287. """
  288. Returns name of the vector field in qdrant collection, used by current fastembed model.
  289. Returns:
  290. Name of the vector field.
  291. """
  292. if self.sparse_embedding_model_name is not None:
  293. model_name = self.sparse_embedding_model_name.split("/")[-1].lower()
  294. return f"fast-sparse-{model_name}"
  295. return None
  296. def _scored_points_to_query_responses(
  297. self,
  298. scored_points: list[types.ScoredPoint],
  299. ) -> list[QueryResponse]:
  300. response = []
  301. vector_field_name = self.get_vector_field_name()
  302. sparse_vector_field_name = self.get_sparse_vector_field_name()
  303. for scored_point in scored_points:
  304. embedding = (
  305. scored_point.vector.get(vector_field_name, None)
  306. if isinstance(scored_point.vector, dict)
  307. else None
  308. )
  309. sparse_embedding = None
  310. if sparse_vector_field_name is not None:
  311. sparse_embedding = (
  312. scored_point.vector.get(sparse_vector_field_name, None)
  313. if isinstance(scored_point.vector, dict)
  314. else None
  315. )
  316. response.append(
  317. QueryResponse(
  318. id=scored_point.id,
  319. embedding=embedding,
  320. sparse_embedding=sparse_embedding,
  321. metadata=scored_point.payload,
  322. document=scored_point.payload.get("document", ""),
  323. score=scored_point.score,
  324. )
  325. )
  326. return response
  327. def _points_iterator(
  328. self,
  329. ids: Optional[Iterable[models.ExtendedPointId]],
  330. metadata: Optional[Iterable[dict[str, Any]]],
  331. encoded_docs: Iterable[tuple[str, list[float]]],
  332. ids_accumulator: list,
  333. sparse_vectors: Optional[Iterable[types.SparseVector]] = None,
  334. ) -> Iterable[models.PointStruct]:
  335. if ids is None:
  336. ids = iter(lambda: uuid.uuid4().hex, None)
  337. if metadata is None:
  338. metadata = iter(lambda: {}, None)
  339. if sparse_vectors is None:
  340. sparse_vectors = iter(lambda: None, True)
  341. vector_name = self.get_vector_field_name()
  342. sparse_vector_name = self.get_sparse_vector_field_name()
  343. for idx, meta, (doc, vector), sparse_vector in zip(
  344. ids, metadata, encoded_docs, sparse_vectors
  345. ):
  346. ids_accumulator.append(idx)
  347. payload = {"document": doc, **meta}
  348. point_vector: dict[str, models.Vector] = {vector_name: vector}
  349. if sparse_vector_name is not None and sparse_vector is not None:
  350. point_vector[sparse_vector_name] = sparse_vector
  351. yield models.PointStruct(id=idx, payload=payload, vector=point_vector)
  352. def _validate_collection_info(self, collection_info: models.CollectionInfo) -> None:
  353. embeddings_size, distance = self._get_model_params(model_name=self.embedding_model_name)
  354. vector_field_name = self.get_vector_field_name()
  355. # Check if collection has compatible vector params
  356. assert isinstance(
  357. collection_info.config.params.vectors, dict
  358. ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}"
  359. assert (
  360. vector_field_name in collection_info.config.params.vectors
  361. ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}, expected {vector_field_name}"
  362. vector_params = collection_info.config.params.vectors[vector_field_name]
  363. assert (
  364. embeddings_size == vector_params.size
  365. ), f"Embedding size mismatch: {embeddings_size} != {vector_params.size}"
  366. assert (
  367. distance == vector_params.distance
  368. ), f"Distance mismatch: {distance} != {vector_params.distance}"
  369. sparse_vector_field_name = self.get_sparse_vector_field_name()
  370. if sparse_vector_field_name is not None:
  371. assert (
  372. sparse_vector_field_name in collection_info.config.params.sparse_vectors
  373. ), f"Collection have incompatible vector params: {collection_info.config.params.vectors}"
  374. if self.sparse_embedding_model_name in IDF_EMBEDDING_MODELS:
  375. modifier = collection_info.config.params.sparse_vectors[
  376. sparse_vector_field_name
  377. ].modifier
  378. assert (
  379. modifier == models.Modifier.IDF
  380. ), f"{self.sparse_embedding_model_name} requires modifier IDF, current modifier is {modifier}"
  381. def get_embedding_size(
  382. self,
  383. model_name: Optional[str] = None,
  384. ) -> int:
  385. """Get the size of the embeddings produced by the specified model.
  386. Args:
  387. model_name: optional, the name of the model to get the embedding size for. If None, the default model will
  388. be used.
  389. Returns:
  390. int: the size of the embeddings produced by the model.
  391. Raises:
  392. ValueError: If sparse model name is passed or model is not found in the supported models.
  393. """
  394. model_name = model_name or self.embedding_model_name
  395. embeddings_size, _ = self._get_model_params(model_name=model_name)
  396. return embeddings_size
  397. def get_fastembed_vector_params(
  398. self,
  399. on_disk: Optional[bool] = None,
  400. quantization_config: Optional[models.QuantizationConfig] = None,
  401. hnsw_config: Optional[models.HnswConfigDiff] = None,
  402. ) -> dict[str, models.VectorParams]:
  403. """
  404. Generates vector configuration, compatible with fastembed models.
  405. Args:
  406. on_disk: if True, vectors will be stored on disk. If None, default value will be used.
  407. quantization_config: Quantization configuration. If None, quantization will be disabled.
  408. hnsw_config: HNSW configuration. If None, default configuration will be used.
  409. Returns:
  410. Configuration for `vectors_config` argument in `create_collection` method.
  411. """
  412. vector_field_name = self.get_vector_field_name()
  413. embeddings_size, distance = self._get_model_params(model_name=self.embedding_model_name)
  414. return {
  415. vector_field_name: models.VectorParams(
  416. size=embeddings_size,
  417. distance=distance,
  418. on_disk=on_disk,
  419. quantization_config=quantization_config,
  420. hnsw_config=hnsw_config,
  421. )
  422. }
  423. def get_fastembed_sparse_vector_params(
  424. self,
  425. on_disk: Optional[bool] = None,
  426. modifier: Optional[models.Modifier] = None,
  427. ) -> Optional[dict[str, models.SparseVectorParams]]:
  428. """
  429. Generates vector configuration, compatible with fastembed sparse models.
  430. Args:
  431. on_disk: if True, vectors will be stored on disk. If None, default value will be used.
  432. modifier: Sparse vector queries modifier. E.g. Modifier.IDF for idf-based rescoring. Default: None.
  433. Returns:
  434. Configuration for `vectors_config` argument in `create_collection` method.
  435. """
  436. vector_field_name = self.get_sparse_vector_field_name()
  437. if self.sparse_embedding_model_name in IDF_EMBEDDING_MODELS:
  438. modifier = models.Modifier.IDF if modifier is None else modifier
  439. if vector_field_name is None:
  440. return None
  441. return {
  442. vector_field_name: models.SparseVectorParams(
  443. index=models.SparseIndexParams(
  444. on_disk=on_disk,
  445. ),
  446. modifier=modifier,
  447. )
  448. }
  449. def add(
  450. self,
  451. collection_name: str,
  452. documents: Iterable[str],
  453. metadata: Optional[Iterable[dict[str, Any]]] = None,
  454. ids: Optional[Iterable[models.ExtendedPointId]] = None,
  455. batch_size: int = 32,
  456. parallel: Optional[int] = None,
  457. **kwargs: Any,
  458. ) -> list[Union[str, int]]:
  459. """
  460. Adds text documents into qdrant collection.
  461. If collection does not exist, it will be created with default parameters.
  462. Metadata in combination with documents will be added as payload.
  463. Documents will be embedded using the specified embedding model.
  464. If you want to use your own vectors, use `upsert` method instead.
  465. Args:
  466. collection_name (str):
  467. Name of the collection to add documents to.
  468. documents (Iterable[str]):
  469. List of documents to embed and add to the collection.
  470. metadata (Iterable[dict[str, Any]], optional):
  471. List of metadata dicts. Defaults to None.
  472. ids (Iterable[models.ExtendedPointId], optional):
  473. List of ids to assign to documents.
  474. If not specified, UUIDs will be generated. Defaults to None.
  475. batch_size (int, optional):
  476. How many documents to embed and upload in single request. Defaults to 32.
  477. parallel (Optional[int], optional):
  478. How many parallel workers to use for embedding. Defaults to None.
  479. If number is specified, data-parallel process will be used.
  480. Raises:
  481. ImportError: If fastembed is not installed.
  482. Returns:
  483. List of IDs of added documents. If no ids provided, UUIDs will be randomly generated on client side.
  484. """
  485. # check if we have fastembed installed
  486. encoded_docs = self._embed_documents(
  487. documents=documents,
  488. embedding_model_name=self.embedding_model_name,
  489. batch_size=batch_size,
  490. embed_type="passage",
  491. parallel=parallel,
  492. )
  493. encoded_sparse_docs = None
  494. if self.sparse_embedding_model_name is not None:
  495. encoded_sparse_docs = self._sparse_embed_documents(
  496. documents=documents,
  497. embedding_model_name=self.sparse_embedding_model_name,
  498. batch_size=batch_size,
  499. parallel=parallel,
  500. )
  501. # Check if collection by same name exists, if not, create it
  502. try:
  503. collection_info = self.get_collection(collection_name=collection_name)
  504. except Exception:
  505. self.create_collection(
  506. collection_name=collection_name,
  507. vectors_config=self.get_fastembed_vector_params(),
  508. sparse_vectors_config=self.get_fastembed_sparse_vector_params(),
  509. )
  510. collection_info = self.get_collection(collection_name=collection_name)
  511. self._validate_collection_info(collection_info)
  512. inserted_ids: list = []
  513. points = self._points_iterator(
  514. ids=ids,
  515. metadata=metadata,
  516. encoded_docs=encoded_docs,
  517. ids_accumulator=inserted_ids,
  518. sparse_vectors=encoded_sparse_docs,
  519. )
  520. self.upload_points(
  521. collection_name=collection_name,
  522. points=points,
  523. wait=True,
  524. parallel=parallel or 1,
  525. batch_size=batch_size,
  526. **kwargs,
  527. )
  528. return inserted_ids
  529. def query(
  530. self,
  531. collection_name: str,
  532. query_text: str,
  533. query_filter: Optional[models.Filter] = None,
  534. limit: int = 10,
  535. **kwargs: Any,
  536. ) -> list[QueryResponse]:
  537. """
  538. Search for documents in a collection.
  539. This method automatically embeds the query text using the specified embedding model.
  540. If you want to use your own query vector, use `search` method instead.
  541. Args:
  542. collection_name: Collection to search in
  543. query_text:
  544. Text to search for. This text will be embedded using the specified embedding model.
  545. And then used as a query vector.
  546. query_filter:
  547. - Exclude vectors which doesn't fit given conditions.
  548. - If `None` - search among all vectors
  549. limit: How many results return
  550. **kwargs: Additional search parameters. See `qdrant_client.models.SearchRequest` for details.
  551. Returns:
  552. list[types.ScoredPoint]: List of scored points.
  553. """
  554. embedding_model_inst = self._get_or_init_model(
  555. model_name=self.embedding_model_name, deprecated=True
  556. )
  557. embeddings = list(embedding_model_inst.query_embed(query=query_text))
  558. query_vector = embeddings[0].tolist()
  559. if self.sparse_embedding_model_name is None:
  560. return self._scored_points_to_query_responses(
  561. self.search(
  562. collection_name=collection_name,
  563. query_vector=models.NamedVector(
  564. name=self.get_vector_field_name(), vector=query_vector
  565. ),
  566. query_filter=query_filter,
  567. limit=limit,
  568. with_payload=True,
  569. **kwargs,
  570. )
  571. )
  572. sparse_embedding_model_inst = self._get_or_init_sparse_model(
  573. model_name=self.sparse_embedding_model_name, deprecated=True
  574. )
  575. sparse_vector = list(sparse_embedding_model_inst.query_embed(query=query_text))[0]
  576. sparse_query_vector = models.SparseVector(
  577. indices=sparse_vector.indices.tolist(),
  578. values=sparse_vector.values.tolist(),
  579. )
  580. dense_request = models.SearchRequest(
  581. vector=models.NamedVector(
  582. name=self.get_vector_field_name(),
  583. vector=query_vector,
  584. ),
  585. filter=query_filter,
  586. limit=limit,
  587. with_payload=True,
  588. **kwargs,
  589. )
  590. sparse_request = models.SearchRequest(
  591. vector=models.NamedSparseVector(
  592. name=self.get_sparse_vector_field_name(),
  593. vector=sparse_query_vector,
  594. ),
  595. filter=query_filter,
  596. limit=limit,
  597. with_payload=True,
  598. **kwargs,
  599. )
  600. dense_request_response, sparse_request_response = self.search_batch(
  601. collection_name=collection_name, requests=[dense_request, sparse_request]
  602. )
  603. return self._scored_points_to_query_responses(
  604. reciprocal_rank_fusion([dense_request_response, sparse_request_response], limit=limit)
  605. )
  606. def query_batch(
  607. self,
  608. collection_name: str,
  609. query_texts: list[str],
  610. query_filter: Optional[models.Filter] = None,
  611. limit: int = 10,
  612. **kwargs: Any,
  613. ) -> list[list[QueryResponse]]:
  614. """
  615. Search for documents in a collection with batched query.
  616. This method automatically embeds the query text using the specified embedding model.
  617. Args:
  618. collection_name: Collection to search in
  619. query_texts:
  620. A list of texts to search for. Each text will be embedded using the specified embedding model.
  621. And then used as a query vector for a separate search requests.
  622. query_filter:
  623. - Exclude vectors which doesn't fit given conditions.
  624. - If `None` - search among all vectors
  625. This filter will be applied to all search requests.
  626. limit: How many results return
  627. **kwargs: Additional search parameters. See `qdrant_client.models.SearchRequest` for details.
  628. Returns:
  629. list[list[QueryResponse]]: List of lists of responses for each query text.
  630. """
  631. embedding_model_inst = self._get_or_init_model(
  632. model_name=self.embedding_model_name, deprecated=True
  633. )
  634. query_vectors = list(embedding_model_inst.query_embed(query=query_texts))
  635. requests = []
  636. for vector in query_vectors:
  637. request = models.SearchRequest(
  638. vector=models.NamedVector(
  639. name=self.get_vector_field_name(), vector=vector.tolist()
  640. ),
  641. filter=query_filter,
  642. limit=limit,
  643. with_payload=True,
  644. **kwargs,
  645. )
  646. requests.append(request)
  647. if self.sparse_embedding_model_name is None:
  648. responses = self.search_batch(
  649. collection_name=collection_name,
  650. requests=requests,
  651. )
  652. return [self._scored_points_to_query_responses(response) for response in responses]
  653. sparse_embedding_model_inst = self._get_or_init_sparse_model(
  654. model_name=self.sparse_embedding_model_name, deprecated=True
  655. )
  656. sparse_query_vectors = [
  657. models.SparseVector(
  658. indices=sparse_vector.indices.tolist(),
  659. values=sparse_vector.values.tolist(),
  660. )
  661. for sparse_vector in sparse_embedding_model_inst.embed(documents=query_texts)
  662. ]
  663. for sparse_vector in sparse_query_vectors:
  664. request = models.SearchRequest(
  665. vector=models.NamedSparseVector(
  666. name=self.get_sparse_vector_field_name(),
  667. vector=sparse_vector,
  668. ),
  669. filter=query_filter,
  670. limit=limit,
  671. with_payload=True,
  672. **kwargs,
  673. )
  674. requests.append(request)
  675. responses = self.search_batch(
  676. collection_name=collection_name,
  677. requests=requests,
  678. )
  679. dense_responses = responses[: len(query_texts)]
  680. sparse_responses = responses[len(query_texts) :]
  681. responses = [
  682. reciprocal_rank_fusion([dense_response, sparse_response], limit=limit)
  683. for dense_response, sparse_response in zip(dense_responses, sparse_responses)
  684. ]
  685. return [self._scored_points_to_query_responses(response) for response in responses]
  686. @classmethod
  687. def _resolve_query(
  688. cls,
  689. query: Union[
  690. types.PointId,
  691. list[float],
  692. list[list[float]],
  693. types.SparseVector,
  694. types.Query,
  695. types.NumpyArray,
  696. models.Document,
  697. models.Image,
  698. models.InferenceObject,
  699. None,
  700. ],
  701. ) -> Optional[models.Query]:
  702. """Resolves query interface into a models.Query object
  703. Args:
  704. query: models.QueryInterface - query as a model or a plain structure like list[float]
  705. Returns:
  706. Optional[models.Query]: query as it was, models.Query(nearest=query) or None
  707. Raises:
  708. ValueError: if query is not of supported type
  709. """
  710. if isinstance(query, get_args(types.Query)) or isinstance(query, grpc.Query):
  711. return query
  712. if isinstance(query, types.SparseVector):
  713. return models.NearestQuery(nearest=query)
  714. if isinstance(query, np.ndarray):
  715. return models.NearestQuery(nearest=query.tolist())
  716. if isinstance(query, list):
  717. return models.NearestQuery(nearest=query)
  718. if isinstance(query, get_args(types.PointId)):
  719. query = (
  720. GrpcToRest.convert_point_id(query) if isinstance(query, grpc.PointId) else query
  721. )
  722. return models.NearestQuery(nearest=query)
  723. if isinstance(query, get_args(INFERENCE_OBJECT_TYPES)):
  724. return models.NearestQuery(nearest=query)
  725. if query is None:
  726. return None
  727. raise ValueError(f"Unsupported query type: {type(query)}")
  728. def _resolve_query_request(self, query: models.QueryRequest) -> models.QueryRequest:
  729. """Resolve QueryRequest query field
  730. Args:
  731. query: models.QueryRequest - query request to resolve
  732. Returns:
  733. models.QueryRequest: A deepcopy of the query request with resolved query field
  734. """
  735. query = deepcopy(query)
  736. query.query = self._resolve_query(query.query)
  737. return query
  738. def _resolve_query_batch_request(
  739. self, requests: Sequence[models.QueryRequest]
  740. ) -> Sequence[models.QueryRequest]:
  741. """Resolve query field for each query request in a batch
  742. Args:
  743. requests: Sequence[models.QueryRequest] - query requests to resolve
  744. Returns:
  745. Sequence[models.QueryRequest]: A list of deep copied query requests with resolved query fields
  746. """
  747. return [self._resolve_query_request(query) for query in requests]
  748. def _embed_models(
  749. self,
  750. raw_models: Union[BaseModel, Iterable[BaseModel]],
  751. is_query: bool = False,
  752. batch_size: Optional[int] = None,
  753. ) -> Iterable[BaseModel]:
  754. FastEmbedMisc.import_fastembed()
  755. yield from self._model_embedder.embed_models(
  756. raw_models=raw_models,
  757. is_query=is_query,
  758. batch_size=batch_size or self.DEFAULT_BATCH_SIZE,
  759. )
  760. def _embed_models_strict(
  761. self,
  762. raw_models: Iterable[Union[dict[str, BaseModel], BaseModel]],
  763. batch_size: Optional[int] = None,
  764. parallel: Optional[int] = None,
  765. ) -> Iterable[BaseModel]:
  766. FastEmbedMisc.import_fastembed()
  767. yield from self._model_embedder.embed_models_strict(
  768. raw_models=raw_models,
  769. batch_size=batch_size or self.DEFAULT_BATCH_SIZE,
  770. parallel=parallel,
  771. )