coding-agent-cli 0.1.0__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.
@@ -0,0 +1,47 @@
1
+ """Run coding-agent CLIs, stream typed events, and return a classified result."""
2
+
3
+ from .api import EventStream, check, run, run_sync, stream
4
+ from .events import (
5
+ SCHEMA_VERSION,
6
+ Event,
7
+ MessageCompleted,
8
+ MessageStarted,
9
+ RunCompleted,
10
+ SessionIdentified,
11
+ TextDelta,
12
+ ToolCompleted,
13
+ ToolStarted,
14
+ UsageUpdated,
15
+ )
16
+ from .registry import Driver, DriverRegistry, EventHandler, SpawnHandler
17
+ from .types import Capabilities, Diagnostic, Outcome, Request, Result, Usage
18
+
19
+ __version__ = "0.1.0"
20
+
21
+ __all__ = [
22
+ "SCHEMA_VERSION",
23
+ "Capabilities",
24
+ "Diagnostic",
25
+ "Driver",
26
+ "EventStream",
27
+ "DriverRegistry",
28
+ "Event",
29
+ "EventHandler",
30
+ "MessageCompleted",
31
+ "MessageStarted",
32
+ "Outcome",
33
+ "Request",
34
+ "Result",
35
+ "RunCompleted",
36
+ "SessionIdentified",
37
+ "SpawnHandler",
38
+ "TextDelta",
39
+ "ToolCompleted",
40
+ "ToolStarted",
41
+ "Usage",
42
+ "UsageUpdated",
43
+ "check",
44
+ "run",
45
+ "run_sync",
46
+ "stream",
47
+ ]
@@ -0,0 +1,66 @@
1
+ """Shared validation of raw ``args`` and portable options."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from collections.abc import Iterable, Mapping, Sequence
6
+ from typing import Any
7
+
8
+ from .types import Diagnostic, Request
9
+
10
+ PORTABLE_FIELDS = ("model", "reasoning", "system_prompt")
11
+
12
+
13
+ def _has_flag(args: Sequence[str], flags: Iterable[str]) -> str | None:
14
+ """Return the first token in ``args`` that is one of ``flags`` or ``flag=value``."""
15
+ flag_set = set(flags)
16
+ for token in args:
17
+ if token in flag_set or token.split("=", 1)[0] in flag_set:
18
+ return token
19
+ return None
20
+
21
+
22
+ def validate_args(args: Sequence[Any]) -> Diagnostic | None:
23
+ """Reject non-string args, NUL bytes, and the reserved bare ``--`` token."""
24
+ if not all(isinstance(value, str) for value in args):
25
+ return Diagnostic(code="invalid_args", message="args must contain only strings")
26
+ if any("\x00" in value for value in args):
27
+ return Diagnostic(code="invalid_args", message="args cannot contain NUL bytes")
28
+ if "--" in args:
29
+ return Diagnostic(code="invalid_args", message="args cannot contain the reserved '--' token")
30
+ return None
31
+
32
+
33
+ def reject_transport_flags(args: Sequence[str], owned: Iterable[str]) -> Diagnostic | None:
34
+ """Reject a raw arg that overrides a flag the driver owns for transport."""
35
+ token = _has_flag(args, owned)
36
+ if token is None:
37
+ return None
38
+ return Diagnostic(
39
+ code="invalid_args", message=f"args cannot override the driver-owned flag {token.split('=', 1)[0]!r}"
40
+ )
41
+
42
+
43
+ def reject_unsupported(req: Request, supported: Iterable[str]) -> Diagnostic | None:
44
+ """Reject a portable option this driver cannot express."""
45
+ allowed = set(supported)
46
+ for name in PORTABLE_FIELDS:
47
+ if getattr(req, name) is not None and name not in allowed:
48
+ return Diagnostic(code="unsupported_option", message=f"driver {req.driver!r} does not support {name}")
49
+ return None
50
+
51
+
52
+ def reject_conflicts(req: Request, flags: Mapping[str, Iterable[str]]) -> Diagnostic | None:
53
+ """Reject a raw arg that sets a portable field already given. ``flags`` maps field to its flags."""
54
+ for name, field_flags in flags.items():
55
+ if getattr(req, name) is None:
56
+ continue
57
+ token = _has_flag(req.args, field_flags)
58
+ if token is not None:
59
+ return Diagnostic(
60
+ code="invalid_args",
61
+ message=f"args flag {token.split('=', 1)[0]!r} conflicts with the {name} option",
62
+ )
63
+ return None
64
+
65
+
66
+ __all__ = ["PORTABLE_FIELDS", "reject_conflicts", "reject_transport_flags", "reject_unsupported", "validate_args"]
@@ -0,0 +1,170 @@
1
+ """Map facts a driver collected about one run to an ``Outcome`` and ``Diagnostic``.
2
+
3
+ Rows follow the design's outcome precedence; the first matching row wins.
4
+ """
5
+
6
+ from __future__ import annotations
7
+
8
+ import re
9
+ from dataclasses import dataclass
10
+
11
+ from .types import Diagnostic, Outcome
12
+
13
+ REQUEST_CODES = frozenset({"bad_cwd", "invalid_args", "unsupported_option", "invalid_prompt"})
14
+ SPAWN_CODES = frozenset({"missing_cli", "spawn_failed"})
15
+ CLI_CODES = ("auth", "invalid_model", "context_overflow", "service_error")
16
+
17
+ _KEYWORDS: dict[str, tuple[str, ...]] = {
18
+ "auth": (
19
+ "authentication",
20
+ "unauthorized",
21
+ "forbidden",
22
+ "api key",
23
+ "not logged in",
24
+ "not logged",
25
+ "login required",
26
+ "auth_unavailable",
27
+ ),
28
+ "invalid_model": ("invalid model", "unknown model", "not a valid model"),
29
+ "context_overflow": ("context_length_exceeded",),
30
+ "service_error": (
31
+ "overloaded",
32
+ "rate limit",
33
+ "rate_limit",
34
+ "too many requests",
35
+ "internal server error",
36
+ "service unavailable",
37
+ "bad gateway",
38
+ ),
39
+ }
40
+ _CONTEXT_OVERFLOW_RE = re.compile(r"\bcontext_length_exceeded\b")
41
+
42
+
43
+ def _normalize(text: str) -> str:
44
+ return text.lower().replace("node_tls_reject_unauthorized", "")
45
+
46
+
47
+ def keyword_code(text: str) -> str:
48
+ """Return the first CLI code whose keywords appear in ``text``, in precedence order, or ""."""
49
+ lowered = _normalize(text)
50
+ if not lowered.strip():
51
+ return ""
52
+ for code in CLI_CODES:
53
+ if code == "context_overflow":
54
+ if _CONTEXT_OVERFLOW_RE.search(lowered):
55
+ return code
56
+ elif any(token in lowered for token in _KEYWORDS[code]):
57
+ return code
58
+ return ""
59
+
60
+
61
+ def http_status_code(status: int | None) -> str:
62
+ """Map a CLI-reported HTTP status to a CLI code, or ""."""
63
+ if status in (401, 403):
64
+ return "auth"
65
+ if status == 429 or (status is not None and 500 <= status <= 599):
66
+ return "service_error"
67
+ return ""
68
+
69
+
70
+ @dataclass(slots=True, kw_only=True)
71
+ class Facts:
72
+ """What a driver observed. ``cli_error`` is structured CLI error data; ``diagnostics`` is free text."""
73
+
74
+ text: str = ""
75
+ session_id: str | None = None
76
+ cancelled: bool = False
77
+ request_error: Diagnostic | None = None
78
+ process_error: str = ""
79
+ process_message: str = ""
80
+ cli_error: Diagnostic | None = None
81
+ diagnostics: str = ""
82
+ returncode: int | None = None
83
+
84
+
85
+ @dataclass(frozen=True, slots=True)
86
+ class Classification:
87
+ outcome: Outcome
88
+ error: Diagnostic | None
89
+ warnings: tuple[Diagnostic, ...] = ()
90
+
91
+
92
+ _CLI_OUTCOMES = {
93
+ "auth": Outcome.AGENT_FATAL,
94
+ "invalid_model": Outcome.REQUEST_FATAL,
95
+ "context_overflow": Outcome.REQUEST_FATAL,
96
+ "service_error": Outcome.RETRYABLE,
97
+ }
98
+
99
+
100
+ def _cli_diagnostic(facts: Facts) -> Diagnostic | None:
101
+ """Structured CLI error data first, then keyword matching on diagnostics."""
102
+ if facts.cli_error is not None and facts.cli_error.code in _CLI_OUTCOMES:
103
+ return facts.cli_error
104
+ code = keyword_code(facts.diagnostics)
105
+ if code:
106
+ return Diagnostic(code=code, message=facts.diagnostics.strip())
107
+ if facts.cli_error is not None:
108
+ code = keyword_code(facts.cli_error.message)
109
+ if code:
110
+ return Diagnostic(code=code, message=facts.cli_error.message)
111
+ return None
112
+
113
+
114
+ def _message(*parts: str) -> str:
115
+ return "\n".join(part.strip() for part in parts if part and part.strip())
116
+
117
+
118
+ def classify(facts: Facts) -> Classification:
119
+ # Row 1
120
+ if facts.cancelled or facts.process_error == "cancelled":
121
+ return Classification(Outcome.CANCELLED, Diagnostic(code="cancelled", message="run cancelled"))
122
+ # Rows 2 and 3
123
+ if facts.request_error is not None:
124
+ return Classification(Outcome.REQUEST_FATAL, facts.request_error)
125
+ if facts.process_error in REQUEST_CODES:
126
+ return Classification(
127
+ Outcome.REQUEST_FATAL, Diagnostic(code=facts.process_error, message=facts.process_message)
128
+ )
129
+ # Row 4
130
+ if facts.process_error in SPAWN_CODES:
131
+ return Classification(Outcome.AGENT_FATAL, Diagnostic(code=facts.process_error, message=facts.process_message))
132
+ # Rows 5 to 7, and 9 after the timeout row
133
+ cli = _cli_diagnostic(facts)
134
+ if cli is not None and cli.code != "service_error":
135
+ return Classification(_CLI_OUTCOMES[cli.code], cli)
136
+ # Row 8
137
+ if facts.process_error == "timeout":
138
+ return Classification(
139
+ Outcome.RETRYABLE, Diagnostic(code="timeout", message=facts.process_message or "run timed out")
140
+ )
141
+ if cli is not None:
142
+ return Classification(Outcome.RETRYABLE, cli)
143
+ # Row 10
144
+ if facts.process_error in {"execution_error", "stdin_failed"}:
145
+ message = _message(facts.process_message, facts.diagnostics if not facts.text.strip() else "")
146
+ return Classification(Outcome.RETRYABLE, Diagnostic(code=facts.process_error, message=message))
147
+ if facts.cli_error is not None:
148
+ return Classification(Outcome.RETRYABLE, Diagnostic(code="execution_error", message=facts.cli_error.message))
149
+ if facts.returncode:
150
+ return Classification(
151
+ Outcome.RETRYABLE,
152
+ Diagnostic(
153
+ code="execution_error",
154
+ message=_message(f"command exited with status {facts.returncode}", facts.diagnostics),
155
+ ),
156
+ )
157
+ if not facts.text.strip():
158
+ return Classification(
159
+ Outcome.RETRYABLE,
160
+ Diagnostic(code="no_output", message=_message("the CLI produced no assistant text", facts.diagnostics)),
161
+ )
162
+ # Row 11
163
+ if not facts.session_id:
164
+ return Classification(
165
+ Outcome.SUCCESS, None, (Diagnostic(code="no_session_id", message="the CLI reported no session ID"),)
166
+ )
167
+ return Classification(Outcome.SUCCESS, None)
168
+
169
+
170
+ __all__ = ["Classification", "Facts", "classify", "http_status_code", "keyword_code"]
@@ -0,0 +1 @@
1
+ """Built-in drivers. Private: use ``DriverRegistry`` to reach them."""
@@ -0,0 +1,330 @@
1
+ """Helpers shared by the built-in drivers: preflight, capability probes, event buffering, and results."""
2
+
3
+ from __future__ import annotations
4
+
5
+ import asyncio
6
+ import json
7
+ import logging
8
+ import os
9
+ import shutil
10
+ import time
11
+ from collections.abc import Mapping, Sequence
12
+ from dataclasses import dataclass
13
+ from pathlib import Path
14
+ from typing import Any, Literal
15
+
16
+ from .._args import reject_conflicts, reject_transport_flags, reject_unsupported, validate_args
17
+ from .._classify import Facts, classify
18
+ from .._process import ProcessOutcome, run_process
19
+ from ..events import Event, ToolStarted
20
+ from ..registry import EventHandler
21
+ from ..types import Capabilities, Diagnostic, Request, Result, Usage
22
+
23
+ PROBE_TIMEOUT_S = 10.0
24
+
25
+ _VERIFIED_COMMANDS: dict[str, frozenset[str]] = {
26
+ "claude": frozenset({"bash"}),
27
+ "codex": frozenset({"bash"}),
28
+ }
29
+
30
+
31
+ def classify_tool_activity(provider: str, tool_name: str) -> Literal["command", "tool_call"]:
32
+ """Classify tool activity by exact, case-sensitive provider and tool-name allow-lists."""
33
+ commands = _VERIFIED_COMMANDS.get(provider)
34
+ if commands is not None and tool_name in commands:
35
+ return "command"
36
+ return "tool_call"
37
+
38
+
39
+ def tool_started(entity_id: str, provider: str, tool_name: str, tool_input: Any = None) -> ToolStarted:
40
+ return ToolStarted(
41
+ entity_id=entity_id,
42
+ tool=tool_name,
43
+ activity=classify_tool_activity(provider, tool_name),
44
+ input=tool_input,
45
+ )
46
+
47
+
48
+ def tool_status(value: Any) -> Literal["completed", "error", "unknown"]:
49
+ if value in ("completed", "error"):
50
+ return value
51
+ if value in ("failed", "failure"):
52
+ return "error"
53
+ return "unknown"
54
+
55
+
56
+ def parse_json_object(line: str) -> dict[str, Any] | None:
57
+ """Return the line's JSON object, or None for non-JSON or non-object lines."""
58
+ try:
59
+ value = json.loads(line)
60
+ except (json.JSONDecodeError, TypeError):
61
+ return None
62
+ return value if isinstance(value, dict) else None
63
+
64
+
65
+ def as_int(val: Any) -> int | None:
66
+ if isinstance(val, bool) or not isinstance(val, int) or val < 0:
67
+ return None
68
+ return val
69
+
70
+
71
+ def read_and_unlink(path: Path, *, errors: str) -> str:
72
+ """Return the file's UTF-8 text, or "" when it cannot be read. The file is removed either way."""
73
+ try:
74
+ return path.read_text(encoding="utf-8", errors=errors)
75
+ except OSError:
76
+ return ""
77
+ finally:
78
+ try:
79
+ path.unlink()
80
+ except OSError:
81
+ pass
82
+
83
+
84
+ def parse_token_component(usage: dict[str, Any], key: str) -> tuple[int | None, bool]:
85
+ """Return (value, malformed). A missing or null key is absent, not malformed."""
86
+ if key not in usage:
87
+ return None, False
88
+ val = usage[key]
89
+ if val is None:
90
+ return None, False
91
+ if isinstance(val, bool) or not isinstance(val, int) or val < 0:
92
+ return None, True
93
+ return val, False
94
+
95
+
96
+ @dataclass(slots=True, frozen=True)
97
+ class TokenUsage:
98
+ """Normalized usage.
99
+
100
+ ``input`` is raw input plus cache reads and writes, or None when the raw input is absent or malformed.
101
+
102
+ ``observed`` is True when any component was present or malformed.
103
+ """
104
+
105
+ input: int | None
106
+ output: int | None
107
+ cache_read: int | None
108
+ cache_write: int | None
109
+ observed: bool
110
+
111
+
112
+ NO_USAGE = TokenUsage(input=None, output=None, cache_read=None, cache_write=None, observed=False)
113
+
114
+
115
+ def read_usage(
116
+ usage: dict[str, Any],
117
+ *,
118
+ input_key: str = "input_tokens",
119
+ output_key: str = "output_tokens",
120
+ cache_read_key: str = "cache_read_input_tokens",
121
+ cache_write_key: str = "cache_creation_input_tokens",
122
+ ) -> TokenUsage:
123
+ raw_in, raw_in_bad = parse_token_component(usage, input_key)
124
+ cr, cr_bad = parse_token_component(usage, cache_read_key)
125
+ cw, cw_bad = parse_token_component(usage, cache_write_key)
126
+ raw_out, raw_out_bad = parse_token_component(usage, output_key)
127
+ if raw_in_bad or cr_bad or cw_bad or raw_in is None:
128
+ norm_in = None
129
+ else:
130
+ norm_in = raw_in + (cr or 0) + (cw or 0)
131
+ observed = any(v is not None for v in (norm_in, cr, cw, raw_out)) or raw_in_bad or cr_bad or cw_bad or raw_out_bad
132
+ return TokenUsage(input=norm_in, output=raw_out, cache_read=cr, cache_write=cw, observed=observed)
133
+
134
+
135
+ class EventBuffer:
136
+ """Collect the events one parsed line produces, then await the sink once per event in order.
137
+
138
+ The parser stays synchronous. ``flush`` runs after each line, so a failing sink fails the run at that line.
139
+ """
140
+
141
+ def __init__(self, sink: EventHandler) -> None:
142
+ self._sink = sink
143
+ self._pending: list[Event] = []
144
+
145
+ def emit(self, event: Event) -> None:
146
+ self._pending.append(event)
147
+
148
+ async def flush(self) -> None:
149
+ while self._pending:
150
+ await self._sink(self._pending.pop(0))
151
+
152
+
153
+ class Diagnostics:
154
+ """Keep the first ``limit`` diagnostic lines, plus the first line that matches ``keep`` beyond it."""
155
+
156
+ def __init__(self, limit: int = 50) -> None:
157
+ self.limit = limit
158
+ self.lines: list[str] = []
159
+ self._kept: set[str] = set()
160
+
161
+ def add(self, line: str, *, marks: Sequence[str] = ()) -> None:
162
+ extra = [mark for mark in marks if mark not in self._kept]
163
+ if len(self.lines) < self.limit or extra:
164
+ self.lines.append(line)
165
+ self._kept.update(marks)
166
+
167
+ def text(self) -> str:
168
+ return "\n".join(self.lines).strip()
169
+
170
+
171
+ def effective_env(req: Request) -> dict[str, str]:
172
+ return {**os.environ, **req.env}
173
+
174
+
175
+ def resolve_executable(req: Request, default: str) -> Path | None:
176
+ """Resolve ``req.executable`` or ``default`` on the request's effective ``PATH``."""
177
+ found = shutil.which(req.executable or default, path=effective_env(req).get("PATH"))
178
+ return None if found is None else Path(found)
179
+
180
+
181
+ def request_failure(req: Request, error: Diagnostic, cancel: asyncio.Event | None = None) -> Result:
182
+ """A result for a run that never spawned."""
183
+ verdict = classify(Facts(cancelled=cancel is not None and cancel.is_set(), request_error=error))
184
+ return Result(outcome=verdict.outcome, session_id=req.session_id, error=verdict.error)
185
+
186
+
187
+ def build_result(facts: Facts, *, model: str | None, usage: Usage) -> Result:
188
+ verdict = classify(facts)
189
+ return Result(
190
+ outcome=verdict.outcome,
191
+ text=facts.text,
192
+ session_id=facts.session_id or None,
193
+ model=model,
194
+ usage=usage,
195
+ returncode=facts.returncode,
196
+ error=verdict.error,
197
+ warnings=verdict.warnings,
198
+ )
199
+
200
+
201
+ def process_facts(process: ProcessOutcome, cancel: asyncio.Event) -> Facts:
202
+ return Facts(
203
+ cancelled=cancel.is_set(),
204
+ process_error=process.error_class,
205
+ process_message=process.error,
206
+ returncode=process.returncode,
207
+ )
208
+
209
+
210
+ def deadline_for(req: Request) -> float | None:
211
+ return None if req.timeout_s is None else time.monotonic() + req.timeout_s
212
+
213
+
214
+ class CliDriver:
215
+ """Shared preflight and capability probe for the built-in drivers.
216
+
217
+ Subclasses set ``name``, ``executable``, ``owned_flags``, ``portable_flags`` (supported
218
+ portable field to the flags that set it), ``help_args``, and ``help_markers``.
219
+ """
220
+
221
+ name = ""
222
+ executable = ""
223
+ owned_flags: tuple[str, ...] = ()
224
+ portable_flags: Mapping[str, tuple[str, ...]] = {}
225
+ help_args: tuple[str, ...] = ("--help",)
226
+ help_markers: tuple[str, ...] = ()
227
+
228
+ def __init__(self) -> None:
229
+ self._probe_cache: set[tuple[str, int]] = set()
230
+ self.log = logging.getLogger(f"coding_agent_cli.drivers.{self.name}")
231
+
232
+ def validate(self, req: Request) -> Diagnostic | None:
233
+ """Driver-specific argument checks. Return a diagnostic to reject the request."""
234
+ return None
235
+
236
+ def preflight(self, req: Request, cancel: asyncio.Event) -> tuple[Path, Path] | Result:
237
+ """Validate the request and resolve ``(cwd, executable)``, or return the failure result."""
238
+ error = (
239
+ validate_args(req.args)
240
+ or reject_unsupported(req, self.portable_flags)
241
+ or reject_conflicts(req, self.portable_flags)
242
+ or reject_transport_flags(req.args, self.owned_flags)
243
+ or self.validate(req)
244
+ )
245
+ cwd = req.cwd.expanduser().absolute()
246
+ if error is None:
247
+ if not cwd.is_dir():
248
+ error = Diagnostic(code="bad_cwd", message=f"cwd {os.fspath(cwd)!r} is not a directory")
249
+ if error is not None:
250
+ return request_failure(req, error, cancel)
251
+ executable = resolve_executable(req, self.executable)
252
+ if executable is None:
253
+ name = req.executable or self.executable
254
+ verdict = classify(
255
+ Facts(
256
+ cancelled=cancel.is_set(),
257
+ process_error="missing_cli",
258
+ process_message=f"{name!r} was not found on PATH",
259
+ )
260
+ )
261
+ return Result(outcome=verdict.outcome, session_id=req.session_id, error=verdict.error)
262
+ return cwd, executable
263
+
264
+ def portable_argv(self, req: Request) -> list[str]:
265
+ """Flags for the portable options the request sets, in ``portable_flags`` order."""
266
+ argv: list[str] = []
267
+ for name, flags in self.portable_flags.items():
268
+ value = getattr(req, name)
269
+ if value is not None:
270
+ argv.extend([flags[0], value])
271
+ return argv
272
+
273
+ async def check(self, req: Request, *, refresh: bool = False) -> Capabilities:
274
+ executable = resolve_executable(req, self.executable)
275
+ if executable is None:
276
+ return Capabilities(available=False)
277
+ try:
278
+ key = (os.fspath(executable), executable.stat().st_mtime_ns)
279
+ except OSError:
280
+ return Capabilities(available=False)
281
+ if not refresh and key in self._probe_cache:
282
+ return Capabilities(available=True, executable=executable, structured_streaming=True)
283
+ supported = await self._probe(req, executable)
284
+ if supported:
285
+ self._probe_cache.add(key)
286
+ return Capabilities(available=True, executable=executable, structured_streaming=supported or None)
287
+
288
+ async def _probe(self, req: Request, executable: Path) -> bool:
289
+ lines: list[str] = []
290
+
291
+ async def collect(line: str) -> None:
292
+ lines.append(line)
293
+
294
+ try:
295
+ outcome = await run_process(
296
+ [os.fspath(executable), *self.help_args],
297
+ None,
298
+ dict(req.env),
299
+ None,
300
+ time.monotonic() + PROBE_TIMEOUT_S,
301
+ None,
302
+ collect,
303
+ )
304
+ except OSError:
305
+ return False
306
+ if outcome.error_class:
307
+ return False
308
+ text = "\n".join(lines)
309
+ return all(marker in text for marker in self.help_markers)
310
+
311
+
312
+ __all__ = [
313
+ "NO_USAGE",
314
+ "CliDriver",
315
+ "Diagnostics",
316
+ "EventBuffer",
317
+ "TokenUsage",
318
+ "as_int",
319
+ "build_result",
320
+ "classify_tool_activity",
321
+ "deadline_for",
322
+ "effective_env",
323
+ "parse_json_object",
324
+ "process_facts",
325
+ "read_and_unlink",
326
+ "read_usage",
327
+ "request_failure",
328
+ "resolve_executable",
329
+ "tool_started",
330
+ ]
@@ -0,0 +1,77 @@
1
+ """Flat token counters: six nullable integers, null when unknown, never zero-filled."""
2
+
3
+ from __future__ import annotations
4
+
5
+ from typing import Any
6
+
7
+ from ..types import Usage
8
+
9
+ CONSUMPTION_KEYS = ("input", "output", "cache_read", "cache_write")
10
+
11
+
12
+ def _count(value: Any) -> int | None:
13
+ if isinstance(value, bool) or not isinstance(value, int) or value < 0:
14
+ return None
15
+ return value
16
+
17
+
18
+ class TokenCounters:
19
+ """Responses sum across ids; CLI totals and context override them per field."""
20
+
21
+ def __init__(self) -> None:
22
+ self._responses: dict[str, dict[str, int]] = {}
23
+ self._totals: dict[str, int] = {}
24
+ self._context: dict[str, int] = {}
25
+
26
+ def record_response(
27
+ self,
28
+ response_id: str,
29
+ *,
30
+ input: Any = None,
31
+ output: Any = None,
32
+ cache_read: Any = None,
33
+ cache_write: Any = None,
34
+ ) -> None:
35
+ values = {"input": input, "output": output, "cache_read": cache_read, "cache_write": cache_write}
36
+ counted = {key: _count(value) for key, value in values.items()}
37
+ self._responses[response_id] = {key: value for key, value in counted.items() if value is not None}
38
+
39
+ def record_totals(
40
+ self,
41
+ *,
42
+ input: Any = None,
43
+ output: Any = None,
44
+ cache_read: Any = None,
45
+ cache_write: Any = None,
46
+ ) -> None:
47
+ values = {"input": input, "output": output, "cache_read": cache_read, "cache_write": cache_write}
48
+ for key, value in values.items():
49
+ count = _count(value)
50
+ if count is not None:
51
+ self._totals[key] = count
52
+
53
+ def record_context(self, *, used: Any = None, window: Any = None) -> None:
54
+ for key, value in (("context_used", used), ("context_window", window)):
55
+ count = _count(value)
56
+ if count is not None:
57
+ self._context[key] = count
58
+
59
+ def clear_context(self) -> None:
60
+ self._context.pop("context_used", None)
61
+
62
+ def usage(self) -> Usage:
63
+ counters: dict[str, int | None] = {}
64
+ for key in CONSUMPTION_KEYS:
65
+ if key in self._totals:
66
+ counters[key] = self._totals[key]
67
+ else:
68
+ values = [response[key] for response in self._responses.values() if key in response]
69
+ counters[key] = sum(values) if values else None
70
+ return Usage(
71
+ context_used=self._context.get("context_used"),
72
+ context_window=self._context.get("context_window"),
73
+ **counters,
74
+ )
75
+
76
+
77
+ __all__ = ["TokenCounters"]