Update memory_drift_diagnostician.py

memory_drift_diagnostician.py

@@ -1,6 +1,6 @@
 import logging
 import numpy as np
-from typing import Dict, Any, List
+from typing import Dict, Any, List, Optional
 from agentic_reliability_framework.runtime.agents.base import BaseAgent, AgentSpecialization
 from ai_event import AIEvent
 
@@ -19,7 +19,15 @@ class MemoryDriftDiagnosticianAgent(BaseAgent):
         self.zscore_threshold = zscore_threshold
         self._retrieval_scores_history: List[float] = []
 
-    async def analyze(self, event: AIEvent) -> Dict[str, Any]:
+    async def analyze(self, event: AIEvent, context_window: Optional[int] = None) -> Dict[str, Any]:
+        """
+        Analyze drift using retrieval scores.
+        Args:
+            event: AIEvent containing retrieval_scores.
+            context_window: Optional override of history window size.
+        Returns:
+            Dictionary with drift detection findings.
+        """
         try:
             if not event.retrieval_scores:
                 return {
@@ -32,7 +40,9 @@ class MemoryDriftDiagnosticianAgent(BaseAgent):
             current_avg = float(np.mean(event.retrieval_scores))
             self._retrieval_scores_history.append(current_avg)
 
-            if len(self._retrieval_scores_history) > self.history_window:
+            # Use provided context_window if given, else default
+            window = context_window if context_window is not None else self.history_window
+            if len(self._retrieval_scores_history) > window:
                 self._retrieval_scores_history.pop(0)
 
             if len(self._retrieval_scores_history) < 10: