fix: warm up the vector

private_gpt/components/vector_store/factory.py

@@ -21,5 +21,8 @@ class VectorStoreFactory(ABC):
         self.settings = settings
         self.embed_dim = embed_dim
 
+    def warm_up(self) -> None:
+        return None
+
     @abstractmethod
     def vector_store(self, collection: str) -> BasePydanticVectorStore: ...

private_gpt/components/vector_store/qdrant_factory.py

@@ -1,4 +1,5 @@
 import typing
+from threading import Lock
 from typing import Any
 
 from llama_index.core.vector_stores.types import BasePydanticVectorStore
@@ -15,6 +16,7 @@ class QdrantVectorStoreFactory(VectorStoreFactory):
     def __init__(self, settings: Settings, embed_dim: int | None = None) -> None:
         super().__init__(settings)
         self._client: Any | None = None
+        self._client_lock = Lock()
         self._embed_dim = embed_dim
 
     def _build_client(self) -> Any:
@@ -33,8 +35,14 @@ class QdrantVectorStoreFactory(VectorStoreFactory):
         return QdrantClientBuilder.build_clients(self.settings)
 
     def _ensure_client(self) -> None:
-        if self._client is None:
-            self._client = self._build_client()
+        if self._client is not None:
+            return
+        with self._client_lock:
+            if self._client is None:
+                self._client = self._build_client()
+
+    def warm_up(self) -> None:
+        self._ensure_client()
 
     def vector_store(self, collection: str) -> BasePydanticVectorStore:
         try:

private_gpt/components/vector_store/vector_store_component.py

@@ -77,6 +77,9 @@ class VectorStoreComponent:
     def vector_store(self, collection: str) -> BasePydanticVectorStore:
         return self._factory.vector_store(collection)
 
+    def warm_up(self) -> None:
+        self._factory.warm_up()
+
     def get_filters(
         self,
         artifacts: list[str] | None = None,

private_gpt/eager_loading.py

@@ -40,7 +40,7 @@ def _warm_base(injector: Injector) -> None:
 def _warm_stores(injector: Injector) -> None:
     logger.debug("Warming stores")
     injector.get(NodeStoreComponent)
-    injector.get(VectorStoreComponent)
+    injector.get(VectorStoreComponent).warm_up()
 
 
 def _warm_streaming(injector: Injector) -> None: