Repository navigation
Expand file tree
/
Copy path_testing_models.py
More file actions
402 lines (342 loc) · 15.4 KB
/
Copy path_testing_models.py
File metadata and controls
402 lines (342 loc) · 15.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
"""Internal fake chat models for local integration tests.
The tool-binding base these build on (`_fake_models._ToolBindingFakeModel`) is
factored out into a use-neutral module so the `dcode tools list` enumeration
path can reuse it without importing this test-named module.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Any
from langchain_core.messages import AIMessage, BaseMessage
from langchain_core.outputs import ChatGeneration, ChatResult
from deepagents_code._fake_models import _ToolBindingFakeModel
if TYPE_CHECKING:
from collections.abc import Callable
from langchain_core.callbacks import CallbackManagerForLLMRun
DCA_TEST_OFFLOAD_GATE_ENV = "DCA_TEST_OFFLOAD_GATE_DIR"
"""Env var pointing at a directory used to gate summary generation.
When set, a summary request writes `<dir>/entered` and then polls for
`<dir>/release` before replying. File-based so the test process can hold the
server's compaction model call open without sharing Python state across the
server subprocess boundary. Only summary prompts are gated; ordinary turns pass
through, which is what lets a test launch a concurrent run *while* `/offload`
is blocked here.
"""
# Prompt markers that drive `ToolCallingIntegrationChatModel`. Each marker is the
# full token (including the trailing `=`); the file path follows on the same line,
# e.g. `DCA_TEST_WRITE_FILE=/tmp/out.txt`. These are the single source of truth
# shared with the integration tests, so the model and tests cannot drift apart.
DCA_TEST_WRITE_FILE_MARKER = "DCA_TEST_WRITE_FILE="
DCA_TEST_DELEGATE_WRITE_MARKER = "DCA_TEST_DELEGATE_WRITE="
DCA_SUBAGENT_WRITE_FILE_MARKER = "DCA_SUBAGENT_WRITE_FILE="
DCA_TEST_GOAL_CRITERIA_MARKER = "DCA_TEST_GOAL_CRITERIA="
# Distinct file contents per write path, so a test asserting on file content can
# prove which branch executed — in particular that subagent mode delegated through
# the `task` tool rather than writing directly.
TOP_LEVEL_WRITE_CONTENT = "auto-approved"
SUBAGENT_WRITE_CONTENT = "auto-approved-subagent"
class DeterministicIntegrationChatModel(_ToolBindingFakeModel):
"""Deterministic chat model for integration tests.
This subclasses `_ToolBindingFakeModel` (itself a `GenericFakeChatModel`) so
the implementation stays aligned with the core fake-chat-model test surface,
while overriding generation to remain prompt-driven and restart-safe for real
CLI server integration tests.
Why the existing `langchain_core` fakes cannot be reused here:
1. Every core fake (`GenericFakeChatModel`, `FakeListChatModel`,
`FakeMessagesListChatModel`) pops from an iterator or cycles an index —
the actual prompt is ignored. App integration tests start and stop the
server process, which resets in-memory state. An iterator-based model
either raises `StopIteration` or replays from the beginning after a
restart, producing wrong or missing responses. This model derives output
solely from the prompt text, so identical input always produces
identical output regardless of process lifecycle.
2. The agent runtime calls `model.bind_tools(schemas)` during
initialization. A bare `GenericFakeChatModel` inherits
`BaseChatModel.bind_tools`, which raises `NotImplementedError` in any
agent-loop context. The inherited `_ToolBindingFakeModel` supplies a
no-op passthrough.
3. The app server reads `model.profile` for capability negotiation (e.g.
`tool_calling`, `max_input_tokens`). A bare fake's `profile` is `None`,
causing silent misconfiguration at runtime. The inherited
`_ToolBindingFakeModel` supplies a minimal profile.
Additionally, the compact middleware issues summarization prompts mid-
conversation. A list-based model cannot distinguish these from normal user
turns without pre-knowledge of exact call ordering, whereas this model
detects summary requests by inspecting the prompt content.
"""
model: str = "fake"
# `messages`, `profile`, and the `bind_tools` passthrough are inherited from
# `_ToolBindingFakeModel`; this model adds only prompt-driven generation.
def _generate(
self,
messages: list[BaseMessage],
stop: list[str] | None = None, # noqa: ARG002
run_manager: CallbackManagerForLLMRun | None = None, # noqa: ARG002
**kwargs: Any, # noqa: ARG002
) -> ChatResult:
"""Produce a deterministic reply derived from the prompt text.
Returns:
A single-message `ChatResult` with deterministic content.
"""
prompt = "\n".join(
text
for message in messages
if (text := self._stringify_message(message)).strip()
)
if self._looks_like_summary_request(prompt):
self._wait_at_summary_gate()
content = "integration summary"
else:
excerpt = " ".join(prompt.split()[-18:])
if excerpt:
content = f"integration reply: {excerpt}"
else:
content = "integration reply"
return ChatResult(
generations=[
ChatGeneration(
message=AIMessage(
content=content,
usage_metadata={
"input_tokens": 100,
"output_tokens": 20,
"total_tokens": 120,
},
)
)
]
)
@property
def _llm_type(self) -> str:
"""LangChain model type identifier."""
return "deterministic-integration"
@staticmethod
def _wait_at_summary_gate() -> None:
"""Hold the summary call open until the test releases it.
No-op unless `DCA_TEST_OFFLOAD_GATE_DIR` names a directory. When set,
write `<dir>/entered` (the test's signal that the offload operation is
mid-summary) and then poll for `<dir>/release`. Every summary request
rewrites the marker, so the test reads it as "a summary is in flight"
rather than "the first summary started". Bounded so a crashed test
cannot wedge the server subprocess indefinitely.
Raises:
TimeoutError: If the gate is not released within 120 seconds.
"""
import os
import time
from pathlib import Path
gate_dir = os.environ.get(DCA_TEST_OFFLOAD_GATE_ENV)
if not gate_dir:
return
gate = Path(gate_dir)
(gate / "entered").write_text("1")
deadline = time.monotonic() + 120
while not (gate / "release").exists():
if time.monotonic() > deadline:
msg = (
"Offload test gate was never released; refusing to block "
"the server summary call forever."
)
raise TimeoutError(msg)
time.sleep(0.05)
@staticmethod
def _stringify_message(message: BaseMessage) -> str:
"""Flatten message content into plain text for deterministic responses.
Returns:
Plain-text content extracted from the message.
"""
content = message.content
if isinstance(content, str):
return content
if isinstance(content, list):
parts: list[str] = []
for block in content:
if isinstance(block, str):
parts.append(block)
elif isinstance(block, dict) and block.get("type") == "text":
text = block.get("text")
if isinstance(text, str):
parts.append(text)
return " ".join(parts)
return str(content)
@staticmethod
def _looks_like_summary_request(prompt: str) -> bool:
"""Detect the middleware's summary-generation prompt.
Returns:
`True` when the prompt appears to be a summarization request.
"""
lowered = prompt.lower()
return (
"messages to summarize" in lowered
or "condense the following conversation" in lowered
or "<summary>" in lowered
)
def _extract_marker_path(prompt: str, marker: str) -> str:
"""Extract the file path that follows a prompt marker on the same line.
Args:
prompt: The flattened prompt text.
marker: The marker token (including its trailing `=`) to locate.
Returns:
The stripped file path immediately following the marker.
Raises:
ValueError: If the marker is present but not followed by a path, so a
malformed test prompt fails loudly here instead of silently
degrading to a `"done"` reply or raising an opaque `IndexError`.
"""
_, _, tail = prompt.partition(marker)
lines = tail.splitlines()
file_path = lines[0].strip() if lines else ""
if not file_path:
msg = (
f"Test model saw marker {marker!r} but found no file path after it; "
f"check the integration-test prompt construction."
)
raise ValueError(msg)
return file_path
def _tool_call_result(name: str, args: dict[str, Any], call_id: str) -> ChatResult:
"""Build a single-tool-call `ChatResult`.
Returns:
A `ChatResult` wrapping an `AIMessage` with exactly one tool call.
"""
return ChatResult(
generations=[
ChatGeneration(
message=AIMessage(
content="",
tool_calls=[
{
"name": name,
"args": args,
"id": call_id,
"type": "tool_call",
}
],
)
)
]
)
class ToolCallingIntegrationChatModel(DeterministicIntegrationChatModel):
"""Deterministic tool-calling model for auto-approve integration tests.
Generation is driven entirely by prompt markers (the module-level `DCA_*`
constants), so output is restart-safe and independent of call ordering — the
same rationale as the parent `DeterministicIntegrationChatModel`:
- `DCA_TEST_WRITE_FILE=<path>` emits a top-level `write_file` call.
- `DCA_TEST_DELEGATE_WRITE=<path>` emits a `task` call delegating to the
`general-purpose` subagent, whose prompt then carries
`DCA_SUBAGENT_WRITE_FILE=<path>` to trigger the subagent's `write_file`.
- `DCA_SUBAGENT_WRITE_FILE=<path>` emits the subagent's `write_file` call.
Each marker fires only on the agent's first turn (no prior `ToolMessage`),
so once the tool result returns the model replies with a plain `"done"` and
the agent loop terminates instead of re-issuing the tool call.
"""
# Only `_generate` is overridden; the inherited `_stream` would bypass this
# marker dispatch entirely, so streaming must be disabled.
disable_streaming: bool = True
def _generate(
self,
messages: list[BaseMessage],
stop: list[str] | None = None, # noqa: ARG002
run_manager: CallbackManagerForLLMRun | None = None, # noqa: ARG002
**kwargs: Any, # noqa: ARG002
) -> ChatResult:
"""Emit a deterministic tool call (or terminal reply) from prompt markers.
The `has_tool_result` guard ensures each marker fires only on the
agent's first turn; once a `ToolMessage` is present the model returns
`"done"` so the agent loop terminates.
A recognized marker with no file path raises `ValueError` (via
`_extract_marker_path`) rather than degrading to `"done"`.
Returns:
A single-message `ChatResult`: an `AIMessage` carrying a `task` or
`write_file` tool call when a marker matches and no tool result is
present yet, otherwise a plain `"done"` reply.
"""
prompt = "\n".join(
text
for message in messages
if (text := self._stringify_message(message)).strip()
)
has_tool_result = any(message.type == "tool" for message in messages)
if not has_tool_result:
for marker, build_tool_call in self._marker_dispatch():
if marker in prompt:
name, args, call_id = build_tool_call(
_extract_marker_path(prompt, marker)
)
return _tool_call_result(name, args, call_id)
return ChatResult(
generations=[ChatGeneration(message=AIMessage(content="done"))]
)
@staticmethod
def _marker_dispatch() -> tuple[
tuple[str, Callable[[str], tuple[str, dict[str, Any], str]]], ...
]:
"""Return ordered `(marker, tool-call builder)` pairs.
The delegate marker is checked before the plain write markers because
its emitted `task` description embeds `DCA_SUBAGENT_WRITE_FILE=`; the
first marker found in the prompt wins.
Returns:
Marker-to-builder pairs in precedence order. Each builder maps an
extracted file path to a `(tool_name, args, call_id)` triple.
"""
return (
(
DCA_TEST_DELEGATE_WRITE_MARKER,
lambda path: (
"task",
{
"description": f"{DCA_SUBAGENT_WRITE_FILE_MARKER}{path}",
"subagent_type": "general-purpose",
},
"call_task",
),
),
(
DCA_SUBAGENT_WRITE_FILE_MARKER,
lambda path: (
"write_file",
{"file_path": path, "content": SUBAGENT_WRITE_CONTENT},
"call_write_file",
),
),
(
DCA_TEST_WRITE_FILE_MARKER,
lambda path: (
"write_file",
{"file_path": path, "content": TOP_LEVEL_WRITE_CONTENT},
"call_write_file",
),
),
)
class GoalCriteriaIntegrationChatModel(DeterministicIntegrationChatModel):
"""Exercise nested criteria generation with a repository read."""
disable_streaming: bool = True
def _generate(
self,
messages: list[BaseMessage],
stop: list[str] | None = None, # noqa: ARG002
run_manager: CallbackManagerForLLMRun | None = None, # noqa: ARG002
**kwargs: Any, # noqa: ARG002
) -> ChatResult:
"""Read the marked file, then return a structured goal proposal.
Returns:
A repository tool call followed by a `GoalProposal` tool call.
"""
prompt = "\n".join(
text
for message in messages
if (text := self._stringify_message(message)).strip()
)
has_tool_result = any(message.type == "tool" for message in messages)
if DCA_TEST_GOAL_CRITERIA_MARKER in prompt and not has_tool_result:
path = _extract_marker_path(prompt, DCA_TEST_GOAL_CRITERIA_MARKER)
return _tool_call_result(
"read_file",
{"file_path": path, "limit": 20},
"call_goal_read",
)
return _tool_call_result(
"GoalProposal",
{
"objective": "verify server-side criteria generation",
"criteria": "- server repository context is available",
},
"call_goal_proposal",
)