Merge pull request #121 from P3ngLiu/feature/v0.2.1/connection_check

Feature/v0.2.1/connection_check

omagent-core/src/omagent_core/engine/orkes/orkes_workflow_client.py

@@ -22,6 +22,11 @@ def __init__(
     ):
         super(OrkesWorkflowClient, self).__init__(configuration)
 
+    def check_connection(self) -> bool:
+        """Check if the connection to Conductor server is valid"""
+        self.metadataResourceApi.get_all_workflows()
+
+
     def start_workflow_by_name(
             self,
             name: str,

omagent-core/src/omagent_core/services/connectors/milvus.py

@@ -1,6 +1,5 @@
 from .base import ConnectorBase
-from pymilvus import Collection, DataType, MilvusClient, connections, utility
-from pymilvus.client import types
+from pymilvus import MilvusClient
 from typing import Any, Optional
 from pydantic import Field
 from omagent_core.utils.registry import registry
@@ -15,11 +14,18 @@ class MilvusConnector(ConnectorBase):
     db: Optional[str] = Field(default="default")
     alias: Optional[str] = Field(default="alias")
 
-    def model_post_init(self, __context: Any) -> None:                
-        self._client = MilvusClient(
-            uri=self.host,
-            user=self.username,
-            password=self.password,
-            db_name=self.db,
-        ) 
-
+    def model_post_init(self, __context: Any) -> None:  
+        try:
+            self._client = MilvusClient(
+                uri=self.host,
+                user=self.username,
+                password=self.password,
+                db_name=self.db,
+            )
+        except Exception as e:
+            raise ConnectionError(f"Connection to Milvus failed. Please check your connector config in container.yaml. \n Error Message: {e}")
+
+    def check_connection(self) -> bool:
+        """Check if the connection to Milvus is valid."""
+        # Try to list collections to verify connection
+        self._client.list_collections()

omagent-core/src/omagent_core/services/connectors/redis.py

@@ -21,4 +21,9 @@ def model_post_init(self, __context: Any) -> None:
             db=self.db,
             decode_responses=False
         )
-        self._client = Redis(connection_pool=pool)
+        self._client = Redis(connection_pool=pool)
+
+    def check_connection(self) -> bool:
+        """Check if Redis connection is valid by executing a simple ping command"""
+        self._client.ping()
+

omagent-core/src/omagent_core/utils/container.py

@@ -203,5 +203,24 @@ def clean_config_dict(config_dict: dict) -> dict:
                     overwrite=True,
                 )
 
+        self.check_connection()
 
+    def check_connection(self):
+        for name, connector in self._connectors.items():
+            try:
+                connector.check_connection()
+            except Exception as e:
+                raise ConnectionError(f"Connection to {name} failed. Please check your connector config in container.yaml. \n Error Message: {e}")
+
+        try:
+            from omagent_core.engine.orkes.orkes_workflow_client import OrkesWorkflowClient
+            conductor_client = OrkesWorkflowClient(self.conductor_config)
+            conductor_client.check_connection()
+        except Exception as e:
+            raise ConnectionError(f"Connection to Conductor failed. Please check your conductor config in container.yaml. \n Error Message: {e}")
+
+        print("--------------------------------")
+        print("All connections passed the connection check")
+        print("--------------------------------")
+
 container = Container()