streamctx 0.3.0__tar.gz → 0.3.1__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- {streamctx-0.3.0/src/streamctx.egg-info → streamctx-0.3.1}/PKG-INFO +1 -1
- {streamctx-0.3.0 → streamctx-0.3.1}/setup.py +1 -1
- streamctx-0.3.1/src/streamctx/__init__.py +156 -0
- streamctx-0.3.1/src/streamctx/compressor.py +100 -0
- streamctx-0.3.1/src/streamctx/differ.py +199 -0
- streamctx-0.3.1/src/streamctx/healer.py +124 -0
- streamctx-0.3.1/src/streamctx/poison_detector.py +245 -0
- streamctx-0.3.1/src/streamctx/pricing.py +68 -0
- streamctx-0.3.1/src/streamctx/reporter.py +128 -0
- streamctx-0.3.1/src/streamctx/storage.py +200 -0
- streamctx-0.3.1/src/streamctx/tracker.py +456 -0
- {streamctx-0.3.0 → streamctx-0.3.1/src/streamctx.egg-info}/PKG-INFO +1 -1
- {streamctx-0.3.0 → streamctx-0.3.1}/src/streamctx.egg-info/SOURCES.txt +9 -0
- streamctx-0.3.1/src/streamctx.egg-info/top_level.txt +1 -0
- streamctx-0.3.0/src/streamctx.egg-info/top_level.txt +0 -1
- {streamctx-0.3.0 → streamctx-0.3.1}/LICENSE +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/README.md +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/setup.cfg +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/src/streamctx.egg-info/dependency_links.txt +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/src/streamctx.egg-info/requires.txt +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_checkpoint.py +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_compression.py +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_diff.py +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_healing.py +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_openai_integration.py +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_poison.py +0 -0
- {streamctx-0.3.0 → streamctx-0.3.1}/tests/test_streaming.py +0 -0
|
@@ -5,7 +5,7 @@ with open("README.md", encoding="utf-8") as f:
|
|
|
5
5
|
|
|
6
6
|
setup(
|
|
7
7
|
name="streamctx",
|
|
8
|
-
version="0.3.
|
|
8
|
+
version="0.3.1",
|
|
9
9
|
description="Context health monitoring for AI agents — detect poisoning, drift, loops",
|
|
10
10
|
long_description=long_description,
|
|
11
11
|
long_description_content_type="text/markdown",
|
|
@@ -0,0 +1,156 @@
|
|
|
1
|
+
|
|
2
|
+
"""
|
|
3
|
+
streamctx — zero-config LLM call tracing with context reuse insights.
|
|
4
|
+
|
|
5
|
+
Usage::
|
|
6
|
+
|
|
7
|
+
import streamctx
|
|
8
|
+
streamctx.start()
|
|
9
|
+
"""
|
|
10
|
+
|
|
11
|
+
from __future__ import annotations
|
|
12
|
+
|
|
13
|
+
from typing import Any
|
|
14
|
+
|
|
15
|
+
from .reporter import print_auto_summary, print_report
|
|
16
|
+
from .tracker import get_tracker
|
|
17
|
+
from .compressor import compress_messages, get_compression_stats
|
|
18
|
+
from .poison_detector import PoisonDetector
|
|
19
|
+
from .differ import ContextDiffer
|
|
20
|
+
|
|
21
|
+
__version__ = "0.3.0"
|
|
22
|
+
__all__ = [
|
|
23
|
+
"start", "wrap", "report", "stop",
|
|
24
|
+
"checkpoint", "resume", "get_session_id",
|
|
25
|
+
"compress", "compression_stats",
|
|
26
|
+
"healing_stats",
|
|
27
|
+
"scan",
|
|
28
|
+
"context_diff",
|
|
29
|
+
"__version__",
|
|
30
|
+
]
|
|
31
|
+
|
|
32
|
+
_detector = PoisonDetector()
|
|
33
|
+
_differ = ContextDiffer()
|
|
34
|
+
|
|
35
|
+
|
|
36
|
+
def start() -> None:
|
|
37
|
+
"""Enable tracing, monkeypatch OpenAI/Anthropic SDKs, and start token tracking."""
|
|
38
|
+
get_tracker().start()
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def wrap(client: Any) -> Any:
|
|
42
|
+
"""Manually wrap an OpenAI or Anthropic client for token tracking."""
|
|
43
|
+
return get_tracker().wrap(client)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
def report() -> None:
|
|
47
|
+
"""Print a terminal summary with token counts, cost, and context reuse."""
|
|
48
|
+
print_report(get_tracker())
|
|
49
|
+
|
|
50
|
+
|
|
51
|
+
def stop() -> None:
|
|
52
|
+
"""Stop tracking, restore SDK patches, and close the session."""
|
|
53
|
+
get_tracker().stop()
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
def checkpoint() -> None:
|
|
57
|
+
"""Manually save a checkpoint of the current conversation state."""
|
|
58
|
+
get_tracker().checkpoint()
|
|
59
|
+
|
|
60
|
+
|
|
61
|
+
def resume(session_id: int) -> list:
|
|
62
|
+
"""Resume from the latest checkpoint of a given session."""
|
|
63
|
+
return get_tracker().resume(session_id)
|
|
64
|
+
|
|
65
|
+
|
|
66
|
+
def get_session_id() -> int | None:
|
|
67
|
+
"""Get the current active session ID."""
|
|
68
|
+
return get_tracker().get_session_id()
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
def compress(
|
|
72
|
+
messages: list,
|
|
73
|
+
max_tokens: int = 2000,
|
|
74
|
+
keep_last_n: int = 4,
|
|
75
|
+
) -> dict:
|
|
76
|
+
"""
|
|
77
|
+
Compress messages to reduce token usage by 40-70%.
|
|
78
|
+
|
|
79
|
+
Usage::
|
|
80
|
+
|
|
81
|
+
result = streamctx.compress(messages, max_tokens=2000)
|
|
82
|
+
compressed = result["messages"]
|
|
83
|
+
print(result["stats"])
|
|
84
|
+
"""
|
|
85
|
+
compressed, original, after = compress_messages(
|
|
86
|
+
messages,
|
|
87
|
+
max_tokens=max_tokens,
|
|
88
|
+
keep_last_n=keep_last_n,
|
|
89
|
+
)
|
|
90
|
+
stats = get_compression_stats(original, after)
|
|
91
|
+
return {
|
|
92
|
+
"messages": compressed,
|
|
93
|
+
"stats": stats,
|
|
94
|
+
}
|
|
95
|
+
|
|
96
|
+
|
|
97
|
+
def compression_stats(original_tokens: int, compressed_tokens: int) -> dict:
|
|
98
|
+
"""Get compression statistics."""
|
|
99
|
+
return get_compression_stats(original_tokens, compressed_tokens)
|
|
100
|
+
|
|
101
|
+
|
|
102
|
+
def healing_stats() -> dict:
|
|
103
|
+
"""Get self-healing statistics."""
|
|
104
|
+
return get_tracker().healing_stats()
|
|
105
|
+
|
|
106
|
+
|
|
107
|
+
def scan(messages: list) -> dict:
|
|
108
|
+
"""
|
|
109
|
+
Scan messages for context poisoning.
|
|
110
|
+
|
|
111
|
+
Usage::
|
|
112
|
+
|
|
113
|
+
result = streamctx.scan(messages)
|
|
114
|
+
print(result["health_score"])
|
|
115
|
+
print(result["warnings"])
|
|
116
|
+
"""
|
|
117
|
+
return _detector.scan(messages)
|
|
118
|
+
|
|
119
|
+
|
|
120
|
+
def context_diff(
|
|
121
|
+
messages_a: list,
|
|
122
|
+
messages_b: list,
|
|
123
|
+
step_a: int = 0,
|
|
124
|
+
step_b: int = 0,
|
|
125
|
+
) -> dict:
|
|
126
|
+
"""
|
|
127
|
+
Compare two sets of messages to see exactly what changed.
|
|
128
|
+
|
|
129
|
+
Shows added, removed, unchanged messages and drift score.
|
|
130
|
+
|
|
131
|
+
Usage::
|
|
132
|
+
|
|
133
|
+
diff = streamctx.context_diff(
|
|
134
|
+
messages_a=step3_messages,
|
|
135
|
+
messages_b=step7_messages,
|
|
136
|
+
step_a=3,
|
|
137
|
+
step_b=7,
|
|
138
|
+
)
|
|
139
|
+
print(diff["summary"])
|
|
140
|
+
print(diff["warnings"])
|
|
141
|
+
print(diff["drift_score"])
|
|
142
|
+
|
|
143
|
+
Returns:
|
|
144
|
+
{
|
|
145
|
+
"added": list,
|
|
146
|
+
"removed": list,
|
|
147
|
+
"unchanged": list,
|
|
148
|
+
"drift_score": int 0-100,
|
|
149
|
+
"summary": str,
|
|
150
|
+
"warnings": list,
|
|
151
|
+
"token_delta": int,
|
|
152
|
+
}
|
|
153
|
+
"""
|
|
154
|
+
return _differ.diff(messages_a, messages_b, step_a, step_b)
|
|
155
|
+
|
|
156
|
+
|
|
@@ -0,0 +1,100 @@
|
|
|
1
|
+
"""Context compression engine for StreamCtx."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
def _estimate_tokens(text: str) -> int:
|
|
9
|
+
if not text:
|
|
10
|
+
return 0
|
|
11
|
+
return max(1, len(text) // 4)
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def _total_tokens(messages: list[dict[str, Any]]) -> int:
|
|
15
|
+
return sum(_estimate_tokens(m.get("content", "")) for m in messages)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def compress_messages(
|
|
19
|
+
messages: list[dict[str, Any]],
|
|
20
|
+
max_tokens: int = 2000,
|
|
21
|
+
keep_system: bool = True,
|
|
22
|
+
keep_last_n: int = 4,
|
|
23
|
+
) -> tuple[list[dict[str, Any]], int, int]:
|
|
24
|
+
if not messages:
|
|
25
|
+
return messages, 0, 0
|
|
26
|
+
|
|
27
|
+
original_tokens = _total_tokens(messages)
|
|
28
|
+
|
|
29
|
+
if original_tokens <= max_tokens:
|
|
30
|
+
return messages, original_tokens, original_tokens
|
|
31
|
+
|
|
32
|
+
# Separate system, middle, recent
|
|
33
|
+
system_msgs = [m for m in messages if m.get("role") == "system"]
|
|
34
|
+
non_system = [m for m in messages if m.get("role") != "system"]
|
|
35
|
+
|
|
36
|
+
if len(non_system) <= keep_last_n:
|
|
37
|
+
recent = non_system
|
|
38
|
+
middle = []
|
|
39
|
+
else:
|
|
40
|
+
recent = non_system[-keep_last_n:]
|
|
41
|
+
middle = non_system[:-keep_last_n]
|
|
42
|
+
|
|
43
|
+
# Compress middle into one short summary
|
|
44
|
+
compressed_middle = _compress_middle(middle)
|
|
45
|
+
|
|
46
|
+
result = system_msgs + compressed_middle + recent
|
|
47
|
+
compressed_tokens = _total_tokens(result)
|
|
48
|
+
|
|
49
|
+
return result, original_tokens, compressed_tokens
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
def _compress_middle(
|
|
53
|
+
messages: list[dict[str, Any]],
|
|
54
|
+
) -> list[dict[str, Any]]:
|
|
55
|
+
if not messages:
|
|
56
|
+
return []
|
|
57
|
+
|
|
58
|
+
# Extract key points only — very aggressive compression
|
|
59
|
+
topics = []
|
|
60
|
+
for msg in messages:
|
|
61
|
+
role = msg.get("role", "user")
|
|
62
|
+
content = msg.get("content", "")
|
|
63
|
+
if not content:
|
|
64
|
+
continue
|
|
65
|
+
# Take only first 15 chars per message
|
|
66
|
+
short = content[:15].replace("\n", " ").strip()
|
|
67
|
+
topics.append(f"{role}: {short}")
|
|
68
|
+
|
|
69
|
+
if not topics:
|
|
70
|
+
return []
|
|
71
|
+
|
|
72
|
+
# Single compressed summary message
|
|
73
|
+
summary = "Earlier context: " + " | ".join(topics)
|
|
74
|
+
|
|
75
|
+
return [{"role": "system", "content": summary}]
|
|
76
|
+
|
|
77
|
+
|
|
78
|
+
def get_compression_stats(
|
|
79
|
+
original_tokens: int,
|
|
80
|
+
compressed_tokens: int,
|
|
81
|
+
) -> dict[str, Any]:
|
|
82
|
+
if original_tokens == 0:
|
|
83
|
+
return {
|
|
84
|
+
"original_tokens": 0,
|
|
85
|
+
"compressed_tokens": 0,
|
|
86
|
+
"saved_tokens": 0,
|
|
87
|
+
"compression_pct": 0,
|
|
88
|
+
}
|
|
89
|
+
|
|
90
|
+
saved = original_tokens - compressed_tokens
|
|
91
|
+
pct = int(round(100 * saved / original_tokens))
|
|
92
|
+
|
|
93
|
+
return {
|
|
94
|
+
"original_tokens": original_tokens,
|
|
95
|
+
"compressed_tokens": compressed_tokens,
|
|
96
|
+
"saved_tokens": saved,
|
|
97
|
+
"compression_pct": pct,
|
|
98
|
+
}
|
|
99
|
+
|
|
100
|
+
|
|
@@ -0,0 +1,199 @@
|
|
|
1
|
+
"""Context Diff Engine for StreamCtx.
|
|
2
|
+
|
|
3
|
+
Compare context between any two steps to see exactly:
|
|
4
|
+
- What was added
|
|
5
|
+
- What was removed
|
|
6
|
+
- What changed
|
|
7
|
+
- Context drift score 0-100
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
from typing import Any, Optional
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def _estimate_tokens(text: str) -> int:
|
|
16
|
+
if not text:
|
|
17
|
+
return 0
|
|
18
|
+
return max(1, len(text) // 4)
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class ContextDiffer:
|
|
22
|
+
"""
|
|
23
|
+
Compare context between two steps in a session.
|
|
24
|
+
|
|
25
|
+
Shows exactly when and how context changed —
|
|
26
|
+
helping developers pinpoint where agent went wrong.
|
|
27
|
+
"""
|
|
28
|
+
|
|
29
|
+
def diff(
|
|
30
|
+
self,
|
|
31
|
+
messages_a: list[dict[str, Any]],
|
|
32
|
+
messages_b: list[dict[str, Any]],
|
|
33
|
+
step_a: int = 0,
|
|
34
|
+
step_b: int = 0,
|
|
35
|
+
) -> dict[str, Any]:
|
|
36
|
+
"""
|
|
37
|
+
Compare two sets of messages and return diff.
|
|
38
|
+
|
|
39
|
+
Returns:
|
|
40
|
+
{
|
|
41
|
+
"added": list of new messages in B,
|
|
42
|
+
"removed": list of messages only in A,
|
|
43
|
+
"unchanged": list of common messages,
|
|
44
|
+
"drift_score": int 0-100,
|
|
45
|
+
"summary": str,
|
|
46
|
+
"warnings": list[str],
|
|
47
|
+
"token_delta": int,
|
|
48
|
+
}
|
|
49
|
+
"""
|
|
50
|
+
# Convert to comparable format
|
|
51
|
+
set_a = {self._msg_key(m): m for m in messages_a}
|
|
52
|
+
set_b = {self._msg_key(m): m for m in messages_b}
|
|
53
|
+
|
|
54
|
+
keys_a = set(set_a.keys())
|
|
55
|
+
keys_b = set(set_b.keys())
|
|
56
|
+
|
|
57
|
+
added = [set_b[k] for k in keys_b - keys_a]
|
|
58
|
+
removed = [set_a[k] for k in keys_a - keys_b]
|
|
59
|
+
unchanged = [set_a[k] for k in keys_a & keys_b]
|
|
60
|
+
|
|
61
|
+
# Token delta
|
|
62
|
+
tokens_a = sum(_estimate_tokens(m.get("content", "")) for m in messages_a)
|
|
63
|
+
tokens_b = sum(_estimate_tokens(m.get("content", "")) for m in messages_b)
|
|
64
|
+
token_delta = tokens_b - tokens_a
|
|
65
|
+
|
|
66
|
+
# Drift score
|
|
67
|
+
drift_score = self._calc_drift(
|
|
68
|
+
added, removed, unchanged, messages_a, messages_b
|
|
69
|
+
)
|
|
70
|
+
|
|
71
|
+
# Warnings
|
|
72
|
+
warnings = self._get_warnings(added, removed, drift_score)
|
|
73
|
+
|
|
74
|
+
# Summary
|
|
75
|
+
summary = self._build_summary(
|
|
76
|
+
added, removed, unchanged, step_a, step_b, drift_score
|
|
77
|
+
)
|
|
78
|
+
|
|
79
|
+
return {
|
|
80
|
+
"step_a": step_a,
|
|
81
|
+
"step_b": step_b,
|
|
82
|
+
"added": added,
|
|
83
|
+
"removed": removed,
|
|
84
|
+
"unchanged": unchanged,
|
|
85
|
+
"drift_score": drift_score,
|
|
86
|
+
"summary": summary,
|
|
87
|
+
"warnings": warnings,
|
|
88
|
+
"token_delta": token_delta,
|
|
89
|
+
"tokens_a": tokens_a,
|
|
90
|
+
"tokens_b": tokens_b,
|
|
91
|
+
}
|
|
92
|
+
|
|
93
|
+
def _msg_key(self, msg: dict[str, Any]) -> str:
|
|
94
|
+
"""Unique key for a message."""
|
|
95
|
+
role = msg.get("role", "")
|
|
96
|
+
content = msg.get("content", "")
|
|
97
|
+
return f"{role}:{content}"
|
|
98
|
+
|
|
99
|
+
def _calc_drift(
|
|
100
|
+
self,
|
|
101
|
+
added: list,
|
|
102
|
+
removed: list,
|
|
103
|
+
unchanged: list,
|
|
104
|
+
messages_a: list,
|
|
105
|
+
messages_b: list,
|
|
106
|
+
) -> int:
|
|
107
|
+
"""
|
|
108
|
+
Calculate drift score 0-100.
|
|
109
|
+
0 = identical, 100 = completely different.
|
|
110
|
+
"""
|
|
111
|
+
total = len(messages_a) + len(messages_b)
|
|
112
|
+
if total == 0:
|
|
113
|
+
return 0
|
|
114
|
+
|
|
115
|
+
changed = len(added) + len(removed)
|
|
116
|
+
drift = int(round(100 * changed / total))
|
|
117
|
+
return min(100, drift)
|
|
118
|
+
|
|
119
|
+
def _get_warnings(
|
|
120
|
+
self,
|
|
121
|
+
added: list,
|
|
122
|
+
removed: list,
|
|
123
|
+
drift_score: int,
|
|
124
|
+
) -> list[str]:
|
|
125
|
+
warnings = []
|
|
126
|
+
|
|
127
|
+
# System prompt removed
|
|
128
|
+
removed_roles = [m.get("role") for m in removed]
|
|
129
|
+
added_roles = [m.get("role") for m in added]
|
|
130
|
+
|
|
131
|
+
if "system" in removed_roles:
|
|
132
|
+
warnings.append(
|
|
133
|
+
"⚠️ System prompt was REMOVED — agent may have lost instructions"
|
|
134
|
+
)
|
|
135
|
+
|
|
136
|
+
if "system" in added_roles:
|
|
137
|
+
warnings.append(
|
|
138
|
+
"⚠️ New system prompt ADDED — context instructions changed"
|
|
139
|
+
)
|
|
140
|
+
|
|
141
|
+
if drift_score >= 70:
|
|
142
|
+
warnings.append(
|
|
143
|
+
f"🚨 High drift ({drift_score}/100) — context changed significantly"
|
|
144
|
+
)
|
|
145
|
+
elif drift_score >= 40:
|
|
146
|
+
warnings.append(
|
|
147
|
+
f"⚡ Medium drift ({drift_score}/100) — monitor context closely"
|
|
148
|
+
)
|
|
149
|
+
|
|
150
|
+
# Check for contradictions
|
|
151
|
+
added_content = " ".join(m.get("content", "").lower() for m in added)
|
|
152
|
+
removed_content = " ".join(m.get("content", "").lower() for m in removed)
|
|
153
|
+
|
|
154
|
+
contradiction_pairs = [
|
|
155
|
+
("use gpt", "use claude"),
|
|
156
|
+
("enabled", "disabled"),
|
|
157
|
+
("yes", "no"),
|
|
158
|
+
("true", "false"),
|
|
159
|
+
]
|
|
160
|
+
for w1, w2 in contradiction_pairs:
|
|
161
|
+
if w1 in added_content and w2 in removed_content:
|
|
162
|
+
warnings.append(
|
|
163
|
+
f"⚠️ Contradiction: '{w1}' added but '{w2}' removed"
|
|
164
|
+
)
|
|
165
|
+
elif w2 in added_content and w1 in removed_content:
|
|
166
|
+
warnings.append(
|
|
167
|
+
f"⚠️ Contradiction: '{w2}' added but '{w1}' removed"
|
|
168
|
+
)
|
|
169
|
+
|
|
170
|
+
return warnings
|
|
171
|
+
|
|
172
|
+
def _build_summary(
|
|
173
|
+
self,
|
|
174
|
+
added: list,
|
|
175
|
+
removed: list,
|
|
176
|
+
unchanged: list,
|
|
177
|
+
step_a: int,
|
|
178
|
+
step_b: int,
|
|
179
|
+
drift_score: int,
|
|
180
|
+
) -> str:
|
|
181
|
+
if drift_score == 0:
|
|
182
|
+
return f"✅ Step {step_a} → {step_b}: No context change"
|
|
183
|
+
elif drift_score < 40:
|
|
184
|
+
return (
|
|
185
|
+
f"✅ Step {step_a} → {step_b}: Minor changes "
|
|
186
|
+
f"(+{len(added)} -{len(removed)} messages, drift: {drift_score}/100)"
|
|
187
|
+
)
|
|
188
|
+
elif drift_score < 70:
|
|
189
|
+
return (
|
|
190
|
+
f"⚡ Step {step_a} → {step_b}: Moderate drift "
|
|
191
|
+
f"(+{len(added)} -{len(removed)} messages, drift: {drift_score}/100)"
|
|
192
|
+
)
|
|
193
|
+
else:
|
|
194
|
+
return (
|
|
195
|
+
f"🚨 Step {step_a} → {step_b}: HIGH DRIFT "
|
|
196
|
+
f"(+{len(added)} -{len(removed)} messages, drift: {drift_score}/100)"
|
|
197
|
+
)
|
|
198
|
+
|
|
199
|
+
|
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
"""Self-healing engine for StreamCtx."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
from typing import Any, Optional
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
class SelfHealingEngine:
|
|
9
|
+
"""
|
|
10
|
+
Detects failed LLM calls and recovers using last valid checkpoint.
|
|
11
|
+
|
|
12
|
+
When a call fails:
|
|
13
|
+
1. Catches the exception
|
|
14
|
+
2. Loads last valid checkpoint messages
|
|
15
|
+
3. Injects recovery context
|
|
16
|
+
4. Retries the call automatically
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
def __init__(self) -> None:
|
|
20
|
+
self._last_valid_messages: list[dict[str, Any]] = []
|
|
21
|
+
self._last_valid_response: Any = None
|
|
22
|
+
self._failure_count: int = 0
|
|
23
|
+
self._recovery_count: int = 0
|
|
24
|
+
|
|
25
|
+
def record_success(
|
|
26
|
+
self,
|
|
27
|
+
messages: list[dict[str, Any]],
|
|
28
|
+
response: Any,
|
|
29
|
+
) -> None:
|
|
30
|
+
"""Save last successful call for recovery."""
|
|
31
|
+
self._last_valid_messages = list(messages)
|
|
32
|
+
self._last_valid_response = response
|
|
33
|
+
|
|
34
|
+
def record_failure(self) -> None:
|
|
35
|
+
"""Track failure count."""
|
|
36
|
+
self._failure_count += 1
|
|
37
|
+
|
|
38
|
+
def can_heal(self) -> bool:
|
|
39
|
+
"""Check if we have valid context to recover from."""
|
|
40
|
+
return len(self._last_valid_messages) > 0
|
|
41
|
+
|
|
42
|
+
def get_recovery_messages(
|
|
43
|
+
self,
|
|
44
|
+
failed_messages: list[dict[str, Any]],
|
|
45
|
+
) -> list[dict[str, Any]]:
|
|
46
|
+
"""
|
|
47
|
+
Build recovery message list by injecting last valid context.
|
|
48
|
+
|
|
49
|
+
Strategy:
|
|
50
|
+
- Take last valid messages as base
|
|
51
|
+
- Add recovery notice
|
|
52
|
+
- Append the last user message from failed call
|
|
53
|
+
"""
|
|
54
|
+
if not self._last_valid_messages:
|
|
55
|
+
return failed_messages
|
|
56
|
+
|
|
57
|
+
# Get last user message from failed call
|
|
58
|
+
last_user_msg = None
|
|
59
|
+
for msg in reversed(failed_messages):
|
|
60
|
+
if msg.get("role") == "user":
|
|
61
|
+
last_user_msg = msg
|
|
62
|
+
break
|
|
63
|
+
|
|
64
|
+
# Build recovery context
|
|
65
|
+
recovery = list(self._last_valid_messages)
|
|
66
|
+
|
|
67
|
+
# Add recovery notice as system message
|
|
68
|
+
recovery_notice = {
|
|
69
|
+
"role": "system",
|
|
70
|
+
"content": (
|
|
71
|
+
"Note: Previous response was interrupted. "
|
|
72
|
+
"Continuing from last valid context."
|
|
73
|
+
),
|
|
74
|
+
}
|
|
75
|
+
recovery.append(recovery_notice)
|
|
76
|
+
|
|
77
|
+
# Re-add the last user message
|
|
78
|
+
if last_user_msg:
|
|
79
|
+
recovery.append(last_user_msg)
|
|
80
|
+
|
|
81
|
+
self._recovery_count += 1
|
|
82
|
+
return recovery
|
|
83
|
+
|
|
84
|
+
def attempt_heal(
|
|
85
|
+
self,
|
|
86
|
+
fn: Any,
|
|
87
|
+
failed_messages: list[dict[str, Any]],
|
|
88
|
+
kwargs: dict[str, Any],
|
|
89
|
+
max_retries: int = 2,
|
|
90
|
+
) -> Any:
|
|
91
|
+
"""
|
|
92
|
+
Attempt to recover from a failed LLM call.
|
|
93
|
+
|
|
94
|
+
Retries with last valid context injected.
|
|
95
|
+
"""
|
|
96
|
+
if not self.can_heal():
|
|
97
|
+
raise RuntimeError("Self-healing failed: no valid checkpoint available.")
|
|
98
|
+
|
|
99
|
+
recovery_messages = self.get_recovery_messages(failed_messages)
|
|
100
|
+
|
|
101
|
+
for attempt in range(max_retries):
|
|
102
|
+
try:
|
|
103
|
+
# Update messages with recovery context
|
|
104
|
+
kwargs_copy = dict(kwargs)
|
|
105
|
+
kwargs_copy["messages"] = recovery_messages
|
|
106
|
+
response = fn(**kwargs_copy)
|
|
107
|
+
self.record_success(recovery_messages, response)
|
|
108
|
+
return response
|
|
109
|
+
except Exception:
|
|
110
|
+
if attempt == max_retries - 1:
|
|
111
|
+
raise
|
|
112
|
+
continue
|
|
113
|
+
|
|
114
|
+
raise RuntimeError("Self-healing exhausted all retries.")
|
|
115
|
+
|
|
116
|
+
def get_stats(self) -> dict[str, Any]:
|
|
117
|
+
"""Return healing statistics."""
|
|
118
|
+
return {
|
|
119
|
+
"failure_count": self._failure_count,
|
|
120
|
+
"recovery_count": self._recovery_count,
|
|
121
|
+
"has_valid_context": self.can_heal(),
|
|
122
|
+
}
|
|
123
|
+
|
|
124
|
+
|