solace-agent-mesh 1.5.1__py3-none-any.whl → 1.6.0__py3-none-any.whl

solace_agent_mesh/evaluation/subscriber.py

@@ -3,129 +3,47 @@ Refactored message subscriber with improved structure and readability.
 This module handles Solace message subscription and processing for evaluation.
 """
 
-import os
 import json
+import logging
+import os
 import threading
 import time
-from typing import Dict, List, Optional, Any, Set, Callable
 from dataclasses import dataclass, field
-from enum import Enum
-import logging
-from dotenv import load_dotenv
+from pathlib import Path
 
+from dotenv import load_dotenv
 from solace.messaging.messaging_service import MessagingService
 from solace.messaging.resources.topic_subscription import TopicSubscription
-from solace.messaging.config.solace_properties import (
-    transport_layer_properties,
-    service_properties,
-    authentication_properties,
-    transport_layer_security_properties,
+
+from .shared import (
+    ALLOWED_TOPIC_INFIXES,
+    BLOCKED_TOPIC_INFIXES,
+    MESSAGE_TIMEOUT,
+    BrokerConfig,
+    BrokerConnectionError,
+    ConfigurationError,
+    ConnectionState,
+    MessageProcessingError,
 )
 
-# Set up logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
+log = logging.getLogger(__name__)
 
 # Load environment variables
 load_dotenv()
 
 
-class ConnectionState(Enum):
-    """Represents the connection state of the subscriber."""
-
-    DISCONNECTED = "disconnected"
-    CONNECTING = "connecting"
-    CONNECTED = "connected"
-    DISCONNECTING = "disconnecting"
-    ERROR = "error"
-
-
-class SubscriberError(Exception):
-    """Base exception for subscriber-related errors."""
-
-    pass
-
-
-class BrokerConnectionError(SubscriberError):
-    """Raised when broker connection fails."""
-
-    pass
-
-
-class MessageProcessingError(SubscriberError):
-    """Raised when message processing fails."""
-
-    pass
-
-
-class ConfigurationError(SubscriberError):
-    """Raised when configuration is invalid."""
-
-    pass
-
-
-@dataclass
-class BrokerConfig:
-    """Broker connection configuration with validation."""
-
-    host: str
-    vpn_name: str
-    username: str
-    password: str
-    cert_validated: bool = False
-    connection_timeout: int = 30
-    reconnect_attempts: int = 3
-    reconnect_delay: float = 1.0
-
-    def __post_init__(self):
-        """Validate broker configuration after initialization."""
-        if not self.host or not self.host.strip():
-            raise ConfigurationError("Broker host cannot be empty")
-
-        if not self.vpn_name or not self.vpn_name.strip():
-            raise ConfigurationError("VPN name cannot be empty")
-
-        if not self.username or not self.username.strip():
-            raise ConfigurationError("Username cannot be empty")
-
-        if not self.password or not self.password.strip():
-            raise ConfigurationError("Password cannot be empty")
-
-        if self.connection_timeout <= 0:
-            raise ConfigurationError("Connection timeout must be positive")
-
-        if self.reconnect_attempts < 0:
-            raise ConfigurationError("Reconnect attempts cannot be negative")
-
-    def to_solace_properties(self) -> Dict[str, Any]:
-        """Convert to Solace messaging properties."""
-        return {
-            transport_layer_properties.HOST: self.host,
-            service_properties.VPN_NAME: self.vpn_name,
-            authentication_properties.SCHEME_BASIC_USER_NAME: self.username,
-            authentication_properties.SCHEME_BASIC_PASSWORD: self.password,
-            transport_layer_security_properties.CERT_VALIDATED: self.cert_validated,
-        }
-
-
 @dataclass
 class SubscriptionConfig:
     """Subscription configuration and topic filters."""
 
     namespace: str
-    allowed_topic_infixes: List[str] = field(
-        default_factory=lambda: [
-            "/agent/request/",
-            "/gateway/status/",
-            "/gateway/response/",
-        ]
+    allowed_topic_infixes: list[str] = field(
+        default_factory=lambda: ALLOWED_TOPIC_INFIXES
     )
-    blocked_topic_infixes: List[str] = field(
-        default_factory=lambda: [
-            "/discovery/"
-        ]
+    blocked_topic_infixes: list[str] = field(
+        default_factory=lambda: BLOCKED_TOPIC_INFIXES
     )
-    message_timeout: int = 1000
+    message_timeout: int = MESSAGE_TIMEOUT
     filter_non_final_status: bool = True
     remove_config_keys: bool = False
 
@@ -148,9 +66,7 @@ class SubscriptionConfig:
     def is_topic_allowed(self, topic: str) -> bool:
         """Check if a topic is allowed based on configured infixes."""
         # return any(infix in topic for infix in self.allowed_topic_infixes)
-        if any(infix in topic for infix in self.blocked_topic_infixes):
-            return False
-        return True
+        return not any(infix in topic for infix in self.blocked_topic_infixes)
 
 
 @dataclass
@@ -158,11 +74,11 @@ class ProcessedMessage:
     """Structured representation of a processed message."""
 
     topic: str
-    payload: Any
+    payload: any
     timestamp: float = field(default_factory=time.time)
-    message_type: Optional[str] = None
+    message_type: str | None = None
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self) -> dict[str, any]:
         """Convert to dictionary for JSON serialization."""
         return {
             "topic": self.topic,
@@ -185,7 +101,7 @@ class MessageSanitizer:
     """Handles message sanitization and cleaning."""
 
     @staticmethod
-    def remove_key_recursive(obj: Any, key_to_remove: str) -> None:
+    def remove_key_recursive(obj: any, key_to_remove: str) -> None:
         """
         Recursively remove a key from nested dictionaries and lists.
 
@@ -206,10 +122,10 @@ class MessageSanitizer:
                 for item in obj:
                     MessageSanitizer.remove_key_recursive(item, key_to_remove)
         except Exception as e:
-            logger.warning(f"Error during key removal: {e}")
+            log.warning(f"Error during key removal: {e}")
 
     @staticmethod
-    def sanitize_message(payload: Any, remove_config: bool = True) -> Any:
+    def sanitize_message(payload: any, remove_config: bool = True) -> any:
         """
         Sanitize message payload by removing unwanted keys.
 
@@ -239,7 +155,7 @@ class MessageProcessor:
         self.processed_count = 0
         self.error_count = 0
 
-    def process_message(self, inbound_message) -> Optional[ProcessedMessage]:
+    def process_message(self, inbound_message) -> ProcessedMessage | None:
         """
         Process an inbound message and return a ProcessedMessage if valid.
 
@@ -280,10 +196,10 @@ class MessageProcessor:
 
         except Exception as e:
             self.error_count += 1
-            logger.warning(f"Error processing message: {e}")
+            log.warning(f"Error processing message: {e}")
             return None
 
-    def _extract_payload(self, inbound_message) -> Optional[Any]:
+    def _extract_payload(self, inbound_message):
         """Extract and parse payload from inbound message."""
         try:
             payload_bytes = inbound_message.get_payload_as_bytes()
@@ -300,10 +216,10 @@ class MessageProcessor:
                 return payload_str
 
         except Exception as e:
-            logger.warning(f"Error extracting payload: {e}")
+            log.warning(f"Error extracting payload: {e}")
             return None
 
-    def _should_filter_status_message(self, topic: str, payload: Any) -> bool:
+    def _should_filter_status_message(self, topic: str, payload: any) -> bool:
         """Check if a status message should be filtered out."""
         if not self.config.filter_non_final_status:
             return False
@@ -325,12 +241,12 @@ class MessageProcessor:
 
         return False
 
-    def _find_part_type(self, data: Any, type_to_find: str) -> bool:
+    def _find_part_type(self, data: any, type_to_find: str) -> bool:
         """Recursively search for a part with a specific type."""
         if isinstance(data, dict):
             if data.get("type") == type_to_find:
                 return True
-            for key, value in data.items():
+            for _key, value in data.items():
                 if self._find_part_type(value, type_to_find):
                     return True
         elif isinstance(data, list):
@@ -350,7 +266,7 @@ class MessageProcessor:
         else:
             return "unknown"
 
-    def get_stats(self) -> Dict[str, int]:
+    def get_stats(self) -> dict[str, int]:
         """Get processing statistics."""
         return {
             "processed_count": self.processed_count,
@@ -362,14 +278,14 @@ class TaskTracker:
     """Tracks task completion and manages active tasks."""
 
     def __init__(
-        self, active_tasks: Set[str], wave_complete_event: Optional[threading.Event]
+        self, active_tasks: set[str], wave_complete_event: threading.Event | None
     ):
         self.active_tasks = active_tasks
         self.wave_complete_event = wave_complete_event
-        self.completed_tasks: List[TaskCompletionEvent] = []
+        self.completed_tasks: list[TaskCompletionEvent] = []
         self._lock = threading.Lock()
 
-    def handle_task_completion(self, topic: str) -> Optional[TaskCompletionEvent]:
+    def handle_task_completion(self, topic: str) -> TaskCompletionEvent | None:
         """
         Handle task completion based on topic.
 
@@ -389,7 +305,7 @@ class TaskTracker:
 
             with self._lock:
                 if task_id in self.active_tasks:
-                    logger.info(f"Task {task_id} completed")
+                    log.info(f"Task {task_id} completed")
                     self.active_tasks.remove(task_id)
 
                     completion_event = TaskCompletionEvent(task_id=task_id, topic=topic)
@@ -397,17 +313,17 @@ class TaskTracker:
 
                     # Check if all tasks are complete
                     if not self.active_tasks and self.wave_complete_event:
-                        logger.info("All tasks completed, setting wave complete event")
+                        log.info("All tasks completed, setting wave complete event")
                         self.wave_complete_event.set()
 
                     return completion_event
 
         except Exception as e:
-            logger.error(f"Error handling task completion: {e}")
+            log.error(f"Error handling task completion: {e}")
 
         return None
 
-    def _extract_task_id(self, topic: str) -> Optional[str]:
+    def _extract_task_id(self, topic: str) -> str | None:
         """Extract task ID from topic."""
         try:
             return topic.split("/")[-1]
@@ -428,9 +344,9 @@ class TaskTracker:
 class MessageStorage:
     """Handles message storage and file operations."""
 
-    def __init__(self, results_path: str):
-        self.results_path = results_path
-        self.messages: List[ProcessedMessage] = []
+    def __init__(self, results_path: str | Path):
+        self.results_path = Path(results_path)
+        self.messages: list[ProcessedMessage] = []
         self._lock = threading.Lock()
 
     def add_message(self, message: ProcessedMessage) -> None:
@@ -453,25 +369,25 @@ class MessageStorage:
         Returns:
             The full path to the saved file
         """
-        output_file = os.path.join(self.results_path, filename)
+        output_file = self.results_path / filename
 
         try:
             # Ensure directory exists
-            os.makedirs(os.path.dirname(output_file), exist_ok=True)
+            output_file.parent.mkdir(parents=True, exist_ok=True)
 
             with self._lock:
                 # Convert messages to dictionaries for JSON serialization
                 message_dicts = [msg.to_dict() for msg in self.messages]
 
-            with open(output_file, "w") as f:
+            with output_file.open("w") as f:
                 json.dump(message_dicts, f, indent=4)
 
-            logger.info(f"Saved {len(message_dicts)} messages to {output_file}")
-            return output_file
+            log.info(f"Saved {len(message_dicts)} messages to {output_file}")
+            return str(output_file)
 
         except Exception as e:
-            logger.error(f"Error saving messages: {e}")
-            raise MessageProcessingError(f"Failed to save messages: {e}")
+            log.error(f"Error saving messages: {e}")
+            raise MessageProcessingError(f"Failed to save messages: {e}") from e
 
     def clear_messages(self) -> None:
         """Clear all stored messages."""
@@ -484,7 +400,7 @@ class BrokerConnectionService:
 
     def __init__(self, config: BrokerConfig):
         self.config = config
-        self.messaging_service: Optional[MessagingService] = None
+        self.messaging_service: MessagingService | None = None
         self.connection_state = ConnectionState.DISCONNECTED
         self._connection_lock = threading.Lock()
 
@@ -492,13 +408,13 @@ class BrokerConnectionService:
         """Connect to the Solace broker."""
         with self._connection_lock:
             if self.connection_state == ConnectionState.CONNECTED:
-                logger.warning("Already connected to broker")
+                log.warning("Already connected to broker")
                 return
 
             self.connection_state = ConnectionState.CONNECTING
 
             try:
-                logger.info("Connecting to Solace PubSub+ Broker...")
+                log.info("Connecting to Solace PubSub+ Broker...")
 
                 broker_props = self.config.to_solace_properties()
                 self.messaging_service = (
@@ -507,37 +423,37 @@ class BrokerConnectionService:
                 self.messaging_service.connect()
 
                 self.connection_state = ConnectionState.CONNECTED
-                logger.info("Successfully connected to broker")
+                log.info("Successfully connected to broker")
 
             except Exception as e:
                 self.connection_state = ConnectionState.ERROR
-                logger.error(f"Failed to connect to broker: {e}")
-                raise BrokerConnectionError(f"Connection failed: {e}")
+                log.error(f"Failed to connect to broker: {e}")
+                raise BrokerConnectionError(f"Connection failed: {e}") from e
 
     def disconnect(self) -> None:
         """Disconnect from the Solace broker."""
         with self._connection_lock:
             if self.connection_state == ConnectionState.DISCONNECTED:
-                logger.warning("Already disconnected from broker")
+                log.warning("Already disconnected from broker")
                 return
 
             self.connection_state = ConnectionState.DISCONNECTING
 
             try:
                 if self.messaging_service:
-                    logger.info("Disconnecting from broker...")
+                    log.info("Disconnecting from broker...")
                     self.messaging_service.disconnect()
                     self.messaging_service = None
 
                 self.connection_state = ConnectionState.DISCONNECTED
-                logger.info("Successfully disconnected from broker")
+                log.info("Successfully disconnected from broker")
 
             except Exception as e:
                 self.connection_state = ConnectionState.ERROR
-                logger.error(f"Error during disconnect: {e}")
-                raise BrokerConnectionError(f"Disconnect failed: {e}")
+                log.error(f"Error during disconnect: {e}")
+                raise BrokerConnectionError(f"Disconnect failed: {e}") from e
 
-    def get_messaging_service(self) -> Optional[MessagingService]:
+    def get_messaging_service(self) -> MessagingService | None:
         """Get the messaging service instance."""
         return self.messaging_service
 
@@ -563,12 +479,12 @@ class SubscriptionManager:
         self._receiver_lock = threading.Lock()
 
     def start_subscription(
-        self, subscription_ready_event: Optional[threading.Event] = None
+        self, subscription_ready_event: threading.Event | None = None
     ) -> None:
         """Start message subscription."""
         with self._receiver_lock:
             if self.subscription_active:
-                logger.warning("Subscription already active")
+                log.warning("Subscription already active")
                 return
 
             if not self.connection_service.is_connected():
@@ -592,17 +508,17 @@ class SubscriptionManager:
                 self.message_receiver.add_subscription(subscription)
 
                 self.subscription_active = True
-                logger.info(f"Started subscription to: {self.config.topic_pattern}")
+                log.info(f"Started subscription to: {self.config.topic_pattern}")
 
                 # Signal that subscription is ready
                 if subscription_ready_event:
                     subscription_ready_event.set()
 
             except Exception as e:
-                logger.error(f"Failed to start subscription: {e}")
-                raise BrokerConnectionError(f"Subscription failed: {e}")
+                log.error(f"Failed to start subscription: {e}")
+                raise BrokerConnectionError(f"Subscription failed: {e}") from e
 
-    def receive_message(self, timeout: Optional[int] = None):
+    def receive_message(self, timeout: int | None = None):
         """
         Receive a message from the subscription.
 
@@ -620,7 +536,7 @@ class SubscriptionManager:
         try:
             return self.message_receiver.receive_message(timeout=timeout_ms)
         except Exception as e:
-            logger.warning(f"Error receiving message: {e}")
+            log.warning(f"Error receiving message: {e}")
             return None
 
     def stop_subscription(self) -> None:
@@ -631,51 +547,51 @@ class SubscriptionManager:
 
             try:
                 if self.message_receiver:
-                    logger.info("Stopping message receiver...")
+                    log.info("Stopping message receiver...")
                     self.message_receiver.terminate()
                     self.message_receiver = None
 
                 self.subscription_active = False
-                logger.info("Subscription stopped")
+                log.info("Subscription stopped")
 
             except Exception as e:
-                logger.error(f"Error stopping subscription: {e}")
+                log.error(f"Error stopping subscription: {e}")
 
     def is_active(self) -> bool:
         """Check if subscription is active."""
         return self.subscription_active
 
 
-class MessageSubscriber(threading.Thread):
+class Subscriber(threading.Thread):
     """
     Main message subscriber class that orchestrates all components.
-
     This is the refactored version of the original Subscriber class,
     maintaining the same interface while providing better structure.
     """
 
     def __init__(
         self,
+        broker_config: BrokerConfig,
         namespace: str,
-        active_tasks: Set[str],
-        wave_complete_event: Optional[threading.Event],
-        subscription_ready_event: Optional[threading.Event],
+        active_tasks: set[str],
+        wave_complete_event: threading.Event | None,
+        subscription_ready_event: threading.Event | None,
         results_path: str,
     ):
         """
         Initialize the message subscriber.
-
         Args:
+            broker_config: The broker configuration object
             namespace: The namespace for topic subscription
             active_tasks: Set of active task IDs to track
             wave_complete_event: Event to set when all tasks complete
             subscription_ready_event: Event to set when subscription is ready
             results_path: Path to save results
         """
-        super().__init__(name="MessageSubscriber")
+        super().__init__(name="Subscriber")
 
         # Initialize configuration
-        self.broker_config = self._create_broker_config()
+        self.broker_config = broker_config
         self.subscription_config = SubscriptionConfig(namespace=namespace)
 
         # Initialize services
@@ -696,24 +612,11 @@ class MessageSubscriber(threading.Thread):
         self.messages_received = 0
         self.messages_processed = 0
 
-    def _create_broker_config(self) -> BrokerConfig:
-        """Create broker configuration from environment variables."""
-        try:
-            return BrokerConfig(
-                host=os.environ.get("SOLACE_BROKER_URL", ""),
-                vpn_name=os.environ.get("SOLACE_BROKER_VPN", ""),
-                username=os.environ.get("SOLACE_BROKER_USERNAME", ""),
-                password=os.environ.get("SOLACE_BROKER_PASSWORD", ""),
-            )
-        except ConfigurationError as e:
-            logger.error(f"Invalid broker configuration: {e}")
-            raise
-
     def run(self) -> None:
         """Main thread execution method."""
         try:
             self._running = True
-            logger.info("Starting message subscriber...")
+            log.info("Starting message subscriber...")
 
             # Connect to broker
             self.connection_service.connect()
@@ -725,13 +628,13 @@ class MessageSubscriber(threading.Thread):
             self._message_processing_loop()
 
         except Exception as e:
-            logger.error(f"Error in subscriber thread: {e}")
+            log.error(f"Error in subscriber thread: {e}")
         finally:
             self._cleanup()
 
     def _message_processing_loop(self) -> None:
         """Main message processing loop."""
-        logger.info("Starting message processing loop...")
+        log.info("Starting message processing loop...")
 
         while self._running:
             try:
@@ -744,7 +647,7 @@ class MessageSubscriber(threading.Thread):
 
             except Exception as e:
                 if self._running:
-                    logger.error(f"Error in message processing loop: {e}")
+                    log.error(f"Error in message processing loop: {e}")
                     # Continue processing other messages
                     continue
 
@@ -764,11 +667,11 @@ class MessageSubscriber(threading.Thread):
                 self.task_tracker.handle_task_completion(processed_message.topic)
 
         except Exception as e:
-            logger.warning(f"Error handling message: {e}")
+            log.warning(f"Error handling message: {e}")
 
     def stop(self) -> None:
         """Stop the subscriber and clean up resources."""
-        logger.info("Stopping message subscriber...")
+        log.info("Stopping message subscriber...")
         self._running = False
 
     def _cleanup(self) -> None:
@@ -787,68 +690,37 @@ class MessageSubscriber(threading.Thread):
             self._log_final_statistics()
 
         except Exception as e:
-            logger.error(f"Error during cleanup: {e}")
+            log.error(f"Error during cleanup: {e}")
 
     def _log_final_statistics(self) -> None:
         """Log final processing statistics."""
         runtime = time.time() - self.start_time
         processor_stats = self.message_processor.get_stats()
 
-        logger.info("=== SUBSCRIBER STATISTICS ===")
-        logger.info(f"Runtime: {runtime:.2f} seconds")
-        logger.info(f"Messages received: {self.messages_received}")
-        logger.info(f"Messages processed: {self.messages_processed}")
-        logger.info(f"Messages stored: {self.message_storage.get_message_count()}")
-        logger.info(f"Processing errors: {processor_stats['error_count']}")
-        logger.info(
+        log.info("=== SUBSCRIBER STATISTICS ===")
+        log.info(f"Runtime: {runtime:.2f} seconds")
+        log.info(f"Messages received: {self.messages_received}")
+        log.info(f"Messages processed: {self.messages_processed}")
+        log.info(f"Messages stored: {self.message_storage.get_message_count()}")
+        log.info(f"Processing errors: {processor_stats['error_count']}")
+        log.info(
             f"Active tasks remaining: {self.task_tracker.get_active_task_count()}"
         )
-        logger.info(f"Tasks completed: {self.task_tracker.get_completed_task_count()}")
-        logger.info("=============================")
+        log.info(f"Tasks completed: {self.task_tracker.get_completed_task_count()}")
+        log.info("=============================")
 
     # Backward compatibility properties
     @property
-    def active_tasks(self) -> Set[str]:
+    def active_tasks(self) -> set[str]:
         """Get active tasks set for backward compatibility."""
         return self.task_tracker.active_tasks
 
     @property
-    def messages(self) -> List[Dict[str, Any]]:
+    def messages(self) -> list[dict[str, any]]:
         """Get messages list for backward compatibility."""
         return [msg.to_dict() for msg in self.message_storage.messages]
 
 
-# Backward compatibility alias
-Subscriber = MessageSubscriber
-
-
-def create_subscriber_from_env(
-    namespace: str,
-    active_tasks: Set[str],
-    wave_complete_event: Optional[threading.Event] = None,
-    subscription_ready_event: Optional[threading.Event] = None,
-    results_path: str = ".",
-) -> MessageSubscriber:
-    """
-    Factory function to create a subscriber with environment-based configuration.
-
-    Args:
-        namespace: The namespace for topic subscription
-        active_tasks: Set of active task IDs to track
-        wave_complete_event: Event to set when all tasks complete
-        subscription_ready_event: Event to set when subscription is ready
-        results_path: Path to save results
-
-    Returns:
-        Configured MessageSubscriber instance
-    """
-    return MessageSubscriber(
-        namespace=namespace,
-        active_tasks=active_tasks,
-        wave_complete_event=wave_complete_event,
-        subscription_ready_event=subscription_ready_event,
-        results_path=results_path,
-    )
 
 
 def main():
@@ -858,7 +730,7 @@ def main():
 
     # Set up signal handling for graceful shutdown
     def signal_handler(signum, frame):
-        print("\nShutting down subscriber...")
+        log.info("\nShutting down subscriber...")
         if "subscriber" in locals():
             subscriber.stop()
             subscriber.join()
@@ -872,7 +744,14 @@ def main():
     subscription_ready = threading.Event()
 
     try:
-        subscriber = create_subscriber_from_env(
+        broker_config = BrokerConfig(
+            host=os.environ.get("SOLACE_BROKER_URL", ""),
+            vpn_name=os.environ.get("SOLACE_BROKER_VPN", ""),
+            username=os.environ.get("SOLACE_BROKER_USERNAME", ""),
+            password=os.environ.get("SOLACE_BROKER_PASSWORD", ""),
+        )
+        subscriber = Subscriber(
+            broker_config=broker_config,
             namespace="test",
             active_tasks=active_tasks,
             subscription_ready_event=subscription_ready,
@@ -883,13 +762,13 @@ def main():
 
         # Wait for subscription to be ready
         subscription_ready.wait(timeout=30)
-        print("Subscriber is ready and running...")
+        log.info("Subscriber is ready and running...")
 
         # Keep running until interrupted
         subscriber.join()
 
     except Exception as e:
-        print(f"Error running subscriber: {e}")
+        log.error(f"Error running subscriber: {e}")
         sys.exit(1)