|
@@ -25,6 +25,8 @@ class MilvusCustomerMemoryIndex(CustomerMemoryIndex):
|
|
|
self._embedding_model: Any | None = None
|
|
self._embedding_model: Any | None = None
|
|
|
self._client_lock = Lock()
|
|
self._client_lock = Lock()
|
|
|
self._embedding_model_lock = Lock()
|
|
self._embedding_model_lock = Lock()
|
|
|
|
|
+ self._collection_lock = Lock()
|
|
|
|
|
+ self._collection_ready = False
|
|
|
|
|
|
|
|
def upsert(self, memory: CustomerMemory) -> None:
|
|
def upsert(self, memory: CustomerMemory) -> None:
|
|
|
try:
|
|
try:
|
|
@@ -134,8 +136,20 @@ class MilvusCustomerMemoryIndex(CustomerMemoryIndex):
|
|
|
return cast(list[list[float]], raw_vectors)
|
|
return cast(list[list[float]], raw_vectors)
|
|
|
|
|
|
|
|
def _ensure_collection(self, client: MilvusClient) -> None:
|
|
def _ensure_collection(self, client: MilvusClient) -> None:
|
|
|
- if client.has_collection(self._collection_name):
|
|
|
|
|
|
|
+ if self._collection_ready:
|
|
|
return
|
|
return
|
|
|
|
|
+ with self._collection_lock:
|
|
|
|
|
+ if self._collection_ready:
|
|
|
|
|
+ return
|
|
|
|
|
+ if not client.has_collection(self._collection_name):
|
|
|
|
|
+ self._create_collection(client)
|
|
|
|
|
+ client.load_collection(
|
|
|
|
|
+ collection_name=self._collection_name,
|
|
|
|
|
+ timeout=30.0,
|
|
|
|
|
+ )
|
|
|
|
|
+ self._collection_ready = True
|
|
|
|
|
+
|
|
|
|
|
+ def _create_collection(self, client: MilvusClient) -> None:
|
|
|
schema = MilvusClient.create_schema(
|
|
schema = MilvusClient.create_schema(
|
|
|
auto_id=False,
|
|
auto_id=False,
|
|
|
enable_dynamic_field=False,
|
|
enable_dynamic_field=False,
|