Context Doctor polish
Improve code quality of context doctor. Add handling of Reasoning kvps. Bump json_repair version
linkliti committed
Aug 26, 2026 at 13:14 UTC
51a02bd4f530a2efcf85035f71b710fa3b008844
6 files changed
+123
-77
plugins/_context_doctor/AGENTS.md
+1
-1
@@ -16,7 +16,7 @@
16
17
- Repaired and fallback JSON is always minified.
18
- Nonempty non-tool output becomes `{"thoughts":[raw]}`; XML-like output becomes `{}` only when suppression is enabled.
19
-- Log kvps and heading always reflect transformed output; `update_log` controls only View Details content.
19
+- Log kvps retain streamed `reasoning` and add transformed output; `update_log` controls only View Details content.
20
- A repaired `response` tool call refreshes the response log item when streaming did not create it.
21
- Runtime setup reads the `json_repair` pin from root `requirements.txt`; do not duplicate its version in plugin code.
22
plugins/_context_doctor/README.md
+1
-1
@@ -6,5 +6,5 @@ Repairs model tool call JSON before sending it to history and tool processing.
6
7
- Uses `json_repair` after native parsing has completed.
8
- Accepts only complete Agent Zero tool calls (`tool_name` and object `tool_args`).
9
-- Stores and displays repaired tool calls as compact JSON.
9
+- Stores repaired tool calls as compact JSON; log kvps retain streamed reasoning and add transformed fields.
10
- Optionally replaces XML-like output with `{}` in cases where model uses native XML tool calls.
plugins/_context_doctor/extensions/python/message_loop_result/_10_context_doctor.py
+16
-3
@@ -7,7 +7,10 @@ from typing import Any, override
7
8
from helpers.extension import Extension
9
from helpers.plugins import get_plugin_config
10
-from plugins._context_doctor.helpers.context_doctor import transform_response, update_log_item
10
+from plugins._context_doctor.helpers.context_doctor import (
11
+ transform_response,
12
+ update_log_item,
13
+)
14
15
16
class ContextDoctor(Extension):
@@ -16,21 +19,27 @@ class ContextDoctor(Extension):
19
if not self.agent or not isinstance(result_data, dict):
20
return
21
22
+ # Extract LLM response
23
llm_result = result_data.get("llm_result")
24
response = getattr(llm_result, "response", None)
25
if not isinstance(response, str) or not response:
26
return
27
28
+ # Repair response before default processing
29
config = get_plugin_config("_context_doctor", agent=self.agent) or {}
30
transformed = transform_response(
31
response, suppress_xml=config.get("suppress_xml", True)
32
)
33
llm_result.response = transformed
34
30
- params = getattr(getattr(self.agent, "loop_data", None), "params_temporary", None)
35
+ # Suppressed XML does not update the generating log
36
+ params = getattr(
37
+ getattr(self.agent, "loop_data", None), "params_temporary", None
38
+ )
39
if not isinstance(params, dict) or transformed == "{}":
40
return
41
42
+ # Refresh log item
43
log_item = params.get("log_item_generating")
44
if log_item is not None:
45
update_log_item(
@@ -41,15 +50,19 @@ class ContextDoctor(Extension):
50
raw_response=response,
51
)
52
53
+ # Extract final text from a response tool call
54
try:
55
parsed = json.loads(transformed)
56
tool_args = parsed.get("tool_args", {}) if isinstance(parsed, dict) else {}
47
- response_text = tool_args.get("text") if parsed.get("tool_name") == "response" else None
57
+ response_text = (
58
+ tool_args.get("text") if parsed.get("tool_name") == "response" else None
59
+ )
60
except (AttributeError, TypeError, ValueError):
61
response_text = None
62
if not isinstance(response_text, str) or not response_text.strip():
63
return
64
65
+ # Create a response log item when streaming did not create one
66
response_item = params.get("log_item_response")
67
if response_item is None and log_item is not None:
68
response_item = self.agent.context.log.log(
plugins/_context_doctor/helpers/context_doctor.py
+92
-68
@@ -1,12 +1,34 @@
1
-"""Repair completed model output into compact JSON."""
1
+"""Repair completed model output and refresh its log entry."""
2
3
from __future__ import annotations
4
5
import json
6
+from collections.abc import Iterable
7
from typing import Any
8
9
+from json_repair import repair_json
10
+
11
+from plugins._context_doctor.helpers.json_repair_patch import apply_patch
12
+
13
+
14
+# One-time, idempotent parser patch on plugin import.
15
+apply_patch()
16
+
17
+
18
+_A0_SALVAGE_SCHEMA = {
19
+ "type": "object",
20
+ "properties": {
21
+ "thoughts": {"type": "array", "items": {"type": "string"}},
22
+ "headline": {"type": "string"},
23
+ "tool_name": {"type": "string"},
24
+ "tool_args": {"type": "object"},
25
+ },
26
+ "required": [],
27
+}
28
+
29
30
def _is_tool_call(value: Any) -> bool:
31
+ """Check required tool-call fields and optional display fields."""
32
return (
33
isinstance(value, dict)
34
and isinstance(value.get("tool_name"), str)
@@ -22,80 +44,70 @@ def _is_tool_call(value: Any) -> bool:
44
)
45
46
25
-_A0_SALVAGE_SCHEMA = {
26
- "type": "object",
27
- "properties": {
28
- "thoughts": {"type": "array", "items": {"type": "string"}},
29
- "headline": {"type": "string"},
30
- "tool_name": {"type": "string"},
31
- "tool_args": {"type": "object"},
32
- },
33
- "required": [],
34
-}
47
+def _is_partial_response(value: Any) -> bool:
48
+ """Check whether repair found thoughts or a headline without a tool call."""
49
+ return isinstance(value, dict) and any(
50
+ key in value for key in ("thoughts", "headline")
51
+ )
52
+
53
54
+def _objects(value: Any) -> Iterable[dict[str, Any]]:
55
+ """Yield dictionary candidates from one repair result."""
56
+ if isinstance(value, dict):
57
+ yield value
58
+ elif isinstance(value, list):
59
+ yield from (item for item in value if isinstance(item, dict))
60
37
-def _a0_completeness_score(value: dict[str, Any]) -> int:
38
- return sum(field in value for field in ("thoughts", "headline", "tool_name", "tool_args")) + bool(
39
- value.get("tool_args")
61
+
62
+def _score_tool_call(value: dict[str, Any]) -> int:
63
+ """Rank candidates by populated A0 tool-call fields."""
64
+ return sum(
65
+ key in value for key in ("thoughts", "headline", "tool_name", "tool_args")
66
+ ) + bool(value.get("tool_args"))
67
+
68
+
69
+def _repair(response: str) -> tuple[dict[str, Any] | None, dict[str, Any] | None]:
70
+ """Repair response and select its best tool or partial result."""
71
+ try:
72
+ salvage = repair_json(
73
+ response,
74
+ return_objects=True,
75
+ schema=_A0_SALVAGE_SCHEMA,
76
+ schema_repair_mode="salvage",
77
+ )
78
+ except Exception:
79
+ salvage = None
80
+ try:
81
+ unstructured = repair_json(response, return_objects=True)
82
+ except Exception:
83
+ unstructured = None
84
+
85
+ candidates = [*_objects(salvage), *_objects(unstructured)]
86
+ tool_calls = [candidate for candidate in candidates if _is_tool_call(candidate)]
87
+ if tool_calls:
88
+ return max(tool_calls, key=_score_tool_call), None
89
+ return None, next(
90
+ (candidate for candidate in candidates if _is_partial_response(candidate)), None
91
)
92
93
43
-def transform_response(response: str, *, suppress_xml: bool) -> str:
44
- """Return compact repaired tool JSON or a compact raw-text fallback."""
45
- if response:
46
- try:
47
- from plugins._context_doctor.helpers.json_repair_patch import apply_patch
48
- from json_repair import repair_json
49
-
50
- apply_patch()
51
- try:
52
- salvage_repaired = repair_json(
53
- response,
54
- return_objects=True,
55
- schema=_A0_SALVAGE_SCHEMA,
56
- schema_repair_mode="salvage",
57
- )
58
- except Exception:
59
- salvage_repaired = repair_json(response, return_objects=True)
60
-
61
- candidates = (
62
- salvage_repaired
63
- if isinstance(salvage_repaired, list)
64
- else [salvage_repaired]
65
- )
66
- try:
67
- no_schema = repair_json(response, return_objects=True)
68
- except Exception:
69
- no_schema = None
70
- if isinstance(no_schema, list):
71
- valid = [item for item in no_schema if _is_tool_call(item)]
72
- if len(valid) > 1 or (
73
- len(valid) == 1 and not _is_tool_call(salvage_repaired)
74
- ):
75
- candidates = no_schema
76
- repaired = max(
77
- (item for item in candidates if _is_tool_call(item)),
78
- key=_a0_completeness_score,
79
- default=None,
80
- )
81
- except Exception:
82
- repaired = None
83
- salvage_repaired = None
94
+def _compact_json(value: dict[str, Any]) -> str:
95
+ """Serialize JSON without ASCII escaping or whitespace."""
96
+ return json.dumps(value, ensure_ascii=False, separators=(",", ":"))
97
85
- if _is_tool_call(repaired):
86
- return json.dumps(repaired, ensure_ascii=False, separators=(",", ":"))
98
88
- for candidate in (repaired, salvage_repaired):
89
- if isinstance(candidate, dict) and (
90
- "thoughts" in candidate or "headline" in candidate
91
- ):
92
- return json.dumps(
93
- candidate, ensure_ascii=False, separators=(",", ":")
94
- )
99
+def transform_response(response: str, *, suppress_xml: bool) -> str:
100
+ """Repair model output, falling back to compact thoughts JSON for raw text."""
101
+ if response:
102
+ tool_call, partial_response = _repair(response)
103
+ if tool_call is not None:
104
+ return _compact_json(tool_call)
105
+ if partial_response is not None:
106
+ return _compact_json(partial_response)
107
108
if suppress_xml and "<" in response and ">" in response:
109
return "{}"
98
- return json.dumps({"thoughts": [response]}, ensure_ascii=False, separators=(",", ":"))
110
+ return _compact_json({"thoughts": [response]})
111
112
113
def update_log_item(
@@ -106,19 +118,31 @@ def update_log_item(
118
update_log: bool,
119
raw_response: str,
120
) -> None:
109
- """Refresh log fields from transformed JSON; optionally replace raw details."""
121
+ """Update generated log fields while retaining streamed reasoning."""
122
try:
123
parsed = json.loads(response)
124
if not isinstance(parsed, dict):
125
return
126
+
127
+ current_kvps = getattr(log_item, "kvps", None)
128
+ kvps = (
129
+ {"reasoning": current_kvps["reasoning"]}
130
+ if isinstance(current_kvps, dict) and "reasoning" in current_kvps
131
+ else {}
132
+ )
133
+ kvps.update(parsed)
134
+
135
heading = parsed.get("headline")
136
if not isinstance(heading, str) or not heading:
137
tool_name = parsed.get("tool_name")
138
heading = f"Using {tool_name}" if isinstance(tool_name, str) else ""
118
- kwargs: dict[str, Any] = {"kvps": parsed}
139
+
140
+ kwargs: dict[str, Any] = {
141
+ "content": response if update_log else raw_response,
142
+ "kvps": kvps,
143
+ }
144
if heading:
145
kwargs["heading"] = f"{getattr(agent, 'agent_name', 'A0')}: {heading}"
121
- kwargs["content"] = response if update_log else raw_response
146
log_item.update(**kwargs)
147
except (AttributeError, TypeError, ValueError):
148
pass
plugins/_context_doctor/tests/test_context_doctor.py
+12
-3
@@ -3,7 +3,10 @@ from types import SimpleNamespace
3
from plugins._context_doctor.extensions.python.message_loop_result._10_context_doctor import (
4
ContextDoctor,
5
)
6
-from plugins._context_doctor.helpers.context_doctor import transform_response, update_log_item
6
+from plugins._context_doctor.helpers.context_doctor import (
7
+ transform_response,
8
+ update_log_item,
9
+)
10
11
12
def test_repairs_and_minifies_tool_call():
@@ -57,7 +60,10 @@ def test_suppresses_xml_when_enabled():
60
61
62
def test_updates_log_kvps_and_heading_while_preserving_raw_content():
60
- log_item = SimpleNamespace(update=lambda **kwargs: setattr(log_item, "data", kwargs))
63
+ log_item = SimpleNamespace(
64
+ kvps={"reasoning": "because"},
65
+ update=lambda **kwargs: setattr(log_item, "data", kwargs),
66
+ )
67
raw = '{"tool_name":"response","tool_args":{"text":"ok",},}'
68
repaired = '{"headline":"Done","tool_name":"response","tool_args":{"text":"ok"}}'
69
@@ -70,6 +76,7 @@ def test_updates_log_kvps_and_heading_while_preserving_raw_content():
76
)
77
78
assert log_item.data["content"] == raw
79
+ assert log_item.data["kvps"]["reasoning"] == "because"
80
assert log_item.data["kvps"]["tool_name"] == "response"
81
assert log_item.data["heading"] == "A0: Done"
82
@@ -85,7 +92,9 @@ def test_extension_replaces_result_refreshes_log_and_response_item(monkeypatch):
92
log_item = SimpleNamespace(
93
id="generating", update=lambda **kwargs: setattr(log_item, "data", kwargs)
94
)
88
- response_item = SimpleNamespace(update=lambda **kwargs: setattr(response_item, "data", kwargs))
95
+ response_item = SimpleNamespace(
96
+ update=lambda **kwargs: setattr(response_item, "data", kwargs)
97
+ )
98
context = SimpleNamespace(log=SimpleNamespace(log=lambda **kwargs: response_item))
99
agent = SimpleNamespace(
100
agent_name="A0",
requirements.txt
+1
-1
@@ -12,7 +12,7 @@ flaredantic==0.1.5
12
GitPython==3.1.43
13
giturlparse==0.14.0
14
inputimeout==1.0.4
15
-json_repair==0.63.3
15
+json_repair==0.63.4
16
kokoro>=0.9.2
17
simpleeval==1.0.3
18
langchain-core==0.3.49