crewplus 0.2.28__py3-none-any.whl → 0.2.29__py3-none-any.whl

crewplus/services/model_load_balancer.py

@@ -76,7 +76,7 @@ class ModelLoadBalancer:
             self.logger.error(f"Failed to load model configuration: {e}", exc_info=True)
             raise RuntimeError(f"Failed to load model configuration: {e}")
 
-    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None):
+    def get_model(self, provider: str = None, model_type: str = None, deployment_name: str = None, with_metadata: bool = False):
         """
         Get a model instance.
 
@@ -104,7 +104,11 @@ class ModelLoadBalancer:
             for model_config in self.models_config:
                 if model_config.get('deployment_name') == deployment_name:
                     model_id = model_config['id']
-                    return self.models[model_id]
+                    model = self.models[model_id]
+                    if with_metadata:
+                        return model, deployment_name
+                    return model
+                
             self.logger.error(f"No model found for deployment name: {deployment_name}")
             raise ValueError(f"No model found for deployment name: {deployment_name}")
 
@@ -116,7 +120,10 @@ class ModelLoadBalancer:
 
             selected_model_config = self._round_robin_selection(candidates)
             model_id = selected_model_config['id']
-            return self.models[model_id]
+            model = self.models[model_id]
+            if with_metadata:
+                return model, selected_model_config.get('deployment_name')
+            return model
 
         raise ValueError("Either 'deployment_name' or both 'provider' and 'model_type' must be provided.")
 

crewplus/vectorstores/milvus/vdb_service.py

@@ -10,6 +10,7 @@ from langchain_milvus import Milvus
 from langchain_core.embeddings import Embeddings
 from langchain_openai import AzureOpenAIEmbeddings
 from pymilvus import MilvusClient
+import time
 
 from ...services.init_services import get_model_balancer
 from .schema_milvus import SchemaMilvus, DEFAULT_SCHEMA
@@ -361,10 +362,9 @@ class VDBService(object):
             "params": {}
         }
         
-        vdb = Milvus(
-            embedding_function=embeddings,
+        vdb = self._create_milvus_instance_with_retry(
             collection_name=collection_name,
-            connection_args=self.connection_args,
+            embeddings=embeddings,
             index_params=index_params
         )
 
@@ -373,6 +373,32 @@ class VDBService(object):
 
         return vdb
 
+    def _create_milvus_instance_with_retry(self, collection_name: str, embeddings: Embeddings, index_params: dict) -> Milvus:
+        """
+        Creates a Milvus instance with a retry mechanism for connection failures.
+        """
+        retries = 2
+        for attempt in range(retries + 1):
+            try:
+                vdb = Milvus(
+                    embedding_function=embeddings,
+                    collection_name=collection_name,
+                    connection_args=self.connection_args,
+                    index_params=index_params
+                )
+                self.logger.info(f"Successfully connected to Milvus for collection '{collection_name}' on attempt {attempt + 1}.")
+                return vdb  # Return on success
+            except Exception as e:
+                self.logger.warning(
+                    f"Attempt {attempt + 1}/{retries + 1} to connect to Milvus for collection '{collection_name}' failed: {e}"
+                )
+                if attempt < retries:
+                    self.logger.info("Retrying in 3 seconds...")
+                    time.sleep(3)
+                else:
+                    self.logger.error(f"Failed to connect to Milvus for collection '{collection_name}' after {retries + 1} attempts.")
+                    raise RuntimeError(f"Could not connect to Milvus after {retries + 1} attempts.") from e
+
     def drop_collection(self, collection_name: str) -> None:
         """
         Deletes a collection from the vector database and removes it from the cache.

{crewplus-0.2.28.dist-info → crewplus-0.2.29.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: crewplus
-Version: 0.2.28
+Version: 0.2.29
 Summary: Base services for CrewPlus AI applications
 Author-Email: Tim Liu <tim@opsmateai.com>
 License: MIT

{crewplus-0.2.28.dist-info → crewplus-0.2.29.dist-info}/RECORD

@@ -1,13 +1,13 @@
-crewplus-0.2.28.dist-info/METADATA,sha256=f8Nk2TJu-nZ9Hsdsf1f4JSKfFF_NgNv3z3zDvJmEyQ0,5327
-crewplus-0.2.28.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
-crewplus-0.2.28.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
-crewplus-0.2.28.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
+crewplus-0.2.29.dist-info/METADATA,sha256=Gjp8NS1iow7bCkTVYioXaFmWPr6RKFBqiCrMEzTdVJg,5327
+crewplus-0.2.29.dist-info/WHEEL,sha256=9P2ygRxDrTJz3gsagc0Z96ukrxjr-LFBGOgv3AuKlCA,90
+crewplus-0.2.29.dist-info/entry_points.txt,sha256=6OYgBcLyFCUgeqLgnvMyOJxPCWzgy7se4rLPKtNonMs,34
+crewplus-0.2.29.dist-info/licenses/LICENSE,sha256=2_NHSHRTKB_cTcT_GXgcenOCtIZku8j343mOgAguTfc,1087
 crewplus/__init__.py,sha256=m46HkZL1Y4toD619NL47Sn2Qe084WFFSFD7e6VoYKZc,284
 crewplus/services/__init__.py,sha256=zUM4ZwUfGMBDx-j7Wehf_KC5yYXPTK8BK_oeO5veIXQ,398
 crewplus/services/azure_chat_model.py,sha256=WMSf4BDO8UcP7ZASNGRJxdTEnuWBmCRSY_4yx_VMbok,5499
 crewplus/services/gemini_chat_model.py,sha256=oVLL07VEjwrHOeP56YSPLnldbfjvTVRrsTI6xcPNt1E,35224
 crewplus/services/init_services.py,sha256=EBpDkIwzuujmdlqjyWvdLQCfhQmfS_OKFz-9Ji8nmAU,1628
-crewplus/services/model_load_balancer.py,sha256=HH_eHxFfxgarPWFGpANg7dgShnWca4q46Jz0b1vJ4Sw,9405
+crewplus/services/model_load_balancer.py,sha256=meJlmyJoKPXA6AlIYiVJETd8bS6zeaDevMuVLs6LFhk,9715
 crewplus/services/tracing_manager.py,sha256=aCU9N4Jvh8pDD3h8kWX4O-Ax8xwdLHnQ4wJ3sf-vLwA,6289
 crewplus/utils/__init__.py,sha256=2Gk1n5srFJQnFfBuYTxktdtKOVZyNrFcNaZKhXk35Pw,142
 crewplus/utils/schema_action.py,sha256=GDaBoVFQD1rXqrLVSMTfXYW1xcUu7eDcHsn57XBSnIg,422
@@ -15,9 +15,9 @@ crewplus/utils/schema_document_updater.py,sha256=frvffxn2vbi71fHFPoGb9hq7gH2azmm
 crewplus/vectorstores/milvus/__init__.py,sha256=egGncAdjlXG6ekTQvKMKqhvKBifrUrPlsSB0-bpvq4A,229
 crewplus/vectorstores/milvus/milvus_schema_manager.py,sha256=2IZT61LVui21Pt5Z3y8YYS2dYcwzkgUKxMq2NA0-lQE,9222
 crewplus/vectorstores/milvus/schema_milvus.py,sha256=IvKdUCH451HJ-F3TUR5jDjqwQlQs4SEXAQ_th4JAnfc,12117
-crewplus/vectorstores/milvus/vdb_service.py,sha256=wCltxZc0aD27iTu7wjveHqQWPEF2VyO4B2WGQCheeVs,21118
+crewplus/vectorstores/milvus/vdb_service.py,sha256=CaUMLIMeOCm2R4t5EKtAupIddFXQu0NSb8RpTkInGd4,22498
 docs/GeminiChatModel.md,sha256=zZYyl6RmjZTUsKxxMiC9O4yV70MC4TD-IGUmWhIDBKA,8677
 docs/ModelLoadBalancer.md,sha256=aGHES1dcXPz4c7Y8kB5-vsCNJjriH2SWmjBkSGoYKiI,4398
 docs/VDBService.md,sha256=Dw286Rrf_fsi13jyD3Bo4Sy7nZ_G7tYm7d8MZ2j9hxk,9375
 docs/index.md,sha256=3tlc15uR8lzFNM5WjdoZLw0Y9o1P1gwgbEnOdIBspqc,1643
-crewplus-0.2.28.dist-info/RECORD,,
+crewplus-0.2.29.dist-info/RECORD,,