From 95f7021beccd519f62d71d870a5798453229a8c3 Mon Sep 17 00:00:00 2001 From: Anush Date: Sat, 16 Nov 2024 02:52:23 +0530 Subject: [PATCH] feat: Qdrant vectorstore (#163) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Anush008 Co-authored-by: Michał Pstrąg --- .gitignore | 1 + docs/api_reference/core/vector-stores.md | 2 + examples/apps/documents_chat.py | 2 +- examples/document-search/chroma.py | 10 +- examples/document-search/chroma_otel.py | 2 +- examples/document-search/multimodal.py | 2 +- examples/document-search/qdrant.py | 107 ++++++++++ packages/ragbits-core/pyproject.toml | 3 + .../ragbits/core/metadata_stores/in_memory.py | 2 +- .../ragbits/core/vector_stores/__init__.py | 33 ++- .../src/ragbits/core/vector_stores/base.py | 18 +- .../src/ragbits/core/vector_stores/chroma.py | 55 ++--- .../ragbits/core/vector_stores/in_memory.py | 17 ++ .../src/ragbits/core/vector_stores/qdrant.py | 197 ++++++++++++++++++ .../tests/unit/vector_stores/test_chroma.py | 21 +- .../unit/vector_stores/test_in_memory.py | 9 +- .../tests/unit/vector_stores/test_qdrant.py | 140 +++++++++++++ .../src/ragbits/document_search/_main.py | 2 +- .../document_search/documents/element.py | 41 ++-- .../retrieval/rerankers/litellm.py | 2 +- .../tests/unit/test_elements.py | 14 +- pyproject.toml | 2 +- uv.lock | 134 +++++++++++- 23 files changed, 706 insertions(+), 110 deletions(-) create mode 100644 examples/document-search/qdrant.py create mode 100644 packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py create mode 100644 packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py diff --git a/.gitignore b/.gitignore index b35330b9..088ec593 100644 --- a/.gitignore +++ b/.gitignore @@ -95,5 +95,6 @@ dist/ # examples chroma/ +qdrant/ .aider* diff --git a/docs/api_reference/core/vector-stores.md b/docs/api_reference/core/vector-stores.md index 34254658..0503cd97 100644 --- a/docs/api_reference/core/vector-stores.md +++ b/docs/api_reference/core/vector-stores.md @@ -9,3 +9,5 @@ ::: ragbits.core.vector_stores.in_memory.InMemoryVectorStore ::: ragbits.core.vector_stores.chroma.ChromaVectorStore + +::: ragbits.core.vector_stores.qdrant.QdrantVectorStore diff --git a/examples/apps/documents_chat.py b/examples/apps/documents_chat.py index 3b1b2a43..8f7f24a4 100644 --- a/examples/apps/documents_chat.py +++ b/examples/apps/documents_chat.py @@ -125,7 +125,7 @@ async def _handle_message( if not self._documents_ingested: yield self.NO_DOCUMENTS_INGESTED_MESSAGE results = await self.document_search.search(message[-1]) - prompt = RAGPrompt(QueryWithContext(query=message, context=[i.get_text_representation() for i in results])) + prompt = RAGPrompt(QueryWithContext(query=message, context=[i.text_representation for i in results])) response = await self._llm.generate(prompt) yield response.answer diff --git a/examples/document-search/chroma.py b/examples/document-search/chroma.py index b641c360..bdf27dc5 100644 --- a/examples/document-search/chroma.py +++ b/examples/document-search/chroma.py @@ -9,7 +9,7 @@ 1. Create a list of documents. 2. Initialize the `LiteLLMEmbeddings` class with the OpenAI `text-embedding-3-small` embedding model. - 3. Initialize the `ChromaVectorStore` class with a `PersistentClient` instance and an index name. + 3. Initialize the `ChromaVectorStore` class with a `EphemeralClient` instance and an index name. 4. Initialize the `DocumentSearch` class with the embedder and the vector store. 5. Ingest the documents into the `DocumentSearch` instance. 6. List all documents in the vector store. @@ -33,7 +33,7 @@ import asyncio -from chromadb import PersistentClient +from chromadb import EphemeralClient from ragbits.core.embeddings.litellm import LiteLLMEmbeddings from ragbits.core.vector_stores.chroma import ChromaVectorStore @@ -72,7 +72,7 @@ async def main() -> None: model="text-embedding-3-small", ) vector_store = ChromaVectorStore( - client=PersistentClient("./chroma"), + client=EphemeralClient(), index_name="jokes", ) document_search = DocumentSearch( @@ -91,7 +91,7 @@ async def main() -> None: query = "I'm boiling my water and I need a joke" vector_store_kwargs = { "k": 2, - "max_distance": None, + "max_distance": 0.6, } results = await document_search.search( query, @@ -100,7 +100,7 @@ async def main() -> None: print() print(f"Documents similar to: {query}") - print([element.get_text_representation() for element in results]) + print([element.text_representation for element in results]) if __name__ == "__main__": diff --git a/examples/document-search/chroma_otel.py b/examples/document-search/chroma_otel.py index 64b5e198..7efd5195 100644 --- a/examples/document-search/chroma_otel.py +++ b/examples/document-search/chroma_otel.py @@ -130,7 +130,7 @@ async def main() -> None: print() print(f"Documents similar to: {query}") - print([element.get_text_representation() for element in results]) + print([element.text_representation for element in results]) if __name__ == "__main__": diff --git a/examples/document-search/multimodal.py b/examples/document-search/multimodal.py index 88dfb571..0e932bd5 100644 --- a/examples/document-search/multimodal.py +++ b/examples/document-search/multimodal.py @@ -91,7 +91,7 @@ async def main() -> None: print("Results for 'Fluffy teady bear toy':") for result in results: document = await result.document_meta.fetch() - print(f"Type: {result.element_type}, Location: {document.local_path}, Text: {result.get_text_representation()}") + print(f"Type: {result.element_type}, Location: {document.local_path}, Text: {result.text_representation}") if __name__ == "__main__": diff --git a/examples/document-search/qdrant.py b/examples/document-search/qdrant.py new file mode 100644 index 00000000..b1ea998e --- /dev/null +++ b/examples/document-search/qdrant.py @@ -0,0 +1,107 @@ +""" +Ragbits Document Search Example: Qdrant + +This example demonstrates how to use the `DocumentSearch` class to search for documents with a more advanced setup. +We will use the `LiteLLMEmbeddings` class to embed the documents and the query, the `QdrantVectorStore` class to store +the embeddings. + +The script performs the following steps: + + 1. Create a list of documents. + 2. Initialize the `LiteLLMEmbeddings` class with the OpenAI `text-embedding-3-small` embedding model. + 3. Initialize the `QdrantVectorStore` class with a `AsyncQdrantClient` in-memory instance and an index name. + 4. Initialize the `DocumentSearch` class with the embedder and the vector store. + 5. Ingest the documents into the `DocumentSearch` instance. + 6. List all documents in the vector store. + 7. Search for documents using a query. + 8. Print the list of all documents and the search results. + +To run the script, execute the following command: + + ```bash + uv run examples/document-search/qdrant.py + ``` +""" + +# /// script +# requires-python = ">=3.10" +# dependencies = [ +# "ragbits-document-search", +# "ragbits-core[litellm,qdrant]", +# ] +# /// + +import asyncio + +from qdrant_client import AsyncQdrantClient + +from ragbits.core.embeddings.litellm import LiteLLMEmbeddings +from ragbits.core.vector_stores.qdrant import QdrantVectorStore +from ragbits.document_search import DocumentSearch, SearchConfig +from ragbits.document_search.documents.document import DocumentMeta + +documents = [ + DocumentMeta.create_text_document_from_literal( + """ + RIP boiled water. You will be mist. + """ + ), + DocumentMeta.create_text_document_from_literal( + """ + Why doesn't James Bond fart in bed? Because it would blow his cover. + """ + ), + DocumentMeta.create_text_document_from_literal( + """ + Why programmers don't like to swim? Because they're scared of the floating points. + """ + ), + DocumentMeta.create_text_document_from_literal( + """ + This one is completely unrelated. + """ + ), +] + + +async def main() -> None: + """ + Run the example. + """ + embedder = LiteLLMEmbeddings( + model="text-embedding-3-small", + ) + vector_store = QdrantVectorStore( + client=AsyncQdrantClient(":memory:"), + index_name="jokes", + ) + document_search = DocumentSearch( + embedder=embedder, + vector_store=vector_store, + ) + + await document_search.ingest(documents) + + all_documents = await vector_store.list() + + print() + print("All documents:") + print([doc.metadata["content"] for doc in all_documents]) + + query = "I'm boiling my water and I need a joke" + vector_store_kwargs = { + "k": 2, + "max_distance": 0.6, + } + results = await document_search.search( + query, + config=SearchConfig(vector_store_kwargs=vector_store_kwargs), + ) + + print() + print(f"Documents similar to: {query}") + print([element.text_representation for element in results]) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/packages/ragbits-core/pyproject.toml b/packages/ragbits-core/pyproject.toml index 5ed2672f..ecfe10f1 100644 --- a/packages/ragbits-core/pyproject.toml +++ b/packages/ragbits-core/pyproject.toml @@ -59,6 +59,9 @@ promptfoo = [ otel = [ "opentelemetry-api~=1.27.0", ] +qdrant = [ + "qdrant-client~=1.12.1", +] [tool.uv] dev-dependencies = [ diff --git a/packages/ragbits-core/src/ragbits/core/metadata_stores/in_memory.py b/packages/ragbits-core/src/ragbits/core/metadata_stores/in_memory.py index 0cd1f89f..6f8a73db 100644 --- a/packages/ragbits-core/src/ragbits/core/metadata_stores/in_memory.py +++ b/packages/ragbits-core/src/ragbits/core/metadata_stores/in_memory.py @@ -23,7 +23,7 @@ async def store(self, ids: list[str], metadatas: list[dict]) -> None: ids: list of unique ids of the entries metadatas: list of dicts with metadata. """ - for _id, metadata in zip(ids, metadatas, strict=False): + for _id, metadata in zip(ids, metadatas, strict=True): self._storage[_id] = metadata @traceable diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/__init__.py b/packages/ragbits-core/src/ragbits/core/vector_stores/__init__.py index b9c16c34..7fb02751 100644 --- a/packages/ragbits-core/src/ragbits/core/vector_stores/__init__.py +++ b/packages/ragbits-core/src/ragbits/core/vector_stores/__init__.py @@ -1,34 +1,27 @@ import sys -from ..metadata_stores import get_metadata_store -from ..utils.config_handling import get_cls_from_config -from .base import VectorStore, VectorStoreEntry, VectorStoreOptions, WhereQuery -from .in_memory import InMemoryVectorStore +from ragbits.core.utils.config_handling import get_cls_from_config +from ragbits.core.vector_stores.base import VectorStore, VectorStoreEntry, VectorStoreOptions, WhereQuery +from ragbits.core.vector_stores.in_memory import InMemoryVectorStore -__all__ = ["InMemoryVectorStore", "VectorStore", "VectorStoreEntry", "WhereQuery"] +__all__ = ["InMemoryVectorStore", "VectorStore", "VectorStoreEntry", "VectorStoreOptions", "WhereQuery"] -module = sys.modules[__name__] - -def get_vector_store(vector_store_config: dict) -> VectorStore: +def get_vector_store(config: dict) -> VectorStore: """ Initializes and returns a VectorStore object based on the provided configuration. Args: - vector_store_config: A dictionary containing configuration details for the VectorStore. + config: A dictionary containing configuration details for the VectorStore. Returns: An instance of the specified VectorStore class, initialized with the provided config (if any) or default arguments. - """ - vector_store_cls = get_cls_from_config(vector_store_config["type"], module) - config = vector_store_config.get("config", {}) - if vector_store_config["type"].endswith("ChromaVectorStore"): - return vector_store_cls.from_config(config) - - metadata_store_config = vector_store_config.get("metadata_store_config") - return vector_store_cls( - default_options=VectorStoreOptions(**config.get("default_options", {})), - metadata_store=get_metadata_store(metadata_store_config), - ) + Raises: + KeyError: If the provided configuration does not contain a valid "type" key. + InvalidConfigurationError: If the provided configuration is invalid. + NotImplementedError: If the specified VectorStore class cannot be created from the provided configuration. + """ + vector_store_cls = get_cls_from_config(config["type"], sys.modules[__name__]) + return vector_store_cls.from_config(config.get("config", {})) diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/base.py b/packages/ragbits-core/src/ragbits/core/vector_stores/base.py index 234ca441..9ba342e4 100644 --- a/packages/ragbits-core/src/ragbits/core/vector_stores/base.py +++ b/packages/ragbits-core/src/ragbits/core/vector_stores/base.py @@ -13,8 +13,8 @@ class VectorStoreEntry(BaseModel): """ id: str + key: str vector: list[float] - content: str metadata: dict @@ -48,6 +48,22 @@ def __init__( self._default_options = default_options or VectorStoreOptions() self._metadata_store = metadata_store + @classmethod + def from_config(cls, config: dict) -> "VectorStore": + """ + Creates and returns an instance of the Reranker class from the given configuration. + + Args: + config: A dictionary containing the configuration for initializing the Reranker instance. + + Returns: + An initialized instance of the Reranker class. + + Raises: + NotImplementedError: If the class cannot be created from the provided configuration. + """ + raise NotImplementedError(f"Cannot create class {cls.__name__} from config.") + @abstractmethod async def store(self, entries: list[VectorStoreEntry]) -> None: """ diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py b/packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py index 00a147d7..e9f55c95 100644 --- a/packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py +++ b/packages/ragbits-core/src/ragbits/core/vector_stores/chroma.py @@ -1,10 +1,7 @@ -from __future__ import annotations - import json from typing import Literal import chromadb -from chromadb import Collection from chromadb.api import ClientAPI from ragbits.core.audit import traceable @@ -16,14 +13,14 @@ class ChromaVectorStore(VectorStore): """ - Class that stores text embeddings using [Chroma](https://docs.trychroma.com/). + Vector store implementation using [Chroma](https://docs.trychroma.com). """ def __init__( self, client: ClientAPI, index_name: str, - distance_method: Literal["l2", "ip", "cosine"] = "l2", + distance_method: Literal["l2", "ip", "cosine"] = "cosine", default_options: VectorStoreOptions | None = None, metadata_store: MetadataStore | None = None, ) -> None: @@ -41,22 +38,13 @@ def __init__( self._client = client self._index_name = index_name self._distance_method = distance_method - self._collection = self._get_chroma_collection() - - def _get_chroma_collection(self) -> Collection: - """ - Gets or creates a collection with the given name and metadata. - - Returns: - The collection. - """ - return self._client.get_or_create_collection( + self._collection = self._client.get_or_create_collection( name=self._index_name, metadata={"hnsw:space": self._distance_method}, ) @classmethod - def from_config(cls, config: dict) -> ChromaVectorStore: + def from_config(cls, config: dict) -> "ChromaVectorStore": """ Creates and returns an instance of the ChromaVectorStore class from the given configuration. @@ -70,7 +58,7 @@ def from_config(cls, config: dict) -> ChromaVectorStore: return cls( client=client_cls(**config["client"].get("config", {})), index_name=config["index_name"], - distance_method=config.get("distance_method", "l2"), + distance_method=config.get("distance_method", "cosine"), default_options=VectorStoreOptions(**config.get("default_options", {})), metadata_store=get_metadata_store(config.get("metadata_store")), ) @@ -83,15 +71,20 @@ async def store(self, entries: list[VectorStoreEntry]) -> None: Args: entries: The entries to store. """ + if not entries: + return + ids = [entry.id for entry in entries] - documents = [entry.content for entry in entries] + documents = [entry.key for entry in entries] embeddings = [entry.vector for entry in entries] metadatas = [entry.metadata for entry in entries] + metadatas = ( [{"__metadata": json.dumps(metadata, default=str)} for metadata in metadatas] if self._metadata_store is None else await self._metadata_store.store(ids, metadatas) # type: ignore ) + self._collection.add(ids=ids, embeddings=embeddings, metadatas=metadatas, documents=documents) # type: ignore @traceable @@ -116,14 +109,13 @@ async def retrieve(self, vector: list[float], options: VectorStoreOptions | None n_results=options.k, include=["metadatas", "embeddings", "distances", "documents"], ) + ids = results.get("ids") or [] - metadatas = results.get("metadatas") or [] embeddings = results.get("embeddings") or [] distances = results.get("distances") or [] documents = results.get("documents") or [] - metadatas = [ - [json.loads(metadata["__metadata"]) for batch in metadatas for metadata in batch] # type: ignore + [json.loads(metadata["__metadata"]) for batch in results.get("metadatas", []) for metadata in batch] # type: ignore if self._metadata_store is None else await self._metadata_store.get(*ids) ] @@ -131,12 +123,12 @@ async def retrieve(self, vector: list[float], options: VectorStoreOptions | None return [ VectorStoreEntry( id=id, - content=document, + key=document, vector=list(embeddings), metadata=metadata, # type: ignore ) - for batch in zip(ids, metadatas, embeddings, distances, documents, strict=False) - for id, metadata, embeddings, distance, document in zip(*batch, strict=False) + for batch in zip(ids, metadatas, embeddings, distances, documents, strict=True) + for id, metadata, embeddings, distance, document in zip(*batch, strict=True) if options.max_distance is None or distance <= options.max_distance ] @@ -162,19 +154,18 @@ async def list( # Cast `where` to chromadb's Where type where_chroma: chromadb.Where | None = dict(where) if where else None - get_results = self._collection.get( + results = self._collection.get( where=where_chroma, limit=limit, offset=offset, include=["metadatas", "embeddings", "documents"], ) - ids = get_results.get("ids") or [] - metadatas = get_results.get("metadatas") or [] - embeddings = get_results.get("embeddings") or [] - documents = get_results.get("documents") or [] + ids = results.get("ids") or [] + embeddings = results.get("embeddings") or [] + documents = results.get("documents") or [] metadatas = ( - [json.loads(metadata["__metadata"]) for metadata in metadatas] # type: ignore + [json.loads(metadata["__metadata"]) for metadata in results.get("metadatas", [])] # type: ignore if self._metadata_store is None else await self._metadata_store.get(ids) ) @@ -182,9 +173,9 @@ async def list( return [ VectorStoreEntry( id=id, - content=document, + key=document, vector=list(embedding), metadata=metadata, # type: ignore ) - for id, metadata, embedding, document in zip(ids, metadatas, embeddings, documents, strict=False) + for id, metadata, embedding, document in zip(ids, metadatas, embeddings, documents, strict=True) ] diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py b/packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py index d1a1c27e..28f69608 100644 --- a/packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py +++ b/packages/ragbits-core/src/ragbits/core/vector_stores/in_memory.py @@ -3,6 +3,7 @@ import numpy as np from ragbits.core.audit import traceable +from ragbits.core.metadata_stores import get_metadata_store from ragbits.core.metadata_stores.base import MetadataStore from ragbits.core.vector_stores.base import VectorStore, VectorStoreEntry, VectorStoreOptions, WhereQuery @@ -27,6 +28,22 @@ def __init__( super().__init__(default_options=default_options, metadata_store=metadata_store) self._storage: dict[str, VectorStoreEntry] = {} + @classmethod + def from_config(cls, config: dict) -> "InMemoryVectorStore": + """ + Creates and returns an instance of the InMemoryVectorStore class from the given configuration. + + Args: + config: A dictionary containing the configuration for initializing the InMemoryVectorStore instance. + + Returns: + An initialized instance of the InMemoryVectorStore class. + """ + return cls( + default_options=VectorStoreOptions(**config.get("default_options", {})), + metadata_store=get_metadata_store(config.get("metadata_store")), + ) + @traceable async def store(self, entries: list[VectorStoreEntry]) -> None: """ diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py b/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py new file mode 100644 index 00000000..d5bfeb4f --- /dev/null +++ b/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py @@ -0,0 +1,197 @@ +import json + +import qdrant_client +from qdrant_client import AsyncQdrantClient +from qdrant_client.models import Distance, Filter, VectorParams + +from ragbits.core.audit import traceable +from ragbits.core.metadata_stores import get_metadata_store +from ragbits.core.metadata_stores.base import MetadataStore +from ragbits.core.utils.config_handling import get_cls_from_config +from ragbits.core.vector_stores.base import VectorStore, VectorStoreEntry, VectorStoreOptions + + +class QdrantVectorStore(VectorStore): + """ + Vector store implementation using [Qdrant](https://qdrant.tech). + """ + + def __init__( + self, + client: AsyncQdrantClient, + index_name: str, + distance_method: Distance = Distance.COSINE, + default_options: VectorStoreOptions | None = None, + metadata_store: MetadataStore | None = None, + ) -> None: + """ + Constructs a new QdrantVectorStore instance. + + Args: + client: An instance of the Qdrant client. + index_name: The name of the index. + distance_method: The distance metric to use when creating the collection. + default_options: The default options for querying the vector store. + metadata_store: The metadata store to use. If None, the metadata will be stored in Qdrant. + """ + super().__init__(default_options=default_options, metadata_store=metadata_store) + self._client = client + self._index_name = index_name + self._distance_method = distance_method + + @classmethod + def from_config(cls, config: dict) -> "QdrantVectorStore": + """ + Creates and returns an instance of the QdrantVectorStore class from the given configuration. + + Args: + config: A dictionary containing the configuration for initializing the QdrantVectorStore instance. + + Returns: + An initialized instance of the QdrantVectorStore class. + """ + client_cls = get_cls_from_config(config["client"]["type"], qdrant_client) + return cls( + client=client_cls(**config["client"].get("config", {})), + index_name=config["index_name"], + distance_method=config.get("distance_method", Distance.COSINE), + default_options=VectorStoreOptions(**config.get("default_options", {})), + metadata_store=get_metadata_store(config.get("metadata_store")), + ) + + @traceable + async def store(self, entries: list[VectorStoreEntry]) -> None: + """ + Stores vector entries in the Qdrant collection. + + Args: + entries: List of VectorStoreEntry objects to store + + Raises: + QdrantException: If upload to collection fails. + """ + if not entries: + return + + if not await self._client.collection_exists(self._index_name): + await self._client.create_collection( + collection_name=self._index_name, + vectors_config=VectorParams(size=len(entries[0].vector), distance=self._distance_method), + ) + + ids = [entry.id for entry in entries] + embeddings = [entry.vector for entry in entries] + payloads = [{"document": entry.key} for entry in entries] + metadatas = [entry.metadata for entry in entries] + + metadatas = ( + [{"metadata": json.dumps(metadata, default=str)} for metadata in metadatas] + if self._metadata_store is None + else await self._metadata_store.store(ids, metadatas) # type: ignore + ) + if metadatas is not None: + payloads = [{**payload, **metadata} for (payload, metadata) in zip(payloads, metadatas, strict=True)] + + self._client.upload_collection( + collection_name=self._index_name, + vectors=embeddings, + payload=payloads, + ids=ids, + wait=True, + ) + + @traceable + async def retrieve(self, vector: list[float], options: VectorStoreOptions | None = None) -> list[VectorStoreEntry]: + """ + Retrieves entries from the Qdrant collection based on vector similarity. + + Args: + vector: The vector to query. + options: The options for querying the vector store. + + Returns: + The retrieved entries. + + Raises: + MetadataNotFoundError: If metadata cannot be retrieved + """ + options = options or self._default_options + score_threshold = 1 - options.max_distance if options.max_distance else None + + results = await self._client.query_points( + collection_name=self._index_name, + query=vector, + limit=options.k, + score_threshold=score_threshold, + with_payload=True, + with_vectors=True, + ) + + ids = [point.id for point in results.points] + vectors = [point.vector for point in results.points] + documents = [point.payload["document"] for point in results.points] # type: ignore + metadatas = ( + [json.loads(point.payload["metadata"]) for point in results.points] # type: ignore + if self._metadata_store is None + else await self._metadata_store.get(ids) # type: ignore + ) + + return [ + VectorStoreEntry( + id=str(id), + key=document, + vector=vector, # type: ignore + metadata=metadata, + ) + for id, document, vector, metadata in zip(ids, documents, vectors, metadatas, strict=True) + ] + + @traceable + async def list( # type: ignore + self, + where: Filter | None = None, # type: ignore + limit: int | None = None, + offset: int = 0, + ) -> list[VectorStoreEntry]: + """ + List entries from the vector store. The entries can be filtered, limited and offset. + + Args: + where: Conditions for filtering results. + Reference: https://qdrant.tech/documentation/concepts/filtering + limit: The maximum number of entries to return. + offset: The number of entries to skip. + + Returns: + The entries. + + Raises: + MetadataNotFoundError: If the metadata is not found. + """ + results = await self._client.query_points( + collection_name=self._index_name, + query_filter=where, + limit=limit or 10, + offset=offset, + with_payload=True, + with_vectors=True, + ) + + ids = [point.id for point in results.points] + vectors = [point.vector for point in results.points] + documents = [point.payload["document"] for point in results.points] # type: ignore + metadatas = ( + [json.loads(point.payload["metadata"]) for point in results.points] # type: ignore + if self._metadata_store is None + else await self._metadata_store.get(ids) # type: ignore + ) + + return [ + VectorStoreEntry( + id=str(id), + key=document, + vector=vector, # type: ignore + metadata=metadata, + ) + for id, document, vector, metadata in zip(ids, documents, vectors, metadatas, strict=True) + ] diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_chroma.py b/packages/ragbits-core/tests/unit/vector_stores/test_chroma.py index b3d038ec..f79a5126 100644 --- a/packages/ragbits-core/tests/unit/vector_stores/test_chroma.py +++ b/packages/ragbits-core/tests/unit/vector_stores/test_chroma.py @@ -14,16 +14,11 @@ def mock_chromadb_store() -> ChromaVectorStore: ) -async def test_get_chroma_collection(mock_chromadb_store: ChromaVectorStore) -> None: - _ = mock_chromadb_store._get_chroma_collection() - assert mock_chromadb_store._client.get_or_create_collection.call_count == 2 # type: ignore - - async def test_store(mock_chromadb_store: ChromaVectorStore) -> None: data = [ VectorStoreEntry( id="test_key", - content="test content", + key="test content", vector=[0.1, 0.2, 0.3], metadata={ "content": "test content", @@ -70,8 +65,7 @@ async def test_retrieve( mock_chromadb_store: ChromaVectorStore, max_distance: float | None, results: list[dict] ) -> None: vector = [0.1, 0.2, 0.3] - mock_collection = mock_chromadb_store._get_chroma_collection() - mock_collection.query.return_value = { # type: ignore + mock_chromadb_store._collection.query.return_value = { # type: ignore "metadatas": [ [ { @@ -93,17 +87,16 @@ async def test_retrieve( entries = await mock_chromadb_store.retrieve(vector, options=VectorStoreOptions(max_distance=max_distance)) assert len(entries) == len(results) - for entry, result in zip(entries, results, strict=False): + for entry, result in zip(entries, results, strict=True): assert entry.metadata["content"] == result["content"] assert entry.metadata["document"]["title"] == result["title"] assert entry.vector == result["vector"] assert entry.id == f"test_id_{results.index(result) + 1}" - assert entry.content == result["content"] + assert entry.key == result["content"] async def test_list(mock_chromadb_store: ChromaVectorStore) -> None: - mock_collection = mock_chromadb_store._get_chroma_collection() - mock_collection.get.return_value = { # type: ignore + mock_chromadb_store._collection.get.return_value = { # type: ignore "metadatas": [ { "__metadata": '{"content": "test content", "document": {"title": "test title", "source":' @@ -125,10 +118,10 @@ async def test_list(mock_chromadb_store: ChromaVectorStore) -> None: assert entries[0].metadata["content"] == "test content" assert entries[0].metadata["document"]["title"] == "test title" assert entries[0].vector == [0.12, 0.25, 0.29] - assert entries[0].content == "test content 1" + assert entries[0].key == "test content 1" assert entries[0].id == "test_id_1" assert entries[1].metadata["content"] == "test content 2" assert entries[1].metadata["document"]["title"] == "test title 2" assert entries[1].vector == [0.13, 0.26, 0.30] - assert entries[1].content == "test content2" + assert entries[1].key == "test content2" assert entries[1].id == "test_id_2" diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py b/packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py index d47e19a2..2c167330 100644 --- a/packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py +++ b/packages/ragbits-core/tests/unit/vector_stores/test_in_memory.py @@ -1,6 +1,7 @@ from pathlib import Path import pytest +from pydantic import computed_field from ragbits.core.vector_stores.base import VectorStoreOptions from ragbits.core.vector_stores.in_memory import InMemoryVectorStore @@ -20,12 +21,14 @@ class AnimalElement(Element): type: str age: int - def get_text_representation(self) -> str: + @computed_field # type: ignore[prop-decorator] + @property + def text_representation(self) -> str: """ Get the text representation of the element. Returns: - The key. + The text representation. """ return self.name @@ -66,7 +69,7 @@ async def test_retrieve(store: InMemoryVectorStore, k: int, max_distance: float entries = await store.retrieve(search_vector, options=VectorStoreOptions(k=k, max_distance=max_distance)) assert len(entries) == len(results) - for entry, result in zip(entries, results, strict=False): + for entry, result in zip(entries, results, strict=True): assert entry.metadata["name"] == result diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py b/packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py new file mode 100644 index 00000000..22cb9391 --- /dev/null +++ b/packages/ragbits-core/tests/unit/vector_stores/test_qdrant.py @@ -0,0 +1,140 @@ +from unittest.mock import AsyncMock + +import pytest +from qdrant_client.http import models + +from ragbits.core.vector_stores.base import VectorStoreEntry +from ragbits.core.vector_stores.qdrant import QdrantVectorStore + + +@pytest.fixture +def mock_qdrant_store() -> QdrantVectorStore: + return QdrantVectorStore( + client=AsyncMock(), + index_name="test_collection", + ) + + +async def test_store(mock_qdrant_store: QdrantVectorStore) -> None: + data = [ + VectorStoreEntry( + id="1c7d6b27-4ef1-537c-ad7c-676edb8bc8a8", + key="test_key", + vector=[0.1, 0.2, 0.3], + metadata={ + "content": "test content", + "document": { + "title": "test title", + "source": {"path": "/test/path"}, + "document_type": "test_type", + }, + }, + ) + ] + + mock_qdrant_store._client.collection_exists.return_value = False # type: ignore + await mock_qdrant_store.store(data) + + mock_qdrant_store._client.collection_exists.assert_called_once() # type: ignore + mock_qdrant_store._client.create_collection.assert_called_once() # type: ignore + mock_qdrant_store._client.upload_collection.assert_called_with( # type: ignore + collection_name="test_collection", + vectors=[[0.1, 0.2, 0.3]], + payload=[ + { + "document": "test_key", + "metadata": '{"content": "test content", ' + '"document": {"title": "test title", "source": {"path": "/test/path"}, "document_type": "test_type"}}', + } + ], + ids=["1c7d6b27-4ef1-537c-ad7c-676edb8bc8a8"], + wait=True, + ) + + +async def test_retrieve(mock_qdrant_store: QdrantVectorStore) -> None: + mock_qdrant_store._client.query_points.return_value = models.QueryResponse( # type: ignore + points=[ + models.ScoredPoint( + version=1, + id="1f908deb-bc9f-4b5a-8b73-2e72d8b44dc5", + vector=[0.12, 0.25, 0.29], + score=0.9, + payload={ + "document": "test_key 1", + "metadata": '{"content": "test content 1",' + '"document": {"title": "test title 1", ' + '"source": {"path": "/test/path-1"}, "document_type": "txt"}}', + }, + ), + models.ScoredPoint( + version=1, + id="827cad0b-058f-4b85-b8ed-ac741948d502", + vector=[0.13, 0.26, 0.30], + score=0.9, + payload={ + "document": "test_key 2", + "metadata": '{"content": "test content 2", ' + '"document": {"title": "test title 2", ' + '"source": {"path": "/test/path-2"}, "document_type": "txt"}}', + }, + ), + ] + ) + + results = [ + {"content": "test content 1", "title": "test title 1", "vector": [0.12, 0.25, 0.29]}, + {"content": "test content 2", "title": "test title 2", "vector": [0.13, 0.26, 0.30]}, + ] + + entries = await mock_qdrant_store.retrieve([0.12, 0.25, 0.29]) + + assert len(entries) == len(results) + for entry, result in zip(entries, results, strict=True): + assert entry.metadata["content"] == result["content"] + assert entry.metadata["document"]["title"] == result["title"] + assert entry.vector == result["vector"] + + +async def test_list(mock_qdrant_store: QdrantVectorStore) -> None: + mock_qdrant_store._client.query_points.return_value = models.QueryResponse( # type: ignore + points=[ + models.ScoredPoint( + version=1, + id="1f908deb-bc9f-4b5a-8b73-2e72d8b44dc5", + vector=[0.12, 0.25, 0.29], + score=0.9, + payload={ + "document": "test_key 1", + "metadata": '{"content": "test content 1",' + '"document": {"title": "test title 1", ' + '"source": {"path": "/test/path-1"}, "document_type": "txt"}}', + }, + ), + models.ScoredPoint( + version=1, + id="827cad0b-058f-4b85-b8ed-ac741948d502", + vector=[0.13, 0.26, 0.30], + score=0.9, + payload={ + "document": "test_key 2", + "metadata": '{"content": "test content 2", ' + '"document": {"title": "test title 2", ' + '"source": {"path": "/test/path-2"}, "document_type": "txt"}}', + }, + ), + ] + ) + + results = [ + {"content": "test content 1", "title": "test title 1", "vector": [0.12, 0.25, 0.29]}, + {"content": "test content 2", "title": "test title 2", "vector": [0.13, 0.26, 0.30]}, + ] + + entries = await mock_qdrant_store.list() + + assert len(entries) == len(results) + for entry, result in zip(entries, results, strict=True): + assert entry.metadata["content"] == result["content"] + assert entry.metadata["document"]["title"] == result["title"] + assert entry.vector == result["vector"] diff --git a/packages/ragbits-document-search/src/ragbits/document_search/_main.py b/packages/ragbits-document-search/src/ragbits/document_search/_main.py index 521ea0ff..daa56dc0 100644 --- a/packages/ragbits-document-search/src/ragbits/document_search/_main.py +++ b/packages/ragbits-document-search/src/ragbits/document_search/_main.py @@ -170,7 +170,7 @@ async def insert_elements(self, elements: list[Element]) -> None: Args: elements: The list of Elements to insert. """ - vectors = await self.embedder.embed_text([element.get_text_for_embedding() for element in elements]) + vectors = await self.embedder.embed_text([element.key for element in elements]) image_elements = [element for element in elements if isinstance(element, ImageElement)] entries = [element.to_vector_db_entry(vector) for element, vector in zip(elements, vectors, strict=False)] diff --git a/packages/ragbits-document-search/src/ragbits/document_search/documents/element.py b/packages/ragbits-document-search/src/ragbits/document_search/documents/element.py index 893d7d3c..63bf82bd 100644 --- a/packages/ragbits-document-search/src/ragbits/document_search/documents/element.py +++ b/packages/ragbits-document-search/src/ragbits/document_search/documents/element.py @@ -42,24 +42,27 @@ def id(self) -> str: id_components = [ self.document_meta.id, self.element_type, - self.get_text_for_embedding(), - self.get_text_representation(), + self.key, + self.text_representation, str(self.location), ] - return str(uuid.uuid5(uuid.NAMESPACE_OID, ";".join(id_components))) - def get_text_for_embedding(self) -> str: + @computed_field # type: ignore[prop-decorator] + @property + def key(self) -> str: """ - Get the text representation of the element for embedding. + Get the representation of the element for embedding. Returns: - The text representation for embedding. + The representation for embedding. """ - return self.get_text_representation() + return self.text_representation + @computed_field # type: ignore[prop-decorator] + @property @abstractmethod - def get_text_representation(self) -> str: + def text_representation(self) -> str: """ Get the text representation of the element. @@ -68,12 +71,10 @@ def get_text_representation(self) -> str: """ @classmethod - def __pydantic_init_subclass__(cls, **kwargs: Any) -> None: # pylint: disable=unused-argument #noqa: ANN401 + def __pydantic_init_subclass__(cls, **kwargs: Any) -> None: # noqa: ANN401 element_type_default = cls.model_fields["element_type"].default - if element_type_default is None: raise ValueError("Element type must be defined") - Element._elements_registry[element_type_default] = cls @classmethod @@ -87,11 +88,9 @@ def from_vector_db_entry(cls, db_entry: VectorStoreEntry) -> "Element": Returns: The element. """ - meta = db_entry.metadata - element_type = meta["element_type"] + element_type = db_entry.metadata["element_type"] element_cls = Element._elements_registry[element_type] - - return element_cls(**meta) + return element_cls(**db_entry.metadata) def to_vector_db_entry(self, vector: list[float]) -> VectorStoreEntry: """ @@ -105,9 +104,9 @@ def to_vector_db_entry(self, vector: list[float]) -> VectorStoreEntry: """ return VectorStoreEntry( id=self.id, + key=self.key, vector=vector, - content=self.get_text_for_embedding(), - metadata=self.model_dump(), + metadata=self.model_dump(exclude={"id", "key"}), ) @@ -119,7 +118,9 @@ class TextElement(Element): element_type: str = "text" content: str - def get_text_representation(self) -> str: + @computed_field # type: ignore[prop-decorator] + @property + def text_representation(self) -> str: """ Get the text representation of the element. @@ -139,7 +140,9 @@ class ImageElement(Element): ocr_extracted_text: str image_bytes: bytes - def get_text_representation(self) -> str: + @computed_field # type: ignore[prop-decorator] + @property + def text_representation(self) -> str: """ Get the text representation of the element. diff --git a/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/litellm.py b/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/litellm.py index c83906ce..a27e982f 100644 --- a/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/litellm.py +++ b/packages/ragbits-document-search/src/ragbits/document_search/retrieval/rerankers/litellm.py @@ -58,7 +58,7 @@ async def rerank( The reranked elements. """ options = self._default_options if options is None else options - documents = [element.get_text_representation() for element in elements] + documents = [element.text_representation for element in elements] response = await litellm.arerank( model=self.model, diff --git a/packages/ragbits-document-search/tests/unit/test_elements.py b/packages/ragbits-document-search/tests/unit/test_elements.py index 7f98307c..f4b1fa78 100644 --- a/packages/ragbits-document-search/tests/unit/test_elements.py +++ b/packages/ragbits-document-search/tests/unit/test_elements.py @@ -1,20 +1,24 @@ +from pydantic import computed_field + from ragbits.core.vector_stores.base import VectorStoreEntry from ragbits.document_search.documents.document import DocumentType from ragbits.document_search.documents.element import Element -def test_resolving_element_type(): +def test_resolving_element_type() -> None: class MyElement(Element): element_type: str = "custom_element" foo: str - def get_text_representation(self) -> str: + @computed_field # type: ignore[prop-decorator] + @property + def text_representation(self) -> str: return self.foo + self.foo element = Element.from_vector_db_entry( db_entry=VectorStoreEntry( id="test id", - content="test content", + key="test content", vector=[0.1, 0.2], metadata={ "element_type": "custom_element", @@ -29,7 +33,7 @@ def get_text_representation(self) -> str: assert isinstance(element, MyElement) assert element.foo == "bar" - assert element.get_text_for_embedding() == "barbar" - assert element.get_text_representation() == "barbar" + assert element.key == "barbar" + assert element.text_representation == "barbar" assert element.document_meta.document_type == DocumentType.TXT assert element.document_meta.source.source_type == "local_file_source" diff --git a/pyproject.toml b/pyproject.toml index b56cc198..0ee989d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -6,7 +6,7 @@ readme = "README.md" requires-python = ">=3.10" dependencies = [ "ragbits-cli", - "ragbits-core[chroma,lab,litellm,local,otel]", + "ragbits-core[chroma,lab,litellm,local,otel,qdrant]", "ragbits-document-search[gcs,huggingface]", "ragbits-evaluate[relari]", "ragbits-guardrails[openai]", diff --git a/uv.lock b/uv.lock index a3b1b1b8..c89e5221 100644 --- a/uv.lock +++ b/uv.lock @@ -1368,6 +1368,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a2/df/133216989fe7e17caeafd7ff5b17cc82c4e722025d0b8d5d2290c11fe2e6/grpcio-1.66.2-cp313-cp313-win_amd64.whl", hash = "sha256:fb70487c95786e345af5e854ffec8cb8cc781bcc5df7930c4fbb7feaa72e1cdf", size = 4278018 }, ] +[[package]] +name = "grpcio-tools" +version = "1.62.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "grpcio" }, + { name = "protobuf" }, + { name = "setuptools" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/54/fa/b69bd8040eafc09b88bb0ec0fea59e8aacd1a801e688af087cead213b0d0/grpcio-tools-1.62.3.tar.gz", hash = "sha256:7c7136015c3d62c3eef493efabaf9e3380e3e66d24ee8e94c01cb71377f57833", size = 4538520 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/eb/eb0a3aa9480c3689d31fd2ad536df6a828e97a60f667c8a93d05bdf07150/grpcio_tools-1.62.3-cp310-cp310-macosx_12_0_universal2.whl", hash = "sha256:2f968b049c2849540751ec2100ab05e8086c24bead769ca734fdab58698408c1", size = 5117556 }, + { url = "https://files.pythonhosted.org/packages/f3/fb/8be3dda485f7fab906bfa02db321c3ecef953a87cdb5f6572ca08b187bcb/grpcio_tools-1.62.3-cp310-cp310-manylinux_2_17_aarch64.whl", hash = "sha256:0a8c0c4724ae9c2181b7dbc9b186df46e4f62cb18dc184e46d06c0ebeccf569e", size = 2719330 }, + { url = "https://files.pythonhosted.org/packages/63/de/6978f8d10066e240141cd63d1fbfc92818d96bb53427074f47a8eda921e1/grpcio_tools-1.62.3-cp310-cp310-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5782883a27d3fae8c425b29a9d3dcf5f47d992848a1b76970da3b5a28d424b26", size = 3070818 }, + { url = "https://files.pythonhosted.org/packages/74/34/bb8f816893fc73fd6d830e895e8638d65d13642bb7a434f9175c5ca7da11/grpcio_tools-1.62.3-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f3d812daffd0c2d2794756bd45a353f89e55dc8f91eb2fc840c51b9f6be62667", size = 2804993 }, + { url = "https://files.pythonhosted.org/packages/78/60/b2198d7db83293cdb9760fc083f077c73e4c182da06433b3b157a1567d06/grpcio_tools-1.62.3-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:b47d0dda1bdb0a0ba7a9a6de88e5a1ed61f07fad613964879954961e36d49193", size = 3684915 }, + { url = "https://files.pythonhosted.org/packages/61/20/56dbdc4ecb14d42a03cd164ff45e6e84572bbe61ee59c50c39f4d556a8d5/grpcio_tools-1.62.3-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:ca246dffeca0498be9b4e1ee169b62e64694b0f92e6d0be2573e65522f39eea9", size = 3297482 }, + { url = "https://files.pythonhosted.org/packages/4a/dc/e417a313c905744ce8cedf1e1edd81c41dc45ff400ae1c45080e18f26712/grpcio_tools-1.62.3-cp310-cp310-win32.whl", hash = "sha256:6a56d344b0bab30bf342a67e33d386b0b3c4e65868ffe93c341c51e1a8853ca5", size = 909793 }, + { url = "https://files.pythonhosted.org/packages/d9/69/75e7ebfd8d755d3e7be5c6d1aa6d13220f5bba3a98965e4b50c329046777/grpcio_tools-1.62.3-cp310-cp310-win_amd64.whl", hash = "sha256:710fecf6a171dcbfa263a0a3e7070e0df65ba73158d4c539cec50978f11dad5d", size = 1052459 }, + { url = "https://files.pythonhosted.org/packages/23/52/2dfe0a46b63f5ebcd976570aa5fc62f793d5a8b169e211c6a5aede72b7ae/grpcio_tools-1.62.3-cp311-cp311-macosx_10_10_universal2.whl", hash = "sha256:703f46e0012af83a36082b5f30341113474ed0d91e36640da713355cd0ea5d23", size = 5147623 }, + { url = "https://files.pythonhosted.org/packages/f0/2e/29fdc6c034e058482e054b4a3c2432f84ff2e2765c1342d4f0aa8a5c5b9a/grpcio_tools-1.62.3-cp311-cp311-manylinux_2_17_aarch64.whl", hash = "sha256:7cc83023acd8bc72cf74c2edbe85b52098501d5b74d8377bfa06f3e929803492", size = 2719538 }, + { url = "https://files.pythonhosted.org/packages/f9/60/abe5deba32d9ec2c76cdf1a2f34e404c50787074a2fee6169568986273f1/grpcio_tools-1.62.3-cp311-cp311-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7ff7d58a45b75df67d25f8f144936a3e44aabd91afec833ee06826bd02b7fbe7", size = 3070964 }, + { url = "https://files.pythonhosted.org/packages/bc/ad/e2b066684c75f8d9a48508cde080a3a36618064b9cadac16d019ca511444/grpcio_tools-1.62.3-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7f2483ea232bd72d98a6dc6d7aefd97e5bc80b15cd909b9e356d6f3e326b6e43", size = 2805003 }, + { url = "https://files.pythonhosted.org/packages/9c/3f/59bf7af786eae3f9d24ee05ce75318b87f541d0950190ecb5ffb776a1a58/grpcio_tools-1.62.3-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:962c84b4da0f3b14b3cdb10bc3837ebc5f136b67d919aea8d7bb3fd3df39528a", size = 3685154 }, + { url = "https://files.pythonhosted.org/packages/f1/79/4dd62478b91e27084c67b35a2316ce8a967bd8b6cb8d6ed6c86c3a0df7cb/grpcio_tools-1.62.3-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:8ad0473af5544f89fc5a1ece8676dd03bdf160fb3230f967e05d0f4bf89620e3", size = 3297942 }, + { url = "https://files.pythonhosted.org/packages/b8/cb/86449ecc58bea056b52c0b891f26977afc8c4464d88c738f9648da941a75/grpcio_tools-1.62.3-cp311-cp311-win32.whl", hash = "sha256:db3bc9fa39afc5e4e2767da4459df82b095ef0cab2f257707be06c44a1c2c3e5", size = 910231 }, + { url = "https://files.pythonhosted.org/packages/45/a4/9736215e3945c30ab6843280b0c6e1bff502910156ea2414cd77fbf1738c/grpcio_tools-1.62.3-cp311-cp311-win_amd64.whl", hash = "sha256:e0898d412a434e768a0c7e365acabe13ff1558b767e400936e26b5b6ed1ee51f", size = 1052496 }, + { url = "https://files.pythonhosted.org/packages/2a/a5/d6887eba415ce318ae5005e8dfac3fa74892400b54b6d37b79e8b4f14f5e/grpcio_tools-1.62.3-cp312-cp312-macosx_10_10_universal2.whl", hash = "sha256:d102b9b21c4e1e40af9a2ab3c6d41afba6bd29c0aa50ca013bf85c99cdc44ac5", size = 5147690 }, + { url = "https://files.pythonhosted.org/packages/8a/7c/3cde447a045e83ceb4b570af8afe67ffc86896a2fe7f59594dc8e5d0a645/grpcio_tools-1.62.3-cp312-cp312-manylinux_2_17_aarch64.whl", hash = "sha256:0a52cc9444df978438b8d2332c0ca99000521895229934a59f94f37ed896b133", size = 2720538 }, + { url = "https://files.pythonhosted.org/packages/88/07/f83f2750d44ac4f06c07c37395b9c1383ef5c994745f73c6bfaf767f0944/grpcio_tools-1.62.3-cp312-cp312-manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:141d028bf5762d4a97f981c501da873589df3f7e02f4c1260e1921e565b376fa", size = 3071571 }, + { url = "https://files.pythonhosted.org/packages/37/74/40175897deb61e54aca716bc2e8919155b48f33aafec8043dda9592d8768/grpcio_tools-1.62.3-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:47a5c093ab256dec5714a7a345f8cc89315cb57c298b276fa244f37a0ba507f0", size = 2806207 }, + { url = "https://files.pythonhosted.org/packages/ec/ee/d8de915105a217cbcb9084d684abdc032030dcd887277f2ef167372287fe/grpcio_tools-1.62.3-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:f6831fdec2b853c9daa3358535c55eed3694325889aa714070528cf8f92d7d6d", size = 3685815 }, + { url = "https://files.pythonhosted.org/packages/fd/d9/4360a6c12be3d7521b0b8c39e5d3801d622fbb81cc2721dbd3eee31e28c8/grpcio_tools-1.62.3-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:e02d7c1a02e3814c94ba0cfe43d93e872c758bd8fd5c2797f894d0c49b4a1dfc", size = 3298378 }, + { url = "https://files.pythonhosted.org/packages/29/3b/7cdf4a9e5a3e0a35a528b48b111355cd14da601413a4f887aa99b6da468f/grpcio_tools-1.62.3-cp312-cp312-win32.whl", hash = "sha256:b881fd9505a84457e9f7e99362eeedd86497b659030cf57c6f0070df6d9c2b9b", size = 910416 }, + { url = "https://files.pythonhosted.org/packages/6c/66/dd3ec249e44c1cc15e902e783747819ed41ead1336fcba72bf841f72c6e9/grpcio_tools-1.62.3-cp312-cp312-win_amd64.whl", hash = "sha256:11c625eebefd1fd40a228fc8bae385e448c7e32a6ae134e43cf13bbc23f902b7", size = 1052856 }, +] + [[package]] name = "h11" version = "0.14.0" @@ -1377,6 +1414,28 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/95/04/ff642e65ad6b90db43e668d70ffb6736436c7ce41fcc549f4e9472234127/h11-0.14.0-py3-none-any.whl", hash = "sha256:e3fe4ac4b851c468cc8363d500db52c2ead036020723024a109d37346efaa761", size = 58259 }, ] +[[package]] +name = "h2" +version = "4.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "hpack" }, + { name = "hyperframe" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/2a/32/fec683ddd10629ea4ea46d206752a95a2d8a48c22521edd70b142488efe1/h2-4.1.0.tar.gz", hash = "sha256:a83aca08fbe7aacb79fec788c9c0bac936343560ed9ec18b82a13a12c28d2abb", size = 2145593 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/2a/e5/db6d438da759efbb488c4f3fbdab7764492ff3c3f953132efa6b9f0e9e53/h2-4.1.0-py3-none-any.whl", hash = "sha256:03a46bcf682256c95b5fd9e9a99c1323584c3eec6440d379b9903d709476bc6d", size = 57488 }, +] + +[[package]] +name = "hpack" +version = "4.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/9b/fda93fb4d957db19b0f6b370e79d586b3e8528b20252c729c476a2c02954/hpack-4.0.0.tar.gz", hash = "sha256:fc41de0c63e687ebffde81187a948221294896f6bdc0ae2312708df339430095", size = 49117 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d5/34/e8b383f35b77c402d28563d2b8f83159319b509bc5f760b15d60b0abf165/hpack-4.0.0-py3-none-any.whl", hash = "sha256:84a076fad3dc9a9f8063ccb8041ef100867b1878b25ef0ee63847a5d53818a6c", size = 32611 }, +] + [[package]] name = "httpcore" version = "1.0.6" @@ -1435,6 +1494,11 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/56/95/9377bcb415797e44274b51d46e3249eba641711cf3348050f76ee7b15ffc/httpx-0.27.2-py3-none-any.whl", hash = "sha256:7bb2708e112d8fdd7829cd4243970f0c223274051cb35ee80c03301ee29a3df0", size = 76395 }, ] +[package.optional-dependencies] +http2 = [ + { name = "h2" }, +] + [[package]] name = "huggingface-hub" version = "0.25.1" @@ -1479,6 +1543,15 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c6/50/e0edd38dcd63fb26a8547f13d28f7a008bc4a3fd4eb4ff030673f22ad41a/hydra_core-1.3.2-py3-none-any.whl", hash = "sha256:fa0238a9e31df3373b35b0bfb672c34cc92718d21f81311d8996a16de1141d8b", size = 154547 }, ] +[[package]] +name = "hyperframe" +version = "6.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5a/2a/4747bff0a17f7281abe73e955d60d80aae537a5d203f417fa1c2e7578ebb/hyperframe-6.0.1.tar.gz", hash = "sha256:ae510046231dc8e9ecb1a6586f63d2347bf4c8905914aa84ba585ae85f28a914", size = 25008 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d7/de/85a784bcc4a3779d1753a7ec2dee5de90e18c7bcf402e71b51fcf150b129/hyperframe-6.0.1-py3-none-any.whl", hash = "sha256:0ec6bafd80d8ad2195c4f03aacba3a8265e57bc4cff261e802bf39970ed02a15", size = 12389 }, +] + [[package]] name = "identify" version = "2.6.1" @@ -2416,6 +2489,8 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/bc/f7/7ec7fddc92e50714ea3745631f79bd9c96424cb2702632521028e57d3a36/multiprocess-0.70.16-py310-none-any.whl", hash = "sha256:c4a9944c67bd49f823687463660a2d6daae94c289adff97e0f9d696ba6371d02", size = 134824 }, { url = "https://files.pythonhosted.org/packages/50/15/b56e50e8debaf439f44befec5b2af11db85f6e0f344c3113ae0be0593a91/multiprocess-0.70.16-py311-none-any.whl", hash = "sha256:af4cabb0dac72abfb1e794fa7855c325fd2b55a10a44628a3c1ad3311c04127a", size = 143519 }, { url = "https://files.pythonhosted.org/packages/0a/7d/a988f258104dcd2ccf1ed40fdc97e26c4ac351eeaf81d76e266c52d84e2f/multiprocess-0.70.16-py312-none-any.whl", hash = "sha256:fc0544c531920dde3b00c29863377f87e1632601092ea2daca74e4beb40faa2e", size = 146741 }, + { url = "https://files.pythonhosted.org/packages/ea/89/38df130f2c799090c978b366cfdf5b96d08de5b29a4a293df7f7429fa50b/multiprocess-0.70.16-py38-none-any.whl", hash = "sha256:a71d82033454891091a226dfc319d0cfa8019a4e888ef9ca910372a446de4435", size = 132628 }, + { url = "https://files.pythonhosted.org/packages/da/d9/f7f9379981e39b8c2511c9e0326d212accacb82f12fbfdc1aa2ce2a7b2b6/multiprocess-0.70.16-py39-none-any.whl", hash = "sha256:a0bafd3ae1b732eac64be2e72038231c1ba97724b60b09400d68f229fcc2fbf3", size = 133351 }, ] [[package]] @@ -3147,6 +3222,18 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/88/5f/e351af9a41f866ac3f1fac4ca0613908d9a41741cfcf2228f4ad853b697d/pluggy-1.5.0-py3-none-any.whl", hash = "sha256:44e1ad92c8ca002de6377e165f3e0f1be63266ab4d554740532335b9d75ea669", size = 20556 }, ] +[[package]] +name = "portalocker" +version = "2.10.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pywin32", marker = "platform_system == 'Windows'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ed/d3/c6c64067759e87af98cc668c1cc75171347d0f1577fab7ca3749134e3cd4/portalocker-2.10.1.tar.gz", hash = "sha256:ef1bf844e878ab08aee7e40184156e1151f228f103aa5c6bd0724cc330960f8f", size = 40891 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9b/fb/a70a4214956182e0d7a9099ab17d50bfcba1056188e9b14f35b9e2b62a0d/portalocker-2.10.1-py3-none-any.whl", hash = "sha256:53a5984ebc86a025552264b459b46a2086e269b21823cb572f8f28ee759e45bf", size = 18423 }, +] + [[package]] name = "posthog" version = "3.6.6" @@ -3211,8 +3298,6 @@ version = "6.0.0" source = { registry = "https://pypi.org/simple" } sdist = { url = "https://files.pythonhosted.org/packages/18/c7/8c6872f7372eb6a6b2e4708b88419fb46b857f7a2e1892966b851cc79fc9/psutil-6.0.0.tar.gz", hash = "sha256:8faae4f310b6d969fa26ca0545338b21f73c6b15db7c4a8d934a5482faa818f2", size = 508067 } wheels = [ - { url = "https://files.pythonhosted.org/packages/c5/66/78c9c3020f573c58101dc43a44f6855d01bbbd747e24da2f0c4491200ea3/psutil-6.0.0-cp27-none-win32.whl", hash = "sha256:02b69001f44cc73c1c5279d02b30a817e339ceb258ad75997325e0e6169d8b35", size = 249766 }, - { url = "https://files.pythonhosted.org/packages/e1/3f/2403aa9558bea4d3854b0e5e567bc3dd8e9fbc1fc4453c0aa9aafeb75467/psutil-6.0.0-cp27-none-win_amd64.whl", hash = "sha256:21f1fb635deccd510f69f485b87433460a603919b45e2a324ad65b0cc74f8fb1", size = 253024 }, { url = "https://files.pythonhosted.org/packages/0b/37/f8da2fbd29690b3557cca414c1949f92162981920699cd62095a984983bf/psutil-6.0.0-cp36-abi3-macosx_10_9_x86_64.whl", hash = "sha256:c588a7e9b1173b6e866756dde596fd4cad94f9399daf99ad8c3258b3cb2b47a0", size = 250961 }, { url = "https://files.pythonhosted.org/packages/35/56/72f86175e81c656a01c4401cd3b1c923f891b31fbcebe98985894176d7c9/psutil-6.0.0-cp36-abi3-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:6ed2440ada7ef7d0d608f20ad89a04ec47d2d3ab7190896cd62ca5fc4fe08bf0", size = 287478 }, { url = "https://files.pythonhosted.org/packages/19/74/f59e7e0d392bc1070e9a70e2f9190d652487ac115bb16e2eff6b22ad1d24/psutil-6.0.0-cp36-abi3-manylinux_2_12_x86_64.manylinux2010_x86_64.manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5fd9a97c8e94059b0ef54a7d4baf13b405011176c3b6ff257c247cae0d560ecd", size = 290455 }, @@ -3591,6 +3676,25 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/11/c3/005fcca25ce078d2cc29fd559379817424e94885510568bc1bc53d7d5846/pytz-2024.2-py2.py3-none-any.whl", hash = "sha256:31c7c1817eb7fae7ca4b8c7ee50c72f93aa2dd863de768e1ef4245d426aa0725", size = 508002 }, ] +[[package]] +name = "pywin32" +version = "308" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/72/a6/3e9f2c474895c1bb61b11fa9640be00067b5c5b363c501ee9c3fa53aec01/pywin32-308-cp310-cp310-win32.whl", hash = "sha256:796ff4426437896550d2981b9c2ac0ffd75238ad9ea2d3bfa67a1abd546d262e", size = 5927028 }, + { url = "https://files.pythonhosted.org/packages/d9/b4/84e2463422f869b4b718f79eb7530a4c1693e96b8a4e5e968de38be4d2ba/pywin32-308-cp310-cp310-win_amd64.whl", hash = "sha256:4fc888c59b3c0bef905ce7eb7e2106a07712015ea1c8234b703a088d46110e8e", size = 6558484 }, + { url = "https://files.pythonhosted.org/packages/9f/8f/fb84ab789713f7c6feacaa08dad3ec8105b88ade8d1c4f0f0dfcaaa017d6/pywin32-308-cp310-cp310-win_arm64.whl", hash = "sha256:a5ab5381813b40f264fa3495b98af850098f814a25a63589a8e9eb12560f450c", size = 7971454 }, + { url = "https://files.pythonhosted.org/packages/eb/e2/02652007469263fe1466e98439831d65d4ca80ea1a2df29abecedf7e47b7/pywin32-308-cp311-cp311-win32.whl", hash = "sha256:5d8c8015b24a7d6855b1550d8e660d8daa09983c80e5daf89a273e5c6fb5095a", size = 5928156 }, + { url = "https://files.pythonhosted.org/packages/48/ef/f4fb45e2196bc7ffe09cad0542d9aff66b0e33f6c0954b43e49c33cad7bd/pywin32-308-cp311-cp311-win_amd64.whl", hash = "sha256:575621b90f0dc2695fec346b2d6302faebd4f0f45c05ea29404cefe35d89442b", size = 6559559 }, + { url = "https://files.pythonhosted.org/packages/79/ef/68bb6aa865c5c9b11a35771329e95917b5559845bd75b65549407f9fc6b4/pywin32-308-cp311-cp311-win_arm64.whl", hash = "sha256:100a5442b7332070983c4cd03f2e906a5648a5104b8a7f50175f7906efd16bb6", size = 7972495 }, + { url = "https://files.pythonhosted.org/packages/00/7c/d00d6bdd96de4344e06c4afbf218bc86b54436a94c01c71a8701f613aa56/pywin32-308-cp312-cp312-win32.whl", hash = "sha256:587f3e19696f4bf96fde9d8a57cec74a57021ad5f204c9e627e15c33ff568897", size = 5939729 }, + { url = "https://files.pythonhosted.org/packages/21/27/0c8811fbc3ca188f93b5354e7c286eb91f80a53afa4e11007ef661afa746/pywin32-308-cp312-cp312-win_amd64.whl", hash = "sha256:00b3e11ef09ede56c6a43c71f2d31857cf7c54b0ab6e78ac659497abd2834f47", size = 6543015 }, + { url = "https://files.pythonhosted.org/packages/9d/0f/d40f8373608caed2255781a3ad9a51d03a594a1248cd632d6a298daca693/pywin32-308-cp312-cp312-win_arm64.whl", hash = "sha256:9b4de86c8d909aed15b7011182c8cab38c8850de36e6afb1f0db22b8959e3091", size = 7976033 }, + { url = "https://files.pythonhosted.org/packages/a9/a4/aa562d8935e3df5e49c161b427a3a2efad2ed4e9cf81c3de636f1fdddfd0/pywin32-308-cp313-cp313-win32.whl", hash = "sha256:1c44539a37a5b7b21d02ab34e6a4d314e0788f1690d65b48e9b0b89f31abbbed", size = 5938579 }, + { url = "https://files.pythonhosted.org/packages/c7/50/b0efb8bb66210da67a53ab95fd7a98826a97ee21f1d22949863e6d588b22/pywin32-308-cp313-cp313-win_amd64.whl", hash = "sha256:fd380990e792eaf6827fcb7e187b2b4b1cede0585e3d0c9e84201ec27b9905e4", size = 6542056 }, + { url = "https://files.pythonhosted.org/packages/26/df/2b63e3e4f2df0224f8aaf6d131f54fe4e8c96400eb9df563e2aae2e1a1f9/pywin32-308-cp313-cp313-win_arm64.whl", hash = "sha256:ef313c46d4c18dfb82a2431e3051ac8f112ccee1a34f29c263c583c568db63cd", size = 7974986 }, +] + [[package]] name = "pyyaml" version = "6.0.2" @@ -3647,6 +3751,24 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/5a/66/bbb1dd374f5c870f59c5bb1db0e18cbe7fa739415a24cbd95b2d1f5ae0c4/pyyaml_env_tag-0.1-py3-none-any.whl", hash = "sha256:af31106dec8a4d68c60207c1886031cbf839b68aa7abccdb19868200532c2069", size = 3911 }, ] +[[package]] +name = "qdrant-client" +version = "1.12.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "grpcio" }, + { name = "grpcio-tools" }, + { name = "httpx", extra = ["http2"] }, + { name = "numpy" }, + { name = "portalocker" }, + { name = "pydantic" }, + { name = "urllib3" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/15/5e/ec560881e086f893947c8798949c72de5cfae9453fd05c2250f8dfeaa571/qdrant_client-1.12.1.tar.gz", hash = "sha256:35e8e646f75b7b883b3d2d0ee4c69c5301000bba41c82aa546e985db0f1aeb72", size = 237441 } +wheels = [ + { url = "https://files.pythonhosted.org/packages/68/c0/eef4fe9dad6d41333f7dc6567fa8144ffc1837c8a0edfc2317d50715335f/qdrant_client-1.12.1-py3-none-any.whl", hash = "sha256:b2d17ce18e9e767471368380dd3bbc4a0e3a0e2061fedc9af3542084b48451e0", size = 267171 }, +] + [[package]] name = "ragbits" version = "0.0.1" @@ -3703,6 +3825,9 @@ otel = [ promptfoo = [ { name = "pyyaml" }, ] +qdrant = [ + { name = "qdrant-client" }, +] [package.dev-dependencies] dev = [ @@ -3723,6 +3848,7 @@ requires-dist = [ { name = "opentelemetry-api", marker = "extra == 'otel'", specifier = "~=1.27.0" }, { name = "pydantic", specifier = ">=2.9.1" }, { name = "pyyaml", marker = "extra == 'promptfoo'", specifier = "~=6.0.2" }, + { name = "qdrant-client", marker = "extra == 'qdrant'", specifier = "~=1.12.1" }, { name = "tomli", specifier = "~=2.0.2" }, { name = "torch", marker = "extra == 'local'", specifier = "~=2.2.1" }, { name = "transformers", marker = "extra == 'local'", specifier = "~=4.44.2" }, @@ -3873,7 +3999,7 @@ version = "0.1.0" source = { virtual = "." } dependencies = [ { name = "ragbits-cli" }, - { name = "ragbits-core", extra = ["chroma", "lab", "litellm", "local", "otel"] }, + { name = "ragbits-core", extra = ["chroma", "lab", "litellm", "local", "otel", "qdrant"] }, { name = "ragbits-document-search", extra = ["gcs", "huggingface"] }, { name = "ragbits-evaluate", extra = ["relari"] }, { name = "ragbits-guardrails", extra = ["openai"] }, @@ -3901,7 +4027,7 @@ dev = [ [package.metadata] requires-dist = [ { name = "ragbits-cli", editable = "packages/ragbits-cli" }, - { name = "ragbits-core", extras = ["chroma", "lab", "litellm", "local", "otel"], editable = "packages/ragbits-core" }, + { name = "ragbits-core", extras = ["chroma", "lab", "litellm", "local", "otel", "qdrant"], editable = "packages/ragbits-core" }, { name = "ragbits-document-search", extras = ["gcs", "huggingface"], editable = "packages/ragbits-document-search" }, { name = "ragbits-evaluate", extras = ["relari"], editable = "packages/ragbits-evaluate" }, { name = "ragbits-guardrails", extras = ["openai"], editable = "packages/ragbits-guardrails" },