1"""
2Context Degradation Detection — Public API
3============================================
4 
5Detect, measure, and diagnose context degradation patterns in LLM agent systems.
6 
7Public API:
8    measure_attention_distribution  — Map attention weight across context positions.
9    detect_lost_in_middle           — Flag critical information in degraded-attention regions.
10    analyze_context_structure       — Assess structural degradation risk factors.
11    PoisoningDetector               — Detect context poisoning indicators (error accumulation,
12                                      contradictions, hallucination markers).
13    ContextHealthAnalyzer           — Run composite health analysis combining attention,
14                                      poisoning, and utilization metrics.
15    analyze_agent_context           — One-call convenience function for agent sessions.
16 
17PRODUCTION NOTES:
18- The attention estimation functions simulate U-shaped attention curves for demonstration
19  purposes. Production systems should extract actual attention weights from model internals
20  when available (e.g., via TransformerLens or model-specific APIs).
21- Token estimation uses simplified heuristics (~1 token per whitespace-split word).
22  Production systems should use model-specific tokenizers for accurate counts.
23- Poisoning and hallucination detection uses pattern matching as a proxy. Production
24  systems may benefit from fine-tuned classifiers or model-based detection.
25"""
26 
27import random
28import re
29from typing import Dict, List, Optional
30 
31__all__ = [
32    "measure_attention_distribution",
33    "detect_lost_in_middle",
34    "analyze_context_structure",
35    "PoisoningDetector",
36    "ContextHealthAnalyzer",
37    "analyze_agent_context",
38]
39 
40 
41# ---------------------------------------------------------------------------
42# Attention Distribution Analysis
43# ---------------------------------------------------------------------------
44 
45def measure_attention_distribution(
46    context_tokens: List[str],
47    query: str,
48) -> List[Dict[str, object]]:
49    """Map simulated attention weight to each context position.
50 
51    Use when: diagnosing whether critical information sits in the
52    low-attention middle region of a long context.
53 
54    Args:
55        context_tokens: Whitespace-split tokens (or chunks) of the context.
56        query: The query or task description the context is meant to support.
57 
58    Returns:
59        List of dicts, one per position, each containing:
60            position (int), attention (float), region (str), tokens (str | None).
61    """
62    n = len(context_tokens)
63    attention_by_position: List[Dict[str, object]] = []
64 
65    for position in range(n):
66        is_beginning = position < n * 0.1
67        is_end = position > n * 0.9
68 
69        attention = _estimate_attention(position, n, is_beginning, is_end)
70 
71        attention_by_position.append({
72            "position": position,
73            "attention": attention,
74            "region": "attention_favored" if (is_beginning or is_end) else "attention_degraded",
75            "tokens": context_tokens[position][:50] if position < 5 or position > n - 5 else None,
76        })
77 
78    return attention_by_position
79 
80 
81def _estimate_attention(
82    position: int,
83    total: int,
84    is_beginning: bool,
85    is_end: bool,
86) -> float:
87    """Estimate attention weight for a single position.
88 
89    Simulates the U-shaped attention curve documented in lost-in-middle research:
90    - Beginning tokens receive high attention (primacy / attention-sink effect).
91    - End tokens receive high attention (recency effect).
92    - Middle tokens receive degraded attention.
93 
94    IMPORTANT: This is a simulation for demonstration. Production systems should
95    extract actual attention weights from model forward passes or use
96    interpretability libraries (e.g., TransformerLens).
97    """
98    if is_beginning:
99        return 0.8 + random.random() * 0.2
100    elif is_end:
101        return 0.7 + random.random() * 0.3
102    else:
103        middle_progress = (position - total * 0.1) / (total * 0.8)
104        base_attention = 0.3 * (1 - middle_progress) + 0.1 * middle_progress
105        return base_attention + random.random() * 0.1
106 
107 
108# ---------------------------------------------------------------------------
109# Lost-in-Middle Detection
110# ---------------------------------------------------------------------------
111 
112def detect_lost_in_middle(
113    critical_positions: List[int],
114    attention_distribution: List[Dict[str, object]],
115) -> Dict[str, object]:
116    """Check if critical information sits in attention-degraded positions.
117 
118    Use when: context has been assembled and you need to verify that
119    high-priority content is not buried in the low-attention middle zone.
120 
121    Args:
122        critical_positions: Indices into the context that hold critical info.
123        attention_distribution: Output of ``measure_attention_distribution``.
124 
125    Returns:
126        Dict with keys: at_risk (list[int]), safe (list[int]),
127        recommendations (list[str]), degradation_score (float 0-1).
128    """
129    results: Dict[str, object] = {
130        "at_risk": [],
131        "safe": [],
132        "recommendations": [],
133        "degradation_score": 0.0,
134    }
135 
136    at_risk_count = 0
137    total_critical = len(critical_positions)
138 
139    for pos in critical_positions:
140        if pos < len(attention_distribution):
141            region = attention_distribution[pos]["region"]
142            if region == "attention_degraded":
143                results["at_risk"].append(pos)
144                at_risk_count += 1
145            else:
146                results["safe"].append(pos)
147 
148    if total_critical > 0:
149        results["degradation_score"] = at_risk_count / total_critical
150 
151    if results["at_risk"]:
152        results["recommendations"].extend([
153            "Move critical information to attention-favored positions",
154            "Use explicit markers to highlight critical information",
155            "Consider splitting context to reduce middle section",
156            f"{at_risk_count}/{total_critical} critical items are in degraded region",
157        ])
158 
159    return results
160 
161 
162# ---------------------------------------------------------------------------
163# Context Structure Analysis
164# ---------------------------------------------------------------------------
165 
166def analyze_context_structure(context: str) -> Dict[str, object]:
167    """Assess structural degradation risk factors in a context string.
168 
169    Use when: evaluating whether a context layout puts too much content
170    in the low-attention middle zone before sending it to a model.
171 
172    Args:
173        context: The full context string to analyze.
174 
175    Returns:
176        Dict with total_lines, sections list, middle_content_ratio,
177        and degradation_risk level (low / medium / high).
178    """
179    lines = context.split("\n")
180    sections: List[Dict[str, object]] = []
181 
182    current_section: Dict[str, object] = {"start": 0, "type": "unknown", "length": 0}
183 
184    for i, line in enumerate(lines):
185        if line.startswith("#"):
186            if current_section["length"] > 0:
187                sections.append(current_section)
188            current_section = {
189                "start": i,
190                "type": "header",
191                "length": 1,
192                "header": line.lstrip("#").strip(),
193            }
194        else:
195            current_section["length"] += 1
196 
197    sections.append(current_section)
198 
199    n = len(lines)
200    middle_start = int(n * 0.3)
201    middle_end = int(n * 0.7)
202 
203    middle_content = sum(
204        s["length"] for s in sections
205        if s["start"] >= middle_start and s["start"] <= middle_end
206    )
207 
208    middle_ratio = middle_content / n if n > 0 else 0
209    return {
210        "total_lines": n,
211        "sections": sections,
212        "middle_content_ratio": middle_ratio,
213        "degradation_risk": (
214            "high" if middle_ratio > 0.5
215            else "medium" if middle_ratio > 0.3
216            else "low"
217        ),
218    }
219 
220 
221# ---------------------------------------------------------------------------
222# Context Poisoning Detection
223# ---------------------------------------------------------------------------
224 
225class PoisoningDetector:
226    """Detect context poisoning indicators via pattern matching.
227 
228    Use when: context quality is suspect — outputs degrade on previously
229    successful tasks, tool calls misalign, or hallucinations persist
230    despite corrections.
231    """
232 
233    def __init__(self) -> None:
234        self.claims: List[Dict[str, object]] = []
235        self.error_patterns: List[str] = [
236            r"error",
237            r"failed",
238            r"exception",
239            r"cannot",
240            r"unable",
241            r"invalid",
242            r"not found",
243        ]
244 
245    def extract_claims(self, text: str) -> List[Dict[str, object]]:
246        """Extract claims from text for verification tracking.
247 
248        Use when: building a provenance chain to trace which claims
249        entered context and whether they have been verified.
250 
251        Args:
252            text: Raw text to extract claims from.
253 
254        Returns:
255            List of claim dicts with id, text, verified status, and
256            error indicator flag.
257        """
258        sentences = text.split(".")
259        claims: List[Dict[str, object]] = []
260 
261        for i, sentence in enumerate(sentences):
262            sentence = sentence.strip()
263            if len(sentence) < 10:
264                continue
265 
266            claims.append({
267                "id": i,
268                "text": sentence,
269                "verified": None,
270                "has_error_indicator": any(
271                    re.search(pattern, sentence, re.IGNORECASE)
272                    for pattern in self.error_patterns
273                ),
274            })
275 
276        self.claims.extend(claims)
277        return claims
278 
279    def detect_poisoning(self, context: str) -> Dict[str, object]:
280        """Detect potential context poisoning indicators.
281 
282        Use when: agent output quality has degraded and context
283        contamination is suspected. Checks for error accumulation,
284        contradictions, and hallucination markers.
285 
286        Args:
287            context: The full context string to analyze.
288 
289        Returns:
290            Dict with poisoning_risk (bool), indicators (list),
291            and overall_risk level (low / medium / high).
292        """
293        indicators: List[Dict[str, object]] = []
294 
295        # Check for error accumulation
296        error_count = sum(
297            1 for pattern in self.error_patterns
298            if re.search(pattern, context, re.IGNORECASE)
299        )
300 
301        if error_count > 3:
302            indicators.append({
303                "type": "error_accumulation",
304                "count": error_count,
305                "severity": "high" if error_count > 5 else "medium",
306                "message": f"Found {error_count} error indicators in context",
307            })
308 
309        # Check for contradiction patterns
310        contradictions = self._detect_contradictions(context)
311        if contradictions:
312            indicators.append({
313                "type": "contradictions",
314                "count": len(contradictions),
315                "examples": contradictions[:3],
316                "severity": "high",
317                "message": f"Found {len(contradictions)} potential contradictions",
318            })
319 
320        # Check for hallucination markers
321        hallucination_markers = self._detect_hallucination_markers(context)
322        if hallucination_markers:
323            indicators.append({
324                "type": "hallucination_markers",
325                "count": len(hallucination_markers),
326                "severity": "medium",
327                "message": f"Found {len(hallucination_markers)} phrases associated with uncertain claims",
328            })
329 
330        return {
331            "poisoning_risk": len(indicators) > 0,
332            "indicators": indicators,
333            "overall_risk": (
334                "high" if len(indicators) > 2
335                else "medium" if len(indicators) > 0
336                else "low"
337            ),
338        }
339 
340    def _detect_contradictions(self, text: str) -> List[str]:
341        """Detect potential contradictions in text."""
342        contradictions: List[str] = []
343 
344        conflict_patterns = [
345            (r"however", r"but"),
346            (r"on the other hand", r"instead"),
347            (r"although", r"yet"),
348            (r"despite", r"nevertheless"),
349        ]
350 
351        for pattern1, pattern2 in conflict_patterns:
352            if re.search(pattern1, text, re.IGNORECASE) and re.search(pattern2, text, re.IGNORECASE):
353                sentences = text.split(".")
354                for sentence in sentences:
355                    if (re.search(pattern1, sentence, re.IGNORECASE)
356                            or re.search(pattern2, sentence, re.IGNORECASE)):
357                        stripped = sentence.strip()
358                        if stripped and len(stripped) < 200:
359                            contradictions.append(stripped[:100])
360 
361        return contradictions[:5]
362 
363    def _detect_hallucination_markers(self, text: str) -> List[str]:
364        """Detect phrases associated with uncertain or hallucinated claims."""
365        markers = [
366            "may have been",
367            "might have",
368            "could potentially",
369            "possibly",
370            "apparently",
371            "reportedly",
372            "it is said that",
373            "sources suggest",
374            "believed to be",
375            "thought to be",
376        ]
377 
378        return [marker for marker in markers if marker in text.lower()]
379 
380 
381# ---------------------------------------------------------------------------
382# Context Health Analyzer
383# ---------------------------------------------------------------------------
384 
385class ContextHealthAnalyzer:
386    """Run composite health analysis on a context string.
387 
388    Use when: performing routine health checks on agent context during
389    long-running sessions, or when setting up automated monitoring that
390    triggers compaction or isolation before degradation hits.
391 
392    Combines attention distribution, poisoning detection, and utilization
393    metrics into a single 0-1 health score with status interpretation.
394    """
395 
396    def __init__(self, context_limit: int = 100_000) -> None:
397        self.context_limit: int = context_limit
398        self.metrics_history: List[Dict[str, object]] = []
399 
400    def analyze(
401        self,
402        context: str,
403        critical_positions: Optional[List[int]] = None,
404    ) -> Dict[str, object]:
405        """Perform comprehensive context health analysis.
406 
407        Use when: a single health-check call is needed that covers
408        attention, poisoning, and utilization in one pass.
409 
410        Args:
411            context: The full context string to analyze.
412            critical_positions: Indices of tokens holding critical info.
413                Defaults to the first 10 positions if not provided.
414 
415        Returns:
416            Dict with health_score (float 0-1), status (str),
417            metrics (dict), issues (dict), and recommendations (list[str]).
418        """
419        tokens = context.split()
420 
421        token_count = len(tokens)
422        utilization = token_count / self.context_limit
423 
424        attention_dist = measure_attention_distribution(
425            tokens[:1000],  # Sample for efficiency
426            "current_task",
427        )
428 
429        degradation = detect_lost_in_middle(
430            critical_positions or list(range(10)),
431            attention_dist,
432        )
433 
434        poisoning = PoisoningDetector().detect_poisoning(context)
435 
436        health_score = self._calculate_health_score(
437            utilization=utilization,
438            degradation=degradation["degradation_score"],
439            poisoning_risk=1.0 if poisoning["poisoning_risk"] else 0.0,
440        )
441 
442        result: Dict[str, object] = {
443            "health_score": health_score,
444            "status": self._interpret_score(health_score),
445            "metrics": {
446                "token_count": token_count,
447                "utilization": utilization,
448                "degradation_score": degradation["degradation_score"],
449                "poisoning_risk": poisoning["overall_risk"],
450            },
451            "issues": {
452                "lost_in_middle": degradation,
453                "poisoning": poisoning,
454            },
455            "recommendations": self._generate_recommendations(
456                utilization, degradation, poisoning
457            ),
458        }
459 
460        self.metrics_history.append(result)
461        return result
462 
463    def _calculate_health_score(
464        self,
465        utilization: float,
466        degradation: float,
467        poisoning_risk: float,
468    ) -> float:
469        """Calculate composite health score (0-1, higher is healthier)."""
470        utilization_penalty = min(utilization * 0.5, 0.3)
471        degradation_penalty = degradation * 0.3
472        poisoning_penalty = poisoning_risk * 0.2
473 
474        score = 1.0 - utilization_penalty - degradation_penalty - poisoning_penalty
475        return max(0.0, min(1.0, score))
476 
477    def _interpret_score(self, score: float) -> str:
478        """Map numeric score to human-readable status."""
479        if score > 0.8:
480            return "healthy"
481        elif score > 0.6:
482            return "warning"
483        elif score > 0.4:
484            return "degraded"
485        else:
486            return "critical"
487 
488    def _generate_recommendations(
489        self,
490        utilization: float,
491        degradation: Dict[str, object],
492        poisoning: Dict[str, object],
493    ) -> List[str]:
494        """Generate actionable recommendations based on analysis."""
495        recommendations: List[str] = []
496 
497        if utilization > 0.8:
498            recommendations.append("Context near limit - consider compaction")
499            recommendations.append("Implement observation masking for tool outputs")
500 
501        if degradation.get("at_risk"):
502            recommendations.append("Critical information in degraded attention region")
503            recommendations.append("Move key information to beginning or end of context")
504 
505        if poisoning["poisoning_risk"]:
506            recommendations.append("Context poisoning indicators detected")
507            recommendations.append("Review and remove potentially erroneous information")
508 
509        if not recommendations:
510            recommendations.append("Context appears healthy - continue monitoring")
511 
512        return recommendations
513 
514 
515# ---------------------------------------------------------------------------
516# Convenience Function
517# ---------------------------------------------------------------------------
518 
519def analyze_agent_context(
520    context: str,
521    context_limit: int = 80_000,
522    critical_positions: Optional[List[int]] = None,
523) -> Dict[str, object]:
524    """One-call health analysis for an agent session.
525 
526    Use when: a quick health check is needed without manually configuring
527    an analyzer instance. Prints a summary and returns the full result dict.
528 
529    Args:
530        context: The full context string to analyze.
531        context_limit: Maximum token budget for this agent's context window.
532        critical_positions: Indices of critical tokens. Defaults to [0..4].
533 
534    Returns:
535        Full health analysis dict from ``ContextHealthAnalyzer.analyze``.
536    """
537    analyzer = ContextHealthAnalyzer(context_limit=context_limit)
538 
539    if critical_positions is None:
540        critical_positions = list(range(5))
541 
542    result = analyzer.analyze(context, critical_positions)
543 
544    print(f"Health Score: {result['health_score']:.2f}")
545    print(f"Status: {result['status']}")
546    print("Recommendations:")
547    for rec in result["recommendations"]:
548        print(f"  - {rec}")
549 
550    return result
551 
552 
553# ---------------------------------------------------------------------------
554# CLI Demo
555# ---------------------------------------------------------------------------
556 
557if __name__ == "__main__":
558    # Demonstrate the public API with synthetic context
559    print("=" * 60)
560    print("Context Degradation Detector — Demo")
561    print("=" * 60)
562 
563    # Build a synthetic context with identifiable sections
564    intro = "System prompt: Analyze quarterly revenue data and produce a report. "
565    middle = "Background information. " * 200  # Filler to simulate long context
566    conclusion = "Key finding: Revenue increased 15% year-over-year. "
567    sample_context = intro + middle + conclusion
568 
569    print(f"\nSample context length: {len(sample_context.split())} tokens")
570 
571    # 1. Structure analysis
572    print("\n--- Structure Analysis ---")
573    structure = analyze_context_structure(sample_context)
574    print(f"  Lines: {structure['total_lines']}")
575    print(f"  Middle content ratio: {structure['middle_content_ratio']:.2f}")
576    print(f"  Degradation risk: {structure['degradation_risk']}")
577 
578    # 2. Attention distribution (first 50 tokens for brevity)
579    print("\n--- Attention Distribution (first 50 tokens) ---")
580    tokens = sample_context.split()[:50]
581    attention = measure_attention_distribution(tokens, "quarterly revenue")
582    favored = sum(1 for a in attention if a["region"] == "attention_favored")
583    degraded = sum(1 for a in attention if a["region"] == "attention_degraded")
584    print(f"  Favored positions: {favored}")
585    print(f"  Degraded positions: {degraded}")
586 
587    # 3. Lost-in-middle detection
588    print("\n--- Lost-in-Middle Detection ---")
589    critical = [0, 1, 2, 25, 26, 48, 49]  # Start, middle, end
590    lim_result = detect_lost_in_middle(critical, attention)
591    print(f"  At risk: {lim_result['at_risk']}")
592    print(f"  Safe: {lim_result['safe']}")
593    print(f"  Degradation score: {lim_result['degradation_score']:.2f}")
594 
595    # 4. Poisoning detection
596    print("\n--- Poisoning Detection ---")
597    poisoned_context = (
598        "The API returned an error. However, the system reportedly "
599        "recovered. But the error persisted and the request failed. "
600        "Unable to parse the response. Sources suggest the endpoint "
601        "may have been deprecated. Although retries succeeded, yet "
602        "the invalid token caused an exception."
603    )
604    detector = PoisoningDetector()
605    poisoning = detector.detect_poisoning(poisoned_context)
606    print(f"  Poisoning risk: {poisoning['poisoning_risk']}")
607    print(f"  Overall risk: {poisoning['overall_risk']}")
608    for indicator in poisoning["indicators"]:
609        print(f"    [{indicator['severity']}] {indicator['message']}")
610 
611    # 5. Full health analysis
612    print("\n--- Full Health Analysis ---")
613    result = analyze_agent_context(sample_context)
614    print(f"\n  Full result keys: {list(result.keys())}")
615