zettcode-engine 0.3.4__py3-none-any.whl
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.
- zettcode/_compat.py +35 -0
- zettcode/budget.py +24 -0
- zettcode/config.py +1088 -0
- zettcode/config_edit.py +457 -0
- zettcode/engine/__init__.py +39 -0
- zettcode/engine/agent/__init__.py +81 -0
- zettcode/engine/agent/base.py +60 -0
- zettcode/engine/agent/context.py +166 -0
- zettcode/engine/agent/exports.py +73 -0
- zettcode/engine/agent/interaction.py +54 -0
- zettcode/engine/agent/sessions.py +218 -0
- zettcode/engine/agent/settings.py +269 -0
- zettcode/engine/agent/turns.py +308 -0
- zettcode/engine/core.py +819 -0
- zettcode/engine/dialogue/__init__.py +6 -0
- zettcode/engine/dialogue/ask.py +113 -0
- zettcode/engine/dialogue/side.py +197 -0
- zettcode/engine/dialogue/title.py +80 -0
- zettcode/engine/export/__init__.py +421 -0
- zettcode/engine/export/trace/script.js +18 -0
- zettcode/engine/export/trace/style.css +136 -0
- zettcode/engine/export/trace/template.html +31 -0
- zettcode/engine/extensions/__init__.py +8 -0
- zettcode/engine/extensions/capabilities.py +68 -0
- zettcode/engine/extensions/compaction.py +259 -0
- zettcode/engine/extensions/context.py +521 -0
- zettcode/engine/extensions/files.py +101 -0
- zettcode/engine/extensions/guidance.py +17 -0
- zettcode/engine/extensions/harness.py +208 -0
- zettcode/engine/extensions/harness_switch.py +18 -0
- zettcode/engine/extensions/mcp_notice.py +49 -0
- zettcode/engine/extensions/relay.py +93 -0
- zettcode/engine/extensions/subagents.py +391 -0
- zettcode/engine/extensions/tool_output.py +153 -0
- zettcode/engine/extensions/truncation.py +156 -0
- zettcode/engine/extensions/usage.py +304 -0
- zettcode/engine/harness_paths.py +74 -0
- zettcode/engine/permission/__init__.py +197 -0
- zettcode/engine/permission/approval.py +66 -0
- zettcode/engine/permission/readonly.py +435 -0
- zettcode/engine/permission/sandbox.py +140 -0
- zettcode/engine/pipeline.py +191 -0
- zettcode/engine/plugins/__init__.py +72 -0
- zettcode/engine/plugins/commands.py +152 -0
- zettcode/engine/plugins/context.py +251 -0
- zettcode/engine/plugins/extension.py +183 -0
- zettcode/engine/plugins/groups.py +74 -0
- zettcode/engine/plugins/loader.py +367 -0
- zettcode/engine/plugins/mixins.py +156 -0
- zettcode/engine/plugins/plugin.py +190 -0
- zettcode/engine/plugins/services.py +68 -0
- zettcode/engine/plugins/settings.py +117 -0
- zettcode/engine/plugins/state.py +64 -0
- zettcode/engine/plugins/switches.py +553 -0
- zettcode/engine/py.typed +0 -0
- zettcode/engine/server_tools.py +179 -0
- zettcode/engine/storage/__init__.py +51 -0
- zettcode/engine/storage/layout.py +161 -0
- zettcode/engine/storage/ledger.py +659 -0
- zettcode/engine/storage/locks.py +182 -0
- zettcode/engine/storage/metadata.py +278 -0
- zettcode/engine/storage/records.py +373 -0
- zettcode/engine/storage/store.py +1056 -0
- zettcode/engine/tools/__init__.py +7 -0
- zettcode/engine/tools/boundary.py +56 -0
- zettcode/engine/tools/files.py +262 -0
- zettcode/engine/tools/images.py +32 -0
- zettcode/engine/tools/output.py +81 -0
- zettcode/engine/tools/shell.py +181 -0
- zettcode/engine/tools/workspace.py +271 -0
- zettcode/error_text.py +217 -0
- zettcode/filelock.py +125 -0
- zettcode/logging_setup.py +196 -0
- zettcode/mention.py +124 -0
- zettcode/modes.py +18 -0
- zettcode/paths.py +113 -0
- zettcode/plugin_tui.py +133 -0
- zettcode/processes.py +126 -0
- zettcode/prompt.py +94 -0
- zettcode/providers.py +215 -0
- zettcode/py.typed +0 -0
- zettcode/registry.py +63 -0
- zettcode/relay.py +267 -0
- zettcode/statusline_parts.py +84 -0
- zettcode/version.py +11 -0
- zettcode/wire.py +30 -0
- zettcode/wording.py +11 -0
- zettcode_engine-0.3.4.dist-info/METADATA +23 -0
- zettcode_engine-0.3.4.dist-info/RECORD +90 -0
- zettcode_engine-0.3.4.dist-info/WHEEL +4 -0
|
@@ -0,0 +1,521 @@
|
|
|
1
|
+
"""What the next request carries, broken down by who contributed it.
|
|
2
|
+
|
|
3
|
+
The runtime assembles one request per model call. ``before_model`` sees that
|
|
4
|
+
request just before it leaves, so an extension can remember it; ``/context``
|
|
5
|
+
then measures the remembered messages and tool schemas and shows how much of
|
|
6
|
+
the budget each source takes: the instructions, the tools, the user's turns,
|
|
7
|
+
the model's turns, and the tool output.
|
|
8
|
+
|
|
9
|
+
Token counts come from ``tiktoken`` when its encoding can be opened. That
|
|
10
|
+
encoding is fetched over the network the first time, which a CLI cannot rely
|
|
11
|
+
on, so the count falls back to a characters-per-token estimate and the report
|
|
12
|
+
says which one it used.
|
|
13
|
+
"""
|
|
14
|
+
|
|
15
|
+
from __future__ import annotations
|
|
16
|
+
|
|
17
|
+
import json
|
|
18
|
+
import os
|
|
19
|
+
from collections.abc import Callable, MutableMapping, Sequence
|
|
20
|
+
from dataclasses import dataclass, field
|
|
21
|
+
from typing import TypeAlias, cast
|
|
22
|
+
|
|
23
|
+
from zett_agent.agent import AgentRunContext
|
|
24
|
+
from zett_agent.extensions.base import AgentExtension
|
|
25
|
+
from zett_agent.extensions.compaction import CompactedMessage
|
|
26
|
+
from zett_agent.messages import (
|
|
27
|
+
AnyMessage,
|
|
28
|
+
AssistantMessage,
|
|
29
|
+
ImageContent,
|
|
30
|
+
SystemMessage,
|
|
31
|
+
TextContent,
|
|
32
|
+
ToolMessage,
|
|
33
|
+
)
|
|
34
|
+
from zett_agent.model import ModelRequest, ToolDefinition
|
|
35
|
+
|
|
36
|
+
from zettcode.budget import CHARS_PER_TOKEN, SOURCES
|
|
37
|
+
from zettcode.engine.pipeline import DaemonCall
|
|
38
|
+
from zettcode.paths import DEFAULT_TOKENIZER_CACHE
|
|
39
|
+
|
|
40
|
+
#: What one image costs, whatever its encoded size. A provider bills an image by
|
|
41
|
+
#: its dimensions and the detail asked for, so counting the bytes of a pasted
|
|
42
|
+
#: screenshot — megabytes of base64 — overstates the context by orders of
|
|
43
|
+
#: magnitude: it measured a request the provider counted at 326 k tokens as 1.74 M
|
|
44
|
+
#: here, and spent 400 ms encoding a picture. The value matches the estimate
|
|
45
|
+
#: zett-agent's compaction uses, so the panel and the compaction threshold do not
|
|
46
|
+
#: disagree about the same image.
|
|
47
|
+
IMAGE_TOKENS = 1_100
|
|
48
|
+
|
|
49
|
+
#: What an image part contributes to the text, since it has no text of its own.
|
|
50
|
+
IMAGE_MARKER = "[image]"
|
|
51
|
+
|
|
52
|
+
#: Seconds to wait for tiktoken's vocabulary before falling back.
|
|
53
|
+
ENCODING_DEADLINE = 2.0
|
|
54
|
+
|
|
55
|
+
|
|
56
|
+
@dataclass(frozen=True, slots=True)
|
|
57
|
+
class Counted:
|
|
58
|
+
"""What one message or tool cost the last time it was counted.
|
|
59
|
+
|
|
60
|
+
Attributes:
|
|
61
|
+
item: The message or tool the count belongs to. Holding it is what makes
|
|
62
|
+
the key sound: a cache keyed by ``id`` can only name a different
|
|
63
|
+
object once the one it named is gone, and this keeps that one alive.
|
|
64
|
+
tokens: What the counter said it was worth.
|
|
65
|
+
"""
|
|
66
|
+
|
|
67
|
+
item: object
|
|
68
|
+
tokens: int
|
|
69
|
+
|
|
70
|
+
|
|
71
|
+
#: Counts from an earlier measurement, keyed by ``id`` of the item they describe.
|
|
72
|
+
Counts: TypeAlias = MutableMapping[int, Counted]
|
|
73
|
+
|
|
74
|
+
|
|
75
|
+
@dataclass(slots=True)
|
|
76
|
+
class SessionCounts:
|
|
77
|
+
"""One session's counts, and the counter that produced them.
|
|
78
|
+
|
|
79
|
+
Attributes:
|
|
80
|
+
measured_by: The counter's label. Numbers from one counter are never
|
|
81
|
+
reused for another: the estimate and the encoder disagree by design.
|
|
82
|
+
counts: One entry per message or tool the last report measured.
|
|
83
|
+
"""
|
|
84
|
+
|
|
85
|
+
measured_by: str
|
|
86
|
+
counts: Counts = field(default_factory=dict)
|
|
87
|
+
|
|
88
|
+
|
|
89
|
+
@dataclass(frozen=True, slots=True)
|
|
90
|
+
class ContextShare:
|
|
91
|
+
"""One source's share of the budget.
|
|
92
|
+
|
|
93
|
+
Attributes:
|
|
94
|
+
name: Row label, one of :data:`SOURCES`.
|
|
95
|
+
tokens: Tokens this source contributes.
|
|
96
|
+
items: Messages, tools, or summaries counted into ``tokens``.
|
|
97
|
+
"""
|
|
98
|
+
|
|
99
|
+
name: str
|
|
100
|
+
tokens: int
|
|
101
|
+
items: int = 0
|
|
102
|
+
|
|
103
|
+
|
|
104
|
+
@dataclass(frozen=True, slots=True)
|
|
105
|
+
class ContextReport:
|
|
106
|
+
"""The budget as a whole: every source, the total, and what is left.
|
|
107
|
+
|
|
108
|
+
A request has two numbers, and they come from different places: how many
|
|
109
|
+
tokens it is, which only the provider knows because it counts what it bills,
|
|
110
|
+
and how that total divides between the sources, which only the tokenizer can
|
|
111
|
+
say because no provider reports a breakdown. The shares are always the
|
|
112
|
+
tokenizer's; the total is the provider's whenever it reported one for this
|
|
113
|
+
very request, and the tokenizer's sum when it did not — a session restored
|
|
114
|
+
from the store, where the stored figure belongs to a call that may predate a
|
|
115
|
+
compaction.
|
|
116
|
+
|
|
117
|
+
Attributes:
|
|
118
|
+
window: Tokens the budget is measured against.
|
|
119
|
+
measured_by: Which counter produced the shares.
|
|
120
|
+
shares: One entry per non-empty source, in :data:`SOURCES` order.
|
|
121
|
+
partial: Whether a part the runtime would send is missing from the
|
|
122
|
+
measurement, which is the case for a restored session measured
|
|
123
|
+
before any reply in this process: its tool schemas are not knowable
|
|
124
|
+
yet, and the page says so instead of reading low silently.
|
|
125
|
+
reported: Tokens the provider counted for this request, if it did.
|
|
126
|
+
"""
|
|
127
|
+
|
|
128
|
+
window: int
|
|
129
|
+
measured_by: str
|
|
130
|
+
shares: tuple[ContextShare, ...]
|
|
131
|
+
partial: bool = False
|
|
132
|
+
reported: int | None = None
|
|
133
|
+
|
|
134
|
+
@property
|
|
135
|
+
def measured(self) -> int:
|
|
136
|
+
"""Return the tokens the tokenizer counted, which the shares add up to."""
|
|
137
|
+
return sum(share.tokens for share in self.shares)
|
|
138
|
+
|
|
139
|
+
@property
|
|
140
|
+
def used(self) -> int:
|
|
141
|
+
"""Return the tokens this request carries: the provider's count if there is one."""
|
|
142
|
+
return self.reported if self.reported is not None else self.measured
|
|
143
|
+
|
|
144
|
+
@property
|
|
145
|
+
def from_provider(self) -> bool:
|
|
146
|
+
"""Return whether the total is the provider's rather than the tokenizer's."""
|
|
147
|
+
return self.reported is not None
|
|
148
|
+
|
|
149
|
+
@property
|
|
150
|
+
def estimated(self) -> bool:
|
|
151
|
+
"""Return whether the shares are a characters-per-token estimate and not a tokenizer's count."""
|
|
152
|
+
return "chars/token" in self.measured_by
|
|
153
|
+
|
|
154
|
+
@property
|
|
155
|
+
def free(self) -> int:
|
|
156
|
+
"""Return the tokens left before the budget is full, never negative."""
|
|
157
|
+
return max(0, self.window - self.used)
|
|
158
|
+
|
|
159
|
+
def percent(self, tokens: int) -> float:
|
|
160
|
+
"""Return ``tokens`` as a share of the budget, in percent."""
|
|
161
|
+
return tokens / self.window * 100 if self.window else 0.0
|
|
162
|
+
|
|
163
|
+
def share(self, tokens: int) -> float:
|
|
164
|
+
"""Return ``tokens`` as a share of what the request is using, in percent.
|
|
165
|
+
|
|
166
|
+
:meth:`percent` answers "how full is the model's window"; this answers
|
|
167
|
+
"where does the context go", which is the question a breakdown is for —
|
|
168
|
+
so it divides by what the tokenizer counted, never by a provider's total
|
|
169
|
+
for the same request, which would leave the rows short of a hundred.
|
|
170
|
+
"""
|
|
171
|
+
measured = self.measured
|
|
172
|
+
return tokens / measured * 100 if measured else 0.0
|
|
173
|
+
|
|
174
|
+
|
|
175
|
+
class Tokenizer:
|
|
176
|
+
"""Count tokens with tiktoken when it can be opened, else by character count."""
|
|
177
|
+
|
|
178
|
+
def __init__(
|
|
179
|
+
self,
|
|
180
|
+
encoding: str = "cl100k_base",
|
|
181
|
+
*,
|
|
182
|
+
deadline: float = ENCODING_DEADLINE,
|
|
183
|
+
loader: Callable[[str], Callable[[str], Sequence[int]]] | None = None,
|
|
184
|
+
) -> None:
|
|
185
|
+
"""Configure the encoding, the deadline, and how it is loaded.
|
|
186
|
+
|
|
187
|
+
Args:
|
|
188
|
+
encoding: tiktoken encoding name; ``cl100k_base`` covers the models
|
|
189
|
+
this shell talks to and their gateways.
|
|
190
|
+
deadline: Seconds to wait for the encoder before estimating.
|
|
191
|
+
loader: Returns the raw encoder, injectable for tests; the default
|
|
192
|
+
loads tiktoken, which may fetch its vocabulary over the network.
|
|
193
|
+
"""
|
|
194
|
+
self.encoding = encoding
|
|
195
|
+
self.deadline = deadline
|
|
196
|
+
self._loader = loader or _tiktoken_encoder
|
|
197
|
+
self._encode: Callable[[str], Sequence[int]] | None = None
|
|
198
|
+
self._resolved = False
|
|
199
|
+
self._call: DaemonCall[Callable[[str], Sequence[int]]] | None = None
|
|
200
|
+
|
|
201
|
+
@property
|
|
202
|
+
def label(self) -> str:
|
|
203
|
+
"""Return the counter's name, for the report to show its own accuracy."""
|
|
204
|
+
return f"tiktoken {self.encoding}" if self._encode is not None else f"~{CHARS_PER_TOKEN} chars/token"
|
|
205
|
+
|
|
206
|
+
@property
|
|
207
|
+
def prepared(self) -> bool:
|
|
208
|
+
"""Return whether an encoder was loaded."""
|
|
209
|
+
return self._encode is not None
|
|
210
|
+
|
|
211
|
+
def count(self, text: str) -> int:
|
|
212
|
+
"""Return the tokens in ``text``, estimating when no encoder loaded."""
|
|
213
|
+
if not text:
|
|
214
|
+
return 0
|
|
215
|
+
if self._encode is not None:
|
|
216
|
+
return len(self._encode(text))
|
|
217
|
+
return max(1, round(len(text) / CHARS_PER_TOKEN))
|
|
218
|
+
|
|
219
|
+
async def prepare(self) -> None:
|
|
220
|
+
"""Load the encoder at most once, giving up after the deadline.
|
|
221
|
+
|
|
222
|
+
The load runs on a daemon thread because opening an encoding may fetch
|
|
223
|
+
its vocabulary — seconds on a first run, and an executor's thread would
|
|
224
|
+
hold the process open for the rest of it if the shell left first. A
|
|
225
|
+
failure or a timeout leaves the character estimate in place rather than
|
|
226
|
+
delaying the page, and the load is left to finish on its own: a later
|
|
227
|
+
report adopts its answer instead of estimating for the whole session.
|
|
228
|
+
"""
|
|
229
|
+
if self._resolved:
|
|
230
|
+
if self._encode is None and self._call is not None:
|
|
231
|
+
self._encode = self._call.peek()
|
|
232
|
+
return
|
|
233
|
+
self._resolved = True
|
|
234
|
+
self._call = DaemonCall(lambda: self._loader(self.encoding))
|
|
235
|
+
self._call.start()
|
|
236
|
+
self._encode = await self._call.wait(self.deadline)
|
|
237
|
+
|
|
238
|
+
|
|
239
|
+
def _tiktoken_encoder(encoding: str) -> Callable[[str], Sequence[int]]:
|
|
240
|
+
"""Return tiktoken's encoder, importing it lazily so startup stays cheap.
|
|
241
|
+
|
|
242
|
+
tiktoken caches the vocabulary it downloads in the OS temp directory, which
|
|
243
|
+
is cleaned whenever the system feels like it: a purge turns the next load
|
|
244
|
+
into a 1.7 MB fetch — five seconds, measured — behind the deadline above, so
|
|
245
|
+
the count would arrive as an estimate after a wait. This program's own cache
|
|
246
|
+
directory survives that, so the fetch happens once. An explicit
|
|
247
|
+
``TIKTOKEN_CACHE_DIR`` wins, because that is the reader's own choice.
|
|
248
|
+
"""
|
|
249
|
+
os.environ.setdefault("TIKTOKEN_CACHE_DIR", str(DEFAULT_TOKENIZER_CACHE))
|
|
250
|
+
import tiktoken
|
|
251
|
+
|
|
252
|
+
return tiktoken.get_encoding(encoding).encode
|
|
253
|
+
|
|
254
|
+
|
|
255
|
+
def message_text(message: AnyMessage) -> str:
|
|
256
|
+
"""Return the text of one message as the model sees it, for counting.
|
|
257
|
+
|
|
258
|
+
Args:
|
|
259
|
+
message: Any context message. Tool calls are counted through their name
|
|
260
|
+
and arguments, a tool result through the text it rendered, and an
|
|
261
|
+
image through :data:`IMAGE_MARKER` rather than its payload.
|
|
262
|
+
"""
|
|
263
|
+
if isinstance(message, ToolMessage):
|
|
264
|
+
return message.text
|
|
265
|
+
if isinstance(message, AssistantMessage):
|
|
266
|
+
parts = [message.reasoning or "", message.content]
|
|
267
|
+
parts.extend(
|
|
268
|
+
f"{call.name} {json.dumps(call.arguments, ensure_ascii=False, default=str)}" for call in message.tool_calls
|
|
269
|
+
)
|
|
270
|
+
return "\n".join(part for part in parts if part)
|
|
271
|
+
return content_text(getattr(message, "content", ""))
|
|
272
|
+
|
|
273
|
+
|
|
274
|
+
def content_text(content: object) -> str:
|
|
275
|
+
"""Return the text of a message's content, images standing in for themselves.
|
|
276
|
+
|
|
277
|
+
Content is a string or a sequence of parts. Turning the whole thing into a
|
|
278
|
+
string with ``str`` renders an image's encoded bytes — a pasted screenshot is
|
|
279
|
+
megabytes of base64 — which is what made a 326 k request measure 1.74 M.
|
|
280
|
+
|
|
281
|
+
Args:
|
|
282
|
+
content: The message's content, as the provider would take it.
|
|
283
|
+
"""
|
|
284
|
+
if isinstance(content, str):
|
|
285
|
+
return content
|
|
286
|
+
if isinstance(content, Sequence):
|
|
287
|
+
return "\n".join(
|
|
288
|
+
part.text
|
|
289
|
+
if isinstance(part, TextContent)
|
|
290
|
+
else (part.alt_text or IMAGE_MARKER)
|
|
291
|
+
if isinstance(part, ImageContent)
|
|
292
|
+
else repr(part)
|
|
293
|
+
for part in content
|
|
294
|
+
)
|
|
295
|
+
return str(content)
|
|
296
|
+
|
|
297
|
+
|
|
298
|
+
def image_count(message: AnyMessage) -> int:
|
|
299
|
+
"""Return how many images one message carries, which are billed apart from its text."""
|
|
300
|
+
content = getattr(message, "content", "")
|
|
301
|
+
if isinstance(content, str) or not isinstance(content, Sequence):
|
|
302
|
+
return 0
|
|
303
|
+
return sum(1 for part in content if isinstance(part, ImageContent))
|
|
304
|
+
|
|
305
|
+
|
|
306
|
+
def tool_text(tool: ToolDefinition) -> str:
|
|
307
|
+
"""Return the schema text of one tool, as the provider receives it."""
|
|
308
|
+
parameters = json.dumps(tool.parameters, ensure_ascii=False, default=str, sort_keys=True)
|
|
309
|
+
return f"{tool.name}\n{tool.description}\n{parameters}"
|
|
310
|
+
|
|
311
|
+
|
|
312
|
+
def count_message(message: AnyMessage, tokenizer: Tokenizer, counts: Counts | None) -> int:
|
|
313
|
+
"""Return the tokens one message costs, reusing an earlier count of it.
|
|
314
|
+
|
|
315
|
+
Its text is measured; each image is not, because a provider bills an image by
|
|
316
|
+
its dimensions and detail rather than by the bytes a paste carries.
|
|
317
|
+
|
|
318
|
+
Args:
|
|
319
|
+
message: The message to count.
|
|
320
|
+
tokenizer: Counter for the text.
|
|
321
|
+
counts: Earlier counts, updated in place; ``None`` counts everything.
|
|
322
|
+
"""
|
|
323
|
+
return count_one(
|
|
324
|
+
message,
|
|
325
|
+
message_text(message),
|
|
326
|
+
tokenizer,
|
|
327
|
+
counts,
|
|
328
|
+
extra=IMAGE_TOKENS * image_count(message),
|
|
329
|
+
)
|
|
330
|
+
|
|
331
|
+
|
|
332
|
+
def count_one(item: object, text: str, tokenizer: Tokenizer, counts: Counts | None, *, extra: int = 0) -> int:
|
|
333
|
+
"""Return the tokens in one counted item, reusing an earlier count of it.
|
|
334
|
+
|
|
335
|
+
Counting is the whole cost of a report — 25 ms for a session of 155 k tokens
|
|
336
|
+
— and the overwhelming majority of what a report walks was counted the last
|
|
337
|
+
time the panel was opened: a turn appends, it does not rewrite. The cache
|
|
338
|
+
turns that into a lookup, and only what arrived since is encoded.
|
|
339
|
+
|
|
340
|
+
Args:
|
|
341
|
+
item: The message or tool the text belongs to. Its identity is the key,
|
|
342
|
+
and the entry holds it, so a key can never name a different object.
|
|
343
|
+
text: What the tokenizer would count.
|
|
344
|
+
tokenizer: Counter to use when there is no entry.
|
|
345
|
+
counts: Earlier counts, updated in place; ``None`` counts everything.
|
|
346
|
+
extra: Tokens the text cannot express, an image's estimate say; they are
|
|
347
|
+
part of the cached number, so a hit returns the same answer.
|
|
348
|
+
"""
|
|
349
|
+
if counts is None:
|
|
350
|
+
return tokenizer.count(text) + extra
|
|
351
|
+
key = id(item)
|
|
352
|
+
entry = counts.get(key)
|
|
353
|
+
if entry is not None and entry.item is item:
|
|
354
|
+
return entry.tokens
|
|
355
|
+
tokens = tokenizer.count(text) + extra
|
|
356
|
+
counts[key] = Counted(item, tokens)
|
|
357
|
+
return tokens
|
|
358
|
+
|
|
359
|
+
|
|
360
|
+
def measure(
|
|
361
|
+
messages: Sequence[AnyMessage],
|
|
362
|
+
tools: Sequence[ToolDefinition],
|
|
363
|
+
*,
|
|
364
|
+
window: int,
|
|
365
|
+
tokenizer: Tokenizer,
|
|
366
|
+
counts: Counts | None = None,
|
|
367
|
+
reported: int | None = None,
|
|
368
|
+
) -> ContextReport:
|
|
369
|
+
"""Break one assembled request into the sources that fill the budget.
|
|
370
|
+
|
|
371
|
+
Args:
|
|
372
|
+
messages: Context in provider order; the first system message is this
|
|
373
|
+
application's prompt and later ones are the environment and tool
|
|
374
|
+
notes an extension added.
|
|
375
|
+
tools: Tool schemas offered with the request.
|
|
376
|
+
window: Token budget the shares are measured against.
|
|
377
|
+
tokenizer: Counter to use; it may still be estimating.
|
|
378
|
+
counts: Counts from an earlier measurement of the same items, updated in
|
|
379
|
+
place as new ones are counted; ``None`` measures everything.
|
|
380
|
+
reported: Tokens the provider counted for this request, if it did; the
|
|
381
|
+
shares are still measured here, because it cannot break them down.
|
|
382
|
+
"""
|
|
383
|
+
totals = dict.fromkeys(SOURCES, 0)
|
|
384
|
+
items = dict.fromkeys(SOURCES, 0)
|
|
385
|
+
system_seen = False
|
|
386
|
+
for message in messages:
|
|
387
|
+
match message:
|
|
388
|
+
case SystemMessage():
|
|
389
|
+
source = "Environment notes" if system_seen else "System prompt"
|
|
390
|
+
system_seen = True
|
|
391
|
+
case CompactedMessage():
|
|
392
|
+
source = "Compaction summary"
|
|
393
|
+
case AssistantMessage():
|
|
394
|
+
source = "Assistant messages"
|
|
395
|
+
case ToolMessage():
|
|
396
|
+
source = "Tool output"
|
|
397
|
+
case _:
|
|
398
|
+
source = "User messages"
|
|
399
|
+
totals[source] += count_message(message, tokenizer, counts)
|
|
400
|
+
items[source] += 1
|
|
401
|
+
if tools:
|
|
402
|
+
totals["Tool schemas"] = sum(count_one(tool, tool_text(tool), tokenizer, counts) for tool in tools)
|
|
403
|
+
items["Tool schemas"] = len(tools)
|
|
404
|
+
shares = tuple(ContextShare(name, totals[name], items[name]) for name in SOURCES if totals[name] or items[name])
|
|
405
|
+
return ContextReport(
|
|
406
|
+
window=max(1, window),
|
|
407
|
+
measured_by=tokenizer.label,
|
|
408
|
+
shares=shares,
|
|
409
|
+
partial=not tools,
|
|
410
|
+
reported=reported,
|
|
411
|
+
)
|
|
412
|
+
|
|
413
|
+
|
|
414
|
+
class ContextExtension(AgentExtension):
|
|
415
|
+
"""Remember the request the runtime assembled for the last model call.
|
|
416
|
+
|
|
417
|
+
Nothing is measured here: tokenizing every request would cost a little on
|
|
418
|
+
every turn, and the report is only wanted when ``/context`` asks for it. The
|
|
419
|
+
extension keeps the assembled messages and tools per session instead.
|
|
420
|
+
|
|
421
|
+
A session restored from the store has not been assembled in this process, so
|
|
422
|
+
the shell primes it with the stored dialogue; the instructions and tool
|
|
423
|
+
schemas come from the last request of any session, because one configuration
|
|
424
|
+
sends the same ones every time.
|
|
425
|
+
"""
|
|
426
|
+
|
|
427
|
+
def __init__(self) -> None:
|
|
428
|
+
"""Start with no captured request and nothing counted."""
|
|
429
|
+
self._requests: dict[str, tuple[tuple[AnyMessage, ...], tuple[ToolDefinition, ...]]] = {}
|
|
430
|
+
#: Requests promised but not built yet: a restored session's dialogue is only
|
|
431
|
+
#: assembled when something asks for it (see :meth:`remember`).
|
|
432
|
+
self._pending: dict[str, Callable[[], Sequence[AnyMessage]]] = {}
|
|
433
|
+
self._template: tuple[tuple[AnyMessage, ...], tuple[ToolDefinition, ...]] | None = None
|
|
434
|
+
#: Per session, the items the last report counted and the counter that
|
|
435
|
+
#: produced those numbers (see :class:`SessionCounts`). This is what makes
|
|
436
|
+
#: opening `/context` twice cost nothing the second time; it is dropped
|
|
437
|
+
#: when the counter changes, because the estimate and the encoder disagree.
|
|
438
|
+
self._counts: dict[str, SessionCounts] = {}
|
|
439
|
+
|
|
440
|
+
async def before_model(self, context: AgentRunContext, request: ModelRequest) -> None:
|
|
441
|
+
"""Keep the request that is about to leave, replacing the previous one."""
|
|
442
|
+
# The runtime replaces an empty session_id before hooks run.
|
|
443
|
+
session_id = cast(str, context.config.session_id)
|
|
444
|
+
self._requests[session_id] = (tuple(request.messages), tuple(request.tools))
|
|
445
|
+
self._pending.pop(session_id, None)
|
|
446
|
+
self._template = (tuple(request.messages), tuple(request.tools))
|
|
447
|
+
|
|
448
|
+
@property
|
|
449
|
+
def instructions(self) -> tuple[AnyMessage, ...]:
|
|
450
|
+
"""Return the instruction messages every request starts with, once seen."""
|
|
451
|
+
return self._template[0] if self._template is not None else ()
|
|
452
|
+
|
|
453
|
+
@property
|
|
454
|
+
def tools(self) -> tuple[ToolDefinition, ...]:
|
|
455
|
+
"""Return the tool schemas every request offers, once one was seen."""
|
|
456
|
+
return self._template[1] if self._template is not None else ()
|
|
457
|
+
|
|
458
|
+
def remember(self, session_id: str, messages: Sequence[AnyMessage] | Callable[[], Sequence[AnyMessage]]) -> None:
|
|
459
|
+
"""Record a context assembled outside a model call, as restoring does.
|
|
460
|
+
|
|
461
|
+
Args:
|
|
462
|
+
session_id: Session the messages belong to.
|
|
463
|
+
messages: Instructions followed by the restored dialogue, in the
|
|
464
|
+
order a request would carry them — or a function that builds that
|
|
465
|
+
list. A restore passes the function: the list needs every message
|
|
466
|
+
since the last checkpoint decoded, which is most of what resuming a
|
|
467
|
+
long session would cost, and only ``/context`` and the context
|
|
468
|
+
figure ever read it. It runs the first time one of them asks, and
|
|
469
|
+
a real request in the meantime replaces it without it ever running.
|
|
470
|
+
"""
|
|
471
|
+
self._requests.pop(session_id, None)
|
|
472
|
+
if callable(messages):
|
|
473
|
+
self._pending[session_id] = messages
|
|
474
|
+
return
|
|
475
|
+
self._pending.pop(session_id, None)
|
|
476
|
+
self._requests[session_id] = (tuple(messages), self.tools)
|
|
477
|
+
|
|
478
|
+
def assembled(self, session_id: str) -> tuple[tuple[AnyMessage, ...], tuple[ToolDefinition, ...]] | None:
|
|
479
|
+
"""Return the last assembled request, or None before the first call."""
|
|
480
|
+
build = self._pending.pop(session_id, None)
|
|
481
|
+
if build is not None:
|
|
482
|
+
self._requests[session_id] = (tuple(build()), self.tools)
|
|
483
|
+
return self._requests.get(session_id)
|
|
484
|
+
|
|
485
|
+
def report(
|
|
486
|
+
self,
|
|
487
|
+
session_id: str,
|
|
488
|
+
*,
|
|
489
|
+
window: int,
|
|
490
|
+
tokenizer: Tokenizer,
|
|
491
|
+
reported: int | None = None,
|
|
492
|
+
) -> ContextReport | None:
|
|
493
|
+
"""Return the breakdown of the last assembled request, if there was one.
|
|
494
|
+
|
|
495
|
+
What was counted for an earlier report is counted again only if the
|
|
496
|
+
counter changed: a turn appends to the context rather than rewriting it,
|
|
497
|
+
so a report spends its time on what arrived since the last one.
|
|
498
|
+
|
|
499
|
+
Args:
|
|
500
|
+
session_id: Session to report on.
|
|
501
|
+
window: Tokens the model can carry.
|
|
502
|
+
tokenizer: Counter that measures the shares.
|
|
503
|
+
reported: Tokens the provider counted for this request, if it counted
|
|
504
|
+
it; it becomes the report's total while the shares stay measured.
|
|
505
|
+
"""
|
|
506
|
+
assembled = self.assembled(session_id)
|
|
507
|
+
if assembled is None:
|
|
508
|
+
return None
|
|
509
|
+
messages, tools = assembled
|
|
510
|
+
label = tokenizer.label
|
|
511
|
+
cached = self._counts.get(session_id)
|
|
512
|
+
if cached is None or cached.measured_by != label:
|
|
513
|
+
cached = SessionCounts(measured_by=label)
|
|
514
|
+
counts = cached.counts
|
|
515
|
+
live = {id(message) for message in messages} | {id(tool) for tool in tools}
|
|
516
|
+
for key in [key for key in counts if key not in live]:
|
|
517
|
+
# The request stopped carrying this item, so its count has no reader
|
|
518
|
+
# and its key must not outlive the object it stands for.
|
|
519
|
+
del counts[key]
|
|
520
|
+
self._counts[session_id] = cached
|
|
521
|
+
return measure(messages, tools, window=window, tokenizer=tokenizer, counts=counts, reported=reported)
|
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
"""The extensions that give an agent its file tools and its shell.
|
|
2
|
+
|
|
3
|
+
The tools themselves are in :mod:`zettcode.engine.tools`; these register them for each request and tell the model
|
|
4
|
+
where it is working:
|
|
5
|
+
|
|
6
|
+
* :class:`FileTools` registers ``read``, ``glob`` and ``grep``, and with ``read_only`` off ``write_file`` and
|
|
7
|
+
``edit``. It is what a read-only subagent, or one that may edit, is given. A subagent has no shell, so searching
|
|
8
|
+
is a tool of its own there.
|
|
9
|
+
* :class:`CodingBundle` is the main agent's: ``read``, ``write_file`` and ``edit``, followed by ``bash``. It has no
|
|
10
|
+
``glob`` or ``grep``, because it has a shell: ``rg``, ``grep`` and ``find`` do that work, and run through the same
|
|
11
|
+
permission level as any other command.
|
|
12
|
+
|
|
13
|
+
Neither offers ``read_file``: that is the text-only reader
|
|
14
|
+
:class:`~zettcode.engine.extensions.capabilities.ModelCapabilities` puts in the place of ``read`` for a model that
|
|
15
|
+
does not take images.
|
|
16
|
+
"""
|
|
17
|
+
|
|
18
|
+
from __future__ import annotations
|
|
19
|
+
|
|
20
|
+
import json
|
|
21
|
+
from pathlib import Path
|
|
22
|
+
|
|
23
|
+
from zett_agent.agent import AgentRunContext
|
|
24
|
+
from zett_agent.extensions.base import AgentExtension
|
|
25
|
+
from zett_agent.messages import SystemMessage
|
|
26
|
+
from zett_agent.tools.base import AgentTool
|
|
27
|
+
|
|
28
|
+
from ..tools.boundary import Boundary
|
|
29
|
+
from ..tools.files import edit, glob, grep, read, write_file
|
|
30
|
+
from ..tools.shell import shell_tool
|
|
31
|
+
|
|
32
|
+
__all__ = ["CodingBundle", "FileTools", "file_tools"]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def file_tools(*, read_only: bool, search: bool = True) -> tuple[AgentTool, ...]:
|
|
36
|
+
"""Return the file tools for one mode.
|
|
37
|
+
|
|
38
|
+
Args:
|
|
39
|
+
read_only: Offer only the tools that read; otherwise ``write_file`` and ``edit`` too.
|
|
40
|
+
search: Offer ``glob`` and ``grep``. An agent with a shell leaves them out and searches with it.
|
|
41
|
+
"""
|
|
42
|
+
reading: tuple[AgentTool, ...] = (read, glob, grep) if search else (read,)
|
|
43
|
+
return reading if read_only else (*reading, write_file, edit)
|
|
44
|
+
|
|
45
|
+
|
|
46
|
+
class FileTools(AgentExtension):
|
|
47
|
+
"""Register the file tools, and say which directory their relative paths start from.
|
|
48
|
+
|
|
49
|
+
The directory is looked up at every request, since the shell can change it between two of them.
|
|
50
|
+
"""
|
|
51
|
+
|
|
52
|
+
def __init__(self, *, read_only: bool = False, search: bool = True) -> None:
|
|
53
|
+
"""Choose the mode: ``read_only`` keeps the tools that change a file out of the agent's hands."""
|
|
54
|
+
self.read_only = read_only
|
|
55
|
+
self.search = search
|
|
56
|
+
|
|
57
|
+
@property
|
|
58
|
+
def tools(self) -> tuple[AgentTool, ...]:
|
|
59
|
+
"""Return the tools of the configured mode."""
|
|
60
|
+
return file_tools(read_only=self.read_only, search=self.search)
|
|
61
|
+
|
|
62
|
+
async def on_tool(self, context: AgentRunContext) -> None:
|
|
63
|
+
"""Register this extension's tools for the request."""
|
|
64
|
+
for registered in self.tools:
|
|
65
|
+
context.register_tool(registered)
|
|
66
|
+
|
|
67
|
+
def _environment(self, directory: Path) -> str:
|
|
68
|
+
"""Describe the directory every relative path is taken from."""
|
|
69
|
+
return (
|
|
70
|
+
"# Filesystem environment\n"
|
|
71
|
+
f"Current working directory: {json.dumps(str(directory), ensure_ascii=False)}.\n"
|
|
72
|
+
"Relative filesystem paths are resolved from this directory. Absolute paths are also allowed."
|
|
73
|
+
)
|
|
74
|
+
|
|
75
|
+
async def on_state(self, context: AgentRunContext) -> None:
|
|
76
|
+
"""Put the request's working directory among the system messages, after the ones already there."""
|
|
77
|
+
instructions = [item for item in context.state.messages if isinstance(item, SystemMessage)]
|
|
78
|
+
message = SystemMessage(content=self._environment(Path.cwd().resolve()))
|
|
79
|
+
context.add_message(message, index=len(instructions))
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
class CodingBundle(FileTools):
|
|
83
|
+
"""The main agent's tools: read, write and edit, then the shell, in the working directory the shell starts in.
|
|
84
|
+
|
|
85
|
+
There is no ``glob`` or ``grep``: the shell searches. With a ``boundary`` the shell tool runs each command inside
|
|
86
|
+
the sandbox whenever the boundary says so; without one every command runs as it was given.
|
|
87
|
+
"""
|
|
88
|
+
|
|
89
|
+
def __init__(self, boundary: Boundary | None = None) -> None:
|
|
90
|
+
"""Build the shell tool once, over the boundary that decides how it runs."""
|
|
91
|
+
super().__init__(read_only=False, search=False)
|
|
92
|
+
self._shell = shell_tool(boundary)
|
|
93
|
+
|
|
94
|
+
@property
|
|
95
|
+
def tools(self) -> tuple[AgentTool, ...]:
|
|
96
|
+
"""Return ``read``, ``write_file`` and ``edit``, followed by ``bash``."""
|
|
97
|
+
return (*file_tools(read_only=False, search=False), self._shell)
|
|
98
|
+
|
|
99
|
+
def _environment(self, directory: Path) -> str:
|
|
100
|
+
"""Add that commands start in the same directory."""
|
|
101
|
+
return f"{super()._environment(directory)}\nShell commands run from this directory."
|
|
@@ -0,0 +1,17 @@
|
|
|
1
|
+
"""Where a tool's guidance goes: with the tool, in its own description, and with what it returns.
|
|
2
|
+
|
|
3
|
+
Every agent this application builds (the main one and each sub-agent) carries its tools' snippets and guidelines the
|
|
4
|
+
same way, so the choice is made once here. The guidance is one line of prose beside the tool's description in the
|
|
5
|
+
request's tool definitions, not a pair of sections in the system prompt, which keeps it next to the tool wherever the
|
|
6
|
+
definition goes (a deferred tool found through tool search, an endpoint that shows the model only the definitions).
|
|
7
|
+
The tools' return types are described there too: a tool that returns something with a shape says what shape.
|
|
8
|
+
"""
|
|
9
|
+
|
|
10
|
+
from __future__ import annotations
|
|
11
|
+
|
|
12
|
+
from zett_agent.extensions.tool_guidelines import ToolGuidancePlacement, ToolGuidelinesExtension
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
def tool_guidance() -> ToolGuidelinesExtension:
|
|
16
|
+
"""Return the extension that folds each tool's snippet, guidelines and return type into its description."""
|
|
17
|
+
return ToolGuidelinesExtension(placement=ToolGuidancePlacement.DESCRIPTION, returns=True)
|