SIENTIAPDE-1184

refactor: enhance shutdown procedures and documentation across activities

- Added shutdown methods to MongoDB and Email classes to ensure proper resource cleanup.
- Updated docstrings for shutdown methods to clarify their purpose and functionality.
- Enhanced documentation for various methods across multiple classes, improving clarity on parameters and return values.
- Improved the main function and other utility functions with detailed docstrings for better understanding and maintainability.

orchestrator/activities/activities.py

@@ -82,4 +82,9 @@ class Activities(  # Couchbase,
                           notification_handler=notification_handler)
 
     def shutdown(self):
+        """
+        Shutdown the MongoDB connection and clean up resources.
+        """
         MongoDB.shutdown(self)
+        Postgres.close(self)
+        Email.shutdown(self)

orchestrator/activities/email.py

@@ -42,6 +42,12 @@ class Email(BaseActivity):
                               logger=logger,
                               notification_handler=notification_handler)
 
+    def shutdown(self):
+        """
+        Shutdown the Email connection and clean up resources.
+        """
+        self.server.quit()
+
     @activity.defn(name="build_email_html")
     async def build_email_html(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
@@ -105,9 +111,15 @@ class Email(BaseActivity):
 
     def try_send_email(self, msg: MIMEMultipart, receivers: str):
         """
-        Sends an email to the receivers.
+        Sends an email to the receivers with automatic reconnection handling.
-        """
 
+        Args:
+            msg (MIMEMultipart): The email message to send.
+            receivers (str): Comma-separated list of email addresses to send to.
+
+        Raises:
+            Exception: If email sending fails after reconnection attempts.
+        """
         try:
             self.server.sendmail(
                 self.sender_email, receivers, msg.as_string())

orchestrator/activities/mongo_db.py

@@ -64,7 +64,7 @@ class MongoDB(BaseActivity):
 
     def shutdown(self):
         """
-        Close the MongoDB client connection.
+        Shutdown the MongoDB connection and clean up resources.
         """
         try:
             if self.client:
@@ -81,6 +81,16 @@ class MongoDB(BaseActivity):
         self.shutdown()
 
     def find(self, collection_name: str, filters: dict[str, Any]) -> list[dict[str, Any]]:
+        """
+        Find documents in a MongoDB collection based on the provided filters.
+
+        Args:
+            collection_name (str): The name of the collection to search in.
+            filters (dict[str, Any]): The query filters to apply.
+
+        Returns:
+            list[dict[str, Any]]: List of documents matching the filters, with _id fields removed.
+        """
         collection = self.database[collection_name]
 
         documents = list(collection.find(filters, {"_id": 0}))

orchestrator/activities/slot_manager.py

@@ -202,6 +202,14 @@ class SlotManager(Redis):
     async def get_last_data_timestamp(self, input_data: dict[str, Any]) -> str | None:
         """
         Gets the last data timestamp from redis.
+
+        Args:
+            input_data (dict[str, Any]): The input data containing:
+                - metadata (dict): Metadata for logging purposes.
+                - mail_type (str): The type of mail to get timestamp for.
+
+        Returns:
+            str | None: The last data timestamp as a string, or None if no timestamp exists.
         """
         metadata = input_data['metadata']
         key = f"notification_last_timestamp:{input_data['mail_type']}"
@@ -233,6 +241,15 @@ class SlotManager(Redis):
     async def put_last_data_timestamp(self, input_data: dict[str, Any]):
         """
         Puts the last data timestamp into redis.
+
+        Args:
+            input_data (dict[str, Any]): The input data containing:
+                - metadata (dict): Metadata for logging purposes.
+                - data (list[dict]): The data to extract timestamp from.
+                - mail_type (str): The type of mail to store timestamp for.
+
+        Returns:
+            str | None: The last data timestamp that was stored, or None if no data exists.
         """
         metadata = input_data['metadata']
         key = f"notification_last_timestamp:{input_data['mail_type']}"
@@ -270,7 +287,18 @@ class SlotManager(Redis):
     @activity.defn(name="filter_notification_alerts")
     async def filter_notification_alerts(self, input_data: dict[str, Any]) -> dict[str, Any]:
         """
-        Filter notification alerts
+        Filter notification alerts based on sending configurations and notification package.
+
+        Args:
+            input_data (dict[str, Any]): The input data containing:
+                - metadata (dict): Metadata for logging purposes.
+                - notification_package (list): The package of notifications to filter.
+                - sending_configs (list): The configurations for sending notifications.
+                    Each config should have 'group_name', 'contents', and optionally 'ignore' fields.
+                - notification_ttl (int): Time to live for notifications in seconds.
+
+        Returns:
+            dict[str, Any]: The filtered receiver groups with their notifications.
         """
         metadata = input_data['metadata']
         notification_package = input_data['notification_package']
@@ -330,7 +358,13 @@ class SlotManager(Redis):
     @activity.defn(name="store_notification_cache")
     async def store_notification_cache(self, input_data: dict[str, Any]) -> None:
         """
-        Store notification cache
+        Store notification cache in Redis to track recently sent notifications.
+
+        Args:
+            input_data (dict[str, Any]): The input data containing:
+                - metadata (dict): Metadata for logging purposes.
+                - log_report (list[dict]): The log report containing notification statuses.
+                - sent_ttl (int): Time to live for sent notification cache in seconds.
         """
         metadata = input_data['metadata']
         log_report = DataFrame(input_data['log_report'])

orchestrator/activities/temporal_manager.py

@@ -42,6 +42,10 @@ class TemporalManager(BaseActivity):
                               notification_handler=notification_handler)
 
     async def connect_to_temporal(self):
+        """
+        Connect to Temporal server namespaces for scouter and laborious workflows.
+        Creates client connections to both namespaces and stores them for later use.
+        """
         self.logger.info(
             f"Connecting to Temporal side namespaces at {self.temporal_host}")
         self.logger.info(f"Scouter namespace: {self.scouter_namespace}")

orchestrator/utils/connectors_config.py

@@ -2,6 +2,12 @@ from os import getenv
 
 
 def build_redis_config():
+    """
+    Build Redis configuration from environment variables.
+
+    Returns:
+        dict: Redis configuration with host, port, username, and password.
+    """
     return {
         'host': getenv('REDIS_HOST', 'localhost'),
         'port': int(getenv('REDIS_PORT', '6379')),
@@ -11,6 +17,12 @@ def build_redis_config():
 
 
 def build_mongodb_config():
+    """
+    Build MongoDB configuration from environment variables.
+
+    Returns:
+        dict: MongoDB configuration with connection string, database name, and TTL index seconds.
+    """
     username = getenv('MONGODB_USERNAME', 'root')
     password = getenv('MONGODB_PASSWORD', 'wKZDbMNU1c')
     uri = getenv('MONGODB_URL', 'localhost:27018')
@@ -24,6 +36,12 @@ def build_mongodb_config():
 
 
 def build_couchbase_config():
+    """
+    Build Couchbase configuration from environment variables.
+
+    Returns:
+        dict: Couchbase configuration with connection string, username, and password.
+    """
     return {
         'connection_string': getenv('COUCHBASE_CONNECTION_STRING', 'couchbase://localhost'),
         'username': getenv('COUCHBASE_USERNAME', 'sientia'),
@@ -32,6 +50,12 @@ def build_couchbase_config():
 
 
 def build_temporal_config():
+    """
+    Build Temporal configuration from environment variables.
+
+    Returns:
+        dict: Temporal configuration with host and namespace settings.
+    """
     return {
         'temporal_host': getenv('TEMPORAL_HOST', 'localhost:7233'),
         'temporal_namespace': getenv('TEMPORAL_NAMESPACE', 'default'),
@@ -41,6 +65,12 @@ def build_temporal_config():
 
 
 def build_postgres_config():
+    """
+    Build PostgreSQL configuration from environment variables.
+
+    Returns:
+        dict: PostgreSQL configuration with connection details and connection pool settings.
+    """
     return {
         'host': getenv('POSTGRES_HOST', 'localhost'),
         'port': int(getenv('POSTGRES_PORT', '5432')),
@@ -53,6 +83,12 @@ def build_postgres_config():
 
 
 def build_email_config():
+    """
+    Build email configuration from environment variables.
+
+    Returns:
+        dict: Email configuration with SMTP server settings and sender credentials.
+    """
     return {
         'sender_email': getenv('EMAIL_SENDER', 'sientia-alerts@aignosi.com'),
         'sender_password': getenv('EMAIL_SENDER_PASSWORD', 'sientia'),

orchestrator/utils/email_builder.py

@@ -18,12 +18,33 @@ class EmailBuilder:
             self.general_template = file.read()
 
     def replace_parameters(self, template: str, parameters: dict) -> str:
+        """
+        Replace parameters in a Jinja2 template with provided values.
+
+        Args:
+            template (str): The Jinja2 template string.
+            parameters (dict): Dictionary of parameters to replace in the template.
+
+        Returns:
+            str: The rendered template with parameters replaced.
+        """
         # Criar um template Jinja2
         template = Template(template)
 
         return template.render(parameters)
 
     def parameters(self, general_events: dict, mail_type: str) -> dict:
+        """
+        Build parameters dictionary for email templates based on general events and mail type.
+
+        Args:
+            general_events (dict): Dictionary containing events categorized by level (ERROR, WARNING, INFO).
+                Each level contains a 'models' key with model-specific event data.
+            mail_type (str): The type of email being sent.
+
+        Returns:
+            dict: Dictionary with mail_type and rendered event sections for each notification level.
+        """
         error_models = general_events.get('ERROR', {}).get('models', [])
         warning_models = general_events.get('WARNING', {}).get('models', [])
         info_models = general_events.get('INFO', {}).get('models', [])
@@ -43,7 +64,16 @@ class EmailBuilder:
 
     def build_email(self, report_data: list[dict], mail_type: str) -> str:
         """
-        Builds the email html.
+        Builds the email HTML by organizing report data by notification level and model.
+
+        Args:
+            report_data (list[dict]): List of notification reports, each containing:
+                - level (str): Notification level (ERROR, WARNING, INFO)
+                - model_name (str): Name of the model
+                - Additional notification details
+
+        Returns:
+            str: Complete HTML email content ready for sending.
         """
         general_events = {}
 

orchestrator/utils/orchestrator_functions.py

@@ -2,6 +2,21 @@ from typing import Any
 
 
 def common_config(config: dict[str, Any]):
+    """
+    Extract common configuration parameters from a pipeline configuration.
+
+    Args:
+        config (dict[str, Any]): Pipeline configuration containing:
+            - workflow_type (str): Type of workflow (e.g., 'scouter', 'predictions_batch')
+            - schedule_name (str): Name of the schedule
+            - frequency (str, optional): Frequency of execution (default: '1m')
+            - max_retry_policy (int, optional): Maximum retry attempts (default: 1)
+            - model_id (str): ID of the model
+            - models (dict): Model configuration containing 'name' field
+
+    Returns:
+        dict[str, Any]: Common configuration dictionary with extracted parameters.
+    """
     return {
         "workflow_type": config['workflow_type'],
         "schedule_name": config['schedule_name'],
@@ -14,6 +29,18 @@ def common_config(config: dict[str, Any]):
 
 
 def minimal_retrain(config: dict[str, Any]):
+    """
+    Build minimal retrain configuration from pipeline config.
+
+    Args:
+        config (dict[str, Any]): Pipeline configuration containing:
+            - schedule_name (str): Name of the schedule
+            - query (str): SQL query for retraining
+            - Additional fields from common_config
+
+    Returns:
+        dict[str, Any]: Minimal retrain configuration with workflow type set to 'minimal_retrain'.
+    """
     return {
         **common_config(config),
         "workflow_type": "minimal_retrain",
@@ -25,6 +52,25 @@ def minimal_retrain(config: dict[str, Any]):
 
 
 def scouter(config: dict[str, Any]):
+    """
+    Build scouter configuration from pipeline config.
+
+    Args:
+        config (dict[str, Any]): Pipeline configuration containing:
+            - filters (list[dict], optional): List of filter configurations
+            - read_tags (list[dict]): List of tag configurations with:
+                - filter_name (str): Name of the filter
+                - policy (str): Filter policy
+                - tag_name (str): Name of the tag
+                - aggr_func (str, optional): Aggregation function (default: 'lts')
+                - data_range (list[int], optional): Data range limits (default: [-100, 100])
+            - tag_retention_minutes (int, optional): Tag retention time in minutes (default: 60)
+            - debug_data_package (bool, optional): Enable debug data package (default: False)
+            - Additional fields from common_config
+
+    Returns:
+        dict[str, Any]: Scouter configuration with topic, filters, tags, and retention settings.
+    """
     filters = {}
     for f in config.get('filters', []):
         filters[f['filter_name']] = {
@@ -53,6 +99,19 @@ def scouter(config: dict[str, Any]):
 
 
 def overlap_filter_config(base_filter_config: dict[str, Any], config: list[dict[str, Any]]):
+    """
+    Overlap filter configuration with base filter config.
+
+    Args:
+        base_filter_config (dict[str, Any]): Base filter configuration to extend.
+        config (list[dict[str, Any]]): List of filter configurations to add, each containing:
+            - filter_name (str): Name of the filter
+            - policy (str): Filter policy
+            - config (dict, optional): Additional filter configuration
+
+    Returns:
+        dict[str, Any]: Extended filter configuration with new filters added.
+    """
     for fil in config:
         base_filter_config[fil['filter_name']] = {
             "policy": fil['policy'],
@@ -63,6 +122,15 @@ def overlap_filter_config(base_filter_config: dict[str, Any], config: list[dict[
 
 
 def process_path_priority(path_priority: list[str]):
+    """
+    Process and normalize path priority list to ensure it contains the required priorities.
+
+    Args:
+        path_priority (list[str]): List of path priorities to process.
+
+    Returns:
+        list[str]: Normalized path priority list with exactly 3 elements: ["STOP", "CONTINUE", "REPEAT"].
+    """
     for priority in path_priority[:]:
         if priority not in ["STOP", "CONTINUE", "REPEAT"]:
             path_priority.remove(priority)
@@ -75,6 +143,26 @@ def process_path_priority(path_priority: list[str]):
 
 
 def predictions_batch(config: dict[str, Any]):
+    """
+    Build predictions batch configuration from pipeline config.
+
+    Args:
+        config (dict[str, Any]): Pipeline configuration containing:
+            - write_tags (list[dict]): List of tag configurations with:
+                - server_id (str): ID of the OPC server
+                - type (str): Tag type ('prediction' or 'confidence')
+                - addr (str): Tag address
+                - data_type (str, optional): Data type (default: 'float')
+            - path_priority (list[str], optional): List of path priorities (default: ["STOP", "CONTINUE", "REPEAT"])
+            - input_filters (list[dict], optional): List of input filter configurations
+            - mlflow_transform_filters (list[dict], optional): List of MLflow transform filter configurations
+            - mlflow_predict_filters (list[dict], optional): List of MLflow predict filter configurations
+            - model_retention_minutes (int, optional): Model retention time in minutes (default: 60)
+            - Additional fields from common_config
+
+    Returns:
+        dict[str, Any]: Predictions batch configuration with OPC output config, filters, and path priority.
+    """
     tags = {}
     for tag in config.get('write_tags', []):
         if tag['server_id'] not in tags:
@@ -156,6 +244,23 @@ def gather_read_tags(pipelines: list[dict[str, Any]]) -> dict[str, Any]:
 
 def build_tag_config(tag: dict[str, Any], slot_config: dict[str, Any],
                      opc_servers: dict[str, Any], i: int):
+    """
+    Build tag configuration for a specific slot and OPC server.
+
+    Args:
+        tag (dict[str, Any]): Tag configuration containing:
+            - server_id (str): ID of the OPC server
+            - tag_address (str): Address of the tag
+        slot_config (dict[str, Any]): Current slot configuration to update.
+        opc_servers (dict[str, Any]): Dictionary of OPC server configurations.
+        i (int): Slot number to configure.
+
+    Returns:
+        dict[str, Any]: Updated slot configuration with the new tag.
+
+    Raises:
+        ValueError: If the specified server_id is not found in opc_servers.
+    """
     server_id = tag['server_id']
 
     if server_id not in opc_servers:

orchestrator/worker/worker.py

@@ -29,6 +29,13 @@ POD_ID = os.getenv("POD_ID")
 
 
 async def main():
+    """
+    Main function to initialize and run the Temporal worker.
+
+    Sets up MongoDB connection, notification handler, Temporal client, and starts
+    multiple workers for different task queues (orchestrator, alerts, reports).
+    Handles graceful shutdown and error handling.
+    """
     host = os.getenv('TEMPORAL_HOST', 'localhost:7233')
     namespace = os.getenv('TEMPORAL_NAMESPACE', 'default')
     logger = get_logger(__name__)
@@ -176,6 +183,12 @@ async def main():
 
 
 def start_prometheus_server():
+    """
+    Start the Prometheus metrics server on the configured port.
+
+    Sets up HTTP server for metrics collection and marks the application as UP.
+    Exits the application if the server fails to start.
+    """
     try:
         port = int(os.getenv("HTTP_METRICS_PORT", 9090))
         start_http_server(port)