refactor(logging): replace logger with agent_logger for consistent logging in ModelWeightManager

agents/ensemble_weights.py

@@ -4,7 +4,6 @@ from utils.registry import MODEL_REGISTRY # Import MODEL_REGISTRY
 from utils.agent_logger import AgentLogger
 
 agent_logger = AgentLogger()
-logger = logging.getLogger(__name__)
 
 class ContextualWeightOverrideAgent:
     def __init__(self):
@@ -49,7 +48,7 @@ class ModelWeightManager:
         num_models = len(MODEL_REGISTRY)
         if num_models > 0:
             if strongest_model_id and strongest_model_id in MODEL_REGISTRY:
-                logger.info(f"Designating '{strongest_model_id}' as the strongest model.")
+                agent_logger.log("weight_optimization", "info", f"Designating '{strongest_model_id}' as the strongest model.")
                 # Assign a high weight to the strongest model (e.g., 50%)
                 strongest_weight_share = 0.5
                 self.base_weights = {strongest_model_id: strongest_weight_share}
@@ -62,12 +61,12 @@ class ModelWeightManager:
                     self.base_weights[strongest_model_id] = 1.0
             else:
                 if strongest_model_id and strongest_model_id not in MODEL_REGISTRY:
-                    logger.warning(f"Strongest model ID '{strongest_model_id}' not found in MODEL_REGISTRY. Distributing weights equally.")
+                    agent_logger.log("weight_optimization", "warning", f"Strongest model ID '{strongest_model_id}' not found in MODEL_REGISTRY. Distributing weights equally.")
                 initial_weight = 1.0 / num_models
                 self.base_weights = {model_id: initial_weight for model_id in MODEL_REGISTRY.keys()}
         else:
             self.base_weights = {} # Handle case with no registered models
-        logger.info(f"Base weights initialized: {self.base_weights}")
+        agent_logger.log("weight_optimization", "info", f"Base weights initialized: {self.base_weights}")
         
         self.situation_weights = {
             "high_confidence": 1.2,    # Boost weights for high confidence predictions
@@ -85,7 +84,7 @@ class ModelWeightManager:
 
         # 1. Apply contextual overrides first
         if context_tags:
-            logger.info(f"Applying contextual overrides for tags: {context_tags}")
+            agent_logger.log("weight_optimization", "info", f"Applying contextual overrides for tags: {context_tags}")
             overrides = self.context_override_agent.get_overrides(context_tags)
             for model_id, multiplier in overrides.items():
                 adjusted_weights[model_id] = adjusted_weights.get(model_id, 0.0) * multiplier
@@ -109,7 +108,7 @@ class ModelWeightManager:
             agent_logger.log("weight_optimization", "info", f"Adjusted weights after conflict reduction: {adjusted_weights}")
         
         # Adjust based on confidence
-        logger.info("Adjusting weights based on model confidence scores.")
+        agent_logger.log("weight_optimization", "info", "Adjusting weights based on model confidence scores.")
         for model, confidence in confidence_scores.items():
             if confidence > 0.8:
                 adjusted_weights[model] *= self.situation_weights["high_confidence"]
@@ -117,28 +116,28 @@ class ModelWeightManager:
             elif confidence < 0.5:
                 adjusted_weights[model] *= self.situation_weights["low_confidence"]
                 agent_logger.log("weight_optimization", "info", f"Model '{model}' has low confidence ({confidence:.2f}). Weight reduced.")
-        logger.info(f"Adjusted weights before normalization: {adjusted_weights}")
+        agent_logger.log("weight_optimization", "info", f"Adjusted weights before normalization: {adjusted_weights}")
         
         normalized_weights = self._normalize_weights(adjusted_weights)
-        logger.info(f"Final normalized adjusted weights: {normalized_weights}")
+        agent_logger.log("weight_optimization", "info", f"Final normalized adjusted weights: {normalized_weights}")
         return normalized_weights
     
     def _has_consensus(self, predictions):
         """Check if models agree on prediction"""
         agent_logger.log("weight_optimization", "info", "Checking for consensus among model predictions.")
         non_none_predictions = [p.get("Label") for p in predictions.values() if p is not None and isinstance(p, dict) and p.get("Label") is not None and p.get("Label") != "Error"]
-        logger.debug(f"Non-none predictions for consensus check: {non_none_predictions}")
+        agent_logger.log("weight_optimization", "debug", f"Non-none predictions for consensus check: {non_none_predictions}")
         result = len(non_none_predictions) > 0 and len(set(non_none_predictions)) == 1
-        logger.info(f"Consensus detected: {result}")
+        agent_logger.log("weight_optimization", "info", f"Consensus detected: {result}")
         return result
     
     def _has_conflicts(self, predictions):
         """Check if models have conflicting predictions"""
         agent_logger.log("weight_optimization", "info", "Checking for conflicts among model predictions.")
         non_none_predictions = [p.get("Label") for p in predictions.values() if p is not None and isinstance(p, dict) and p.get("Label") is not None and p.get("Label") != "Error"]
-        logger.debug(f"Non-none predictions for conflict check: {non_none_predictions}")
+        agent_logger.log("weight_optimization", "debug", f"Non-none predictions for conflict check: {non_none_predictions}")
         result = len(non_none_predictions) > 1 and len(set(non_none_predictions)) > 1
-        logger.info(f"Conflicts detected: {result}")
+        agent_logger.log("weight_optimization", "info", f"Conflicts detected: {result}")
         return result
     
     def _normalize_weights(self, weights):
@@ -155,4 +154,4 @@ class ModelWeightManager:
                 return {} # No models registered
         normalized = {k: v/total for k, v in weights.items()}
         agent_logger.log("weight_optimization", "info", f"Weights normalized. Total sum: {sum(normalized.values()):.2f}")
-        return normalized 
+        return normalized