diff --git a/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py b/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py index 8d38830f1..d9ffa181f 100644 --- a/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py +++ b/packages/ragbits-core/src/ragbits/core/vector_stores/qdrant.py @@ -1,6 +1,7 @@ import json import typing +import httpx import qdrant_client from qdrant_client import AsyncQdrantClient, models from qdrant_client.models import Distance, Filter, VectorParams @@ -59,6 +60,9 @@ def from_config(cls, config: dict) -> Self: """ client_options = ObjectContructionConfig.model_validate(config["client"]) client_cls = import_by_path(client_options.type, qdrant_client) + if "limits" in client_options.config: + limits = httpx.Limits(**client_options.config["limits"]) + client_options.config["limits"] = limits config["client"] = client_cls(**client_options.config) return super().from_config(config) diff --git a/packages/ragbits-core/tests/unit/vector_stores/test_from_config.py b/packages/ragbits-core/tests/unit/vector_stores/test_from_config.py index 32f9e9e6d..112343c6d 100644 --- a/packages/ragbits-core/tests/unit/vector_stores/test_from_config.py +++ b/packages/ragbits-core/tests/unit/vector_stores/test_from_config.py @@ -1,3 +1,4 @@ +import httpx from chromadb import ClientAPI from qdrant_client import AsyncQdrantClient from qdrant_client.local.async_qdrant_local import AsyncQdrantLocal @@ -61,7 +62,7 @@ def test_subclass_from_config_chroma_client(): assert store.default_options.max_distance == 0.22 -def test_subclass_from_config_drant_client(): +def test_subclass_from_config_qdrant_client(): config = ObjectContructionConfig.model_validate( { "type": "ragbits.core.vector_stores.qdrant:QdrantVectorStore", @@ -70,6 +71,7 @@ def test_subclass_from_config_drant_client(): "type": "AsyncQdrantClient", "config": { "location": ":memory:", + "limits": {"keepalive_expiry": 20, "max_keepalive_connections": 0}, }, }, "index_name": "some_index", @@ -84,6 +86,9 @@ def test_subclass_from_config_drant_client(): assert isinstance(store, QdrantVectorStore) assert store._index_name == "some_index" assert isinstance(store._client, AsyncQdrantClient) + assert isinstance(store._client.init_options["limits"], httpx.Limits) + assert store._client.init_options["limits"].keepalive_expiry == 20 + assert store._client.init_options["limits"].max_keepalive_connections == 0 assert isinstance(store._client._client, AsyncQdrantLocal) assert store.default_options.k == 10 assert store.default_options.max_distance == 0.22