agentabi 0.2.0__tar.gz → 0.2.2__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.
- {agentabi-0.2.0/src/agentabi.egg-info → agentabi-0.2.2}/PKG-INFO +10 -5
- {agentabi-0.2.0 → agentabi-0.2.2}/README.md +7 -2
- {agentabi-0.2.0 → agentabi-0.2.2}/pyproject.toml +2 -2
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/__init__.py +24 -1
- agentabi-0.2.2/src/agentabi/middleware/__init__.py +119 -0
- agentabi-0.2.2/src/agentabi/middleware/logging.py +112 -0
- agentabi-0.2.2/src/agentabi/middleware/timeout.py +84 -0
- agentabi-0.2.2/src/agentabi/middleware/usage.py +134 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/base.py +51 -1
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/claude_native.py +55 -1
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/codex_native.py +72 -18
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/codex_sdk.py +45 -25
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/gemini_native.py +48 -7
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/opencode_native.py +51 -13
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/session.py +94 -3
- {agentabi-0.2.0 → agentabi-0.2.2/src/agentabi.egg-info}/PKG-INFO +10 -5
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/SOURCES.txt +5 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/requires.txt +2 -2
- agentabi-0.2.2/tests/test_middleware.py +512 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/LICENSE +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/setup.cfg +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/auto_detect.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/__init__.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/claude_sdk.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/gemini_sdk.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/providers/registry.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/py.typed +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/__init__.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/__init__.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/capabilities.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/events.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/helpers.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/permissions.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/session.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/task.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi/types/ir/type_guards.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/dependency_links.txt +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/src/agentabi.egg-info/top_level.txt +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/tests/test_ir_types.py +0 -0
- {agentabi-0.2.0 → agentabi-0.2.2}/tests/test_session.py +0 -0
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
Metadata-Version: 2.4
|
|
2
2
|
Name: agentabi
|
|
3
|
-
Version: 0.2.
|
|
3
|
+
Version: 0.2.2
|
|
4
4
|
Summary: Unified async Python API for driving coding agent CLIs — Claude Code, Codex, Gemini CLI, and more
|
|
5
5
|
Author-email: Peng Ding <pding.dp@foxmail.com>
|
|
6
6
|
License-Expression: MIT
|
|
@@ -38,14 +38,19 @@ Provides-Extra: dev
|
|
|
38
38
|
Requires-Dist: pytest>=7.0.0; extra == "dev"
|
|
39
39
|
Requires-Dist: pytest-asyncio>=0.21.0; extra == "dev"
|
|
40
40
|
Requires-Dist: pytest-cov; extra == "dev"
|
|
41
|
-
Requires-Dist: ruff
|
|
42
|
-
Requires-Dist: ty
|
|
41
|
+
Requires-Dist: ruff==0.15.20; extra == "dev"
|
|
42
|
+
Requires-Dist: ty==0.0.54; extra == "dev"
|
|
43
43
|
Requires-Dist: build>=1.2.2.post1; extra == "dev"
|
|
44
44
|
Requires-Dist: twine>=6.1.0; extra == "dev"
|
|
45
45
|
Dynamic: license-file
|
|
46
46
|
|
|
47
47
|
# agentabi
|
|
48
48
|
|
|
49
|
+
[](https://pypi.org/project/agentabi/)
|
|
50
|
+
[](https://github.com/Oaklight/agentabi/releases/latest)
|
|
51
|
+
[](https://github.com/Oaklight/agentabi/actions/workflows/ci.yml)
|
|
52
|
+
[](https://opensource.org/licenses/MIT)
|
|
53
|
+
|
|
49
54
|
Unified interface layer for agentic coding CLIs.
|
|
50
55
|
|
|
51
56
|
One interface. Any coding agent.
|
|
@@ -138,10 +143,10 @@ print(caps["supports_streaming"]) # True
|
|
|
138
143
|
|
|
139
144
|
```
|
|
140
145
|
agentabi → Agent CLI unified interface → like an OS ABI
|
|
141
|
-
|
|
146
|
+
llm-rosetta → LLM API format conversion → like a compiler IR
|
|
142
147
|
```
|
|
143
148
|
|
|
144
|
-
- [
|
|
149
|
+
- [llm-rosetta](https://github.com/Oaklight/llm-rosetta) — LLM API format conversion layer for translating between LLM provider API formats (OpenAI, Anthropic, Google)
|
|
145
150
|
|
|
146
151
|
## License
|
|
147
152
|
|
|
@@ -1,5 +1,10 @@
|
|
|
1
1
|
# agentabi
|
|
2
2
|
|
|
3
|
+
[](https://pypi.org/project/agentabi/)
|
|
4
|
+
[](https://github.com/Oaklight/agentabi/releases/latest)
|
|
5
|
+
[](https://github.com/Oaklight/agentabi/actions/workflows/ci.yml)
|
|
6
|
+
[](https://opensource.org/licenses/MIT)
|
|
7
|
+
|
|
3
8
|
Unified interface layer for agentic coding CLIs.
|
|
4
9
|
|
|
5
10
|
One interface. Any coding agent.
|
|
@@ -92,10 +97,10 @@ print(caps["supports_streaming"]) # True
|
|
|
92
97
|
|
|
93
98
|
```
|
|
94
99
|
agentabi → Agent CLI unified interface → like an OS ABI
|
|
95
|
-
|
|
100
|
+
llm-rosetta → LLM API format conversion → like a compiler IR
|
|
96
101
|
```
|
|
97
102
|
|
|
98
|
-
- [
|
|
103
|
+
- [llm-rosetta](https://github.com/Oaklight/llm-rosetta) — LLM API format conversion layer for translating between LLM provider API formats (OpenAI, Anthropic, Google)
|
|
99
104
|
|
|
100
105
|
## License
|
|
101
106
|
|
|
@@ -19,6 +19,14 @@ Sync convenience:
|
|
|
19
19
|
from agentabi import run_sync
|
|
20
20
|
result = run_sync(prompt="List Python files", agent="claude_code")
|
|
21
21
|
|
|
22
|
+
Middleware:
|
|
23
|
+
from agentabi.middleware import LoggingMiddleware, TimeoutMiddleware
|
|
24
|
+
|
|
25
|
+
session = Session(
|
|
26
|
+
agent="claude_code",
|
|
27
|
+
middleware=[LoggingMiddleware(), TimeoutMiddleware(120)],
|
|
28
|
+
)
|
|
29
|
+
|
|
22
30
|
Discovery:
|
|
23
31
|
from agentabi import detect_agents, get_agent_capabilities
|
|
24
32
|
agents = detect_agents()
|
|
@@ -26,6 +34,14 @@ Discovery:
|
|
|
26
34
|
"""
|
|
27
35
|
|
|
28
36
|
from .auto_detect import detect_agents, get_agent_capabilities, get_default_agent
|
|
37
|
+
from .middleware import (
|
|
38
|
+
LoggingMiddleware,
|
|
39
|
+
Middleware,
|
|
40
|
+
StreamHandler,
|
|
41
|
+
TimeoutMiddleware,
|
|
42
|
+
UsageMeterMiddleware,
|
|
43
|
+
chain_middleware,
|
|
44
|
+
)
|
|
29
45
|
from .providers.base import Provider
|
|
30
46
|
from .providers.registry import AgentNotAvailable, get_provider
|
|
31
47
|
from .session import Session, run_sync
|
|
@@ -53,7 +69,7 @@ from .types.ir import (
|
|
|
53
69
|
UsageEvent,
|
|
54
70
|
)
|
|
55
71
|
|
|
56
|
-
__version__ = "0.2.
|
|
72
|
+
__version__ = "0.2.2"
|
|
57
73
|
|
|
58
74
|
__all__ = [
|
|
59
75
|
# Consumer API
|
|
@@ -92,4 +108,11 @@ __all__ = [
|
|
|
92
108
|
"PermissionConfig",
|
|
93
109
|
"PermissionLevel",
|
|
94
110
|
"PermissionRequest",
|
|
111
|
+
# Middleware
|
|
112
|
+
"Middleware",
|
|
113
|
+
"StreamHandler",
|
|
114
|
+
"chain_middleware",
|
|
115
|
+
"LoggingMiddleware",
|
|
116
|
+
"TimeoutMiddleware",
|
|
117
|
+
"UsageMeterMiddleware",
|
|
95
118
|
]
|
|
@@ -0,0 +1,119 @@
|
|
|
1
|
+
"""
|
|
2
|
+
agentabi - Middleware Pipeline
|
|
3
|
+
|
|
4
|
+
Composable middleware for intercepting Session.stream() calls.
|
|
5
|
+
Middleware can inspect/modify TaskConfig, filter/transform IREvents,
|
|
6
|
+
and run side effects (logging, metering, timeouts).
|
|
7
|
+
|
|
8
|
+
Architecture:
|
|
9
|
+
Middleware follows the onion model (like ASGI). Each middleware
|
|
10
|
+
wraps a ``StreamHandler`` and returns a new ``StreamHandler``.
|
|
11
|
+
The first middleware in the list is the outermost wrapper.
|
|
12
|
+
|
|
13
|
+
Quick start::
|
|
14
|
+
|
|
15
|
+
from agentabi.middleware import LoggingMiddleware, TimeoutMiddleware
|
|
16
|
+
|
|
17
|
+
session = Session(
|
|
18
|
+
agent="claude_code",
|
|
19
|
+
middleware=[LoggingMiddleware(), TimeoutMiddleware(120)],
|
|
20
|
+
)
|
|
21
|
+
|
|
22
|
+
Custom middleware::
|
|
23
|
+
|
|
24
|
+
from agentabi.middleware import Middleware, StreamHandler
|
|
25
|
+
|
|
26
|
+
class MyMiddleware:
|
|
27
|
+
def __call__(self, handler: StreamHandler) -> StreamHandler:
|
|
28
|
+
async def wrapper(task):
|
|
29
|
+
task = {**task, "prompt": task["prompt"].strip()}
|
|
30
|
+
async for event in handler(task):
|
|
31
|
+
yield event
|
|
32
|
+
return wrapper
|
|
33
|
+
"""
|
|
34
|
+
|
|
35
|
+
from __future__ import annotations
|
|
36
|
+
|
|
37
|
+
from collections.abc import AsyncIterator, Sequence
|
|
38
|
+
from typing import Callable
|
|
39
|
+
|
|
40
|
+
from ..types.ir.events import IREvent
|
|
41
|
+
from ..types.ir.task import TaskConfig
|
|
42
|
+
from .logging import LoggingMiddleware
|
|
43
|
+
from .timeout import TimeoutMiddleware
|
|
44
|
+
from .usage import UsageMeterMiddleware
|
|
45
|
+
|
|
46
|
+
# ============================================================================
|
|
47
|
+
# Core type aliases
|
|
48
|
+
# ============================================================================
|
|
49
|
+
|
|
50
|
+
StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
|
|
51
|
+
"""A callable that takes a TaskConfig and returns an async iterator of IREvents.
|
|
52
|
+
|
|
53
|
+
This is the fundamental unit of the middleware pipeline. The provider's
|
|
54
|
+
``stream()`` method is the innermost StreamHandler; each middleware wraps
|
|
55
|
+
it with additional behavior.
|
|
56
|
+
"""
|
|
57
|
+
|
|
58
|
+
Middleware = Callable[[StreamHandler], StreamHandler]
|
|
59
|
+
"""A callable that wraps a StreamHandler with additional behavior.
|
|
60
|
+
|
|
61
|
+
Middleware receives the next handler in the chain and returns a new
|
|
62
|
+
handler. The wrapper can:
|
|
63
|
+
|
|
64
|
+
- Modify TaskConfig before passing to the next handler (copy-on-mutate
|
|
65
|
+
recommended: ``task = {**task, "key": new_value}``)
|
|
66
|
+
- Filter or transform yielded IREvents
|
|
67
|
+
- Run side effects (logging, timing, metering)
|
|
68
|
+
- Short-circuit the pipeline (e.g., timeout)
|
|
69
|
+
"""
|
|
70
|
+
|
|
71
|
+
|
|
72
|
+
# ============================================================================
|
|
73
|
+
# Pipeline chaining
|
|
74
|
+
# ============================================================================
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
def chain_middleware(
|
|
78
|
+
handler: StreamHandler,
|
|
79
|
+
middleware: Sequence[Middleware],
|
|
80
|
+
) -> StreamHandler:
|
|
81
|
+
"""Chain middleware around a base handler using the onion model.
|
|
82
|
+
|
|
83
|
+
The first middleware in the sequence becomes the outermost wrapper.
|
|
84
|
+
When the pipeline executes, control flows inward through each
|
|
85
|
+
middleware, reaches the base handler, then flows back outward.
|
|
86
|
+
|
|
87
|
+
Args:
|
|
88
|
+
handler: The base StreamHandler (typically provider.stream).
|
|
89
|
+
middleware: Sequence of middleware to apply. Empty sequence
|
|
90
|
+
returns the handler unchanged.
|
|
91
|
+
|
|
92
|
+
Returns:
|
|
93
|
+
A new StreamHandler with all middleware applied.
|
|
94
|
+
|
|
95
|
+
Example::
|
|
96
|
+
|
|
97
|
+
pipeline = chain_middleware(
|
|
98
|
+
provider.stream,
|
|
99
|
+
[LoggingMiddleware(), TimeoutMiddleware(30)],
|
|
100
|
+
)
|
|
101
|
+
async for event in pipeline(task):
|
|
102
|
+
...
|
|
103
|
+
"""
|
|
104
|
+
for mw in reversed(middleware):
|
|
105
|
+
handler = mw(handler)
|
|
106
|
+
return handler
|
|
107
|
+
|
|
108
|
+
|
|
109
|
+
__all__ = [
|
|
110
|
+
# Core types
|
|
111
|
+
"StreamHandler",
|
|
112
|
+
"Middleware",
|
|
113
|
+
# Pipeline
|
|
114
|
+
"chain_middleware",
|
|
115
|
+
# Built-in middleware
|
|
116
|
+
"LoggingMiddleware",
|
|
117
|
+
"TimeoutMiddleware",
|
|
118
|
+
"UsageMeterMiddleware",
|
|
119
|
+
]
|
|
@@ -0,0 +1,112 @@
|
|
|
1
|
+
"""
|
|
2
|
+
agentabi - Logging Middleware
|
|
3
|
+
|
|
4
|
+
Logs task execution and IR events at configurable levels.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import logging
|
|
10
|
+
import time
|
|
11
|
+
from collections.abc import AsyncIterator
|
|
12
|
+
from typing import Callable
|
|
13
|
+
|
|
14
|
+
from ..types.ir.events import IREvent
|
|
15
|
+
from ..types.ir.task import TaskConfig
|
|
16
|
+
|
|
17
|
+
StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
|
|
18
|
+
|
|
19
|
+
_DEFAULT_LOGGER = logging.getLogger("agentabi.middleware")
|
|
20
|
+
|
|
21
|
+
|
|
22
|
+
class LoggingMiddleware:
|
|
23
|
+
"""Middleware that logs task execution and IR events.
|
|
24
|
+
|
|
25
|
+
Logs task start/end at INFO level and individual events at DEBUG
|
|
26
|
+
level. Uses the ``agentabi.middleware`` logger by default.
|
|
27
|
+
|
|
28
|
+
Args:
|
|
29
|
+
logger: Logger instance to use. Defaults to ``agentabi.middleware``.
|
|
30
|
+
level: Log level for event-level messages. Defaults to DEBUG.
|
|
31
|
+
log_events: Whether to log individual IR events. Defaults to True.
|
|
32
|
+
log_event_types: If set, only log events whose ``type`` is in
|
|
33
|
+
this set. Defaults to None (log all events).
|
|
34
|
+
|
|
35
|
+
Example::
|
|
36
|
+
|
|
37
|
+
import logging
|
|
38
|
+
logging.basicConfig(level=logging.DEBUG)
|
|
39
|
+
|
|
40
|
+
session = Session(
|
|
41
|
+
agent="claude_code",
|
|
42
|
+
middleware=[LoggingMiddleware()],
|
|
43
|
+
)
|
|
44
|
+
"""
|
|
45
|
+
|
|
46
|
+
def __init__(
|
|
47
|
+
self,
|
|
48
|
+
*,
|
|
49
|
+
logger: logging.Logger | None = None,
|
|
50
|
+
level: int = logging.DEBUG,
|
|
51
|
+
log_events: bool = True,
|
|
52
|
+
log_event_types: set[str] | None = None,
|
|
53
|
+
) -> None:
|
|
54
|
+
self._logger = logger or _DEFAULT_LOGGER
|
|
55
|
+
self._level = level
|
|
56
|
+
self._log_events = log_events
|
|
57
|
+
self._log_event_types = log_event_types
|
|
58
|
+
|
|
59
|
+
def __call__(self, handler: StreamHandler) -> StreamHandler:
|
|
60
|
+
"""Wrap a StreamHandler with logging."""
|
|
61
|
+
logger = self._logger
|
|
62
|
+
level = self._level
|
|
63
|
+
log_events = self._log_events
|
|
64
|
+
log_event_types = self._log_event_types
|
|
65
|
+
|
|
66
|
+
async def wrapper(task: TaskConfig) -> AsyncIterator[IREvent]:
|
|
67
|
+
prompt_preview = task["prompt"][:80]
|
|
68
|
+
agent = task.get("agent", "unknown")
|
|
69
|
+
logger.info("Task started: agent=%s prompt=%r", agent, prompt_preview)
|
|
70
|
+
|
|
71
|
+
start = time.monotonic()
|
|
72
|
+
event_count = 0
|
|
73
|
+
failed = False
|
|
74
|
+
|
|
75
|
+
try:
|
|
76
|
+
async for event in handler(task):
|
|
77
|
+
event_count += 1
|
|
78
|
+
if log_events:
|
|
79
|
+
etype = event.get("type", "unknown")
|
|
80
|
+
if log_event_types is None or etype in log_event_types:
|
|
81
|
+
logger.log(level, "Event #%d: type=%s", event_count, etype)
|
|
82
|
+
yield event
|
|
83
|
+
except Exception:
|
|
84
|
+
failed = True
|
|
85
|
+
raise
|
|
86
|
+
finally:
|
|
87
|
+
elapsed = time.monotonic() - start
|
|
88
|
+
if failed:
|
|
89
|
+
logger.error(
|
|
90
|
+
"Task failed: agent=%s elapsed=%.2fs events=%d",
|
|
91
|
+
agent,
|
|
92
|
+
elapsed,
|
|
93
|
+
event_count,
|
|
94
|
+
)
|
|
95
|
+
else:
|
|
96
|
+
logger.info(
|
|
97
|
+
"Task completed: agent=%s elapsed=%.2fs events=%d",
|
|
98
|
+
agent,
|
|
99
|
+
elapsed,
|
|
100
|
+
event_count,
|
|
101
|
+
)
|
|
102
|
+
|
|
103
|
+
return wrapper
|
|
104
|
+
|
|
105
|
+
def __repr__(self) -> str:
|
|
106
|
+
return (
|
|
107
|
+
f"LoggingMiddleware(logger={self._logger.name!r}, "
|
|
108
|
+
f"level={self._level}, log_events={self._log_events})"
|
|
109
|
+
)
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
__all__ = ["LoggingMiddleware"]
|
|
@@ -0,0 +1,84 @@
|
|
|
1
|
+
"""
|
|
2
|
+
agentabi - Timeout Middleware
|
|
3
|
+
|
|
4
|
+
Enforces a wall-clock timeout on streaming task execution.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import asyncio
|
|
10
|
+
import time
|
|
11
|
+
from collections.abc import AsyncIterator
|
|
12
|
+
from typing import Callable
|
|
13
|
+
|
|
14
|
+
from ..types.ir.events import IREvent
|
|
15
|
+
from ..types.ir.task import TaskConfig
|
|
16
|
+
|
|
17
|
+
StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
class TimeoutMiddleware:
|
|
21
|
+
"""Middleware that enforces a wall-clock timeout on task execution.
|
|
22
|
+
|
|
23
|
+
Wraps each ``__anext__`` call on the event stream with
|
|
24
|
+
``asyncio.wait_for``, so the timeout fires even if the stream
|
|
25
|
+
stalls mid-iteration (not just between events).
|
|
26
|
+
|
|
27
|
+
Raises ``asyncio.TimeoutError`` if the total elapsed time exceeds
|
|
28
|
+
the configured timeout.
|
|
29
|
+
|
|
30
|
+
Args:
|
|
31
|
+
timeout: Maximum wall-clock seconds for the entire stream.
|
|
32
|
+
Must be positive.
|
|
33
|
+
|
|
34
|
+
Example::
|
|
35
|
+
|
|
36
|
+
session = Session(
|
|
37
|
+
agent="claude_code",
|
|
38
|
+
middleware=[TimeoutMiddleware(120)],
|
|
39
|
+
)
|
|
40
|
+
|
|
41
|
+
try:
|
|
42
|
+
result = await session.run(prompt="Long task")
|
|
43
|
+
except asyncio.TimeoutError:
|
|
44
|
+
print("Task timed out after 120s")
|
|
45
|
+
|
|
46
|
+
Note:
|
|
47
|
+
The timeout is per-stream invocation, not cumulative across
|
|
48
|
+
multiple calls. Each ``stream()`` or ``run()`` gets a fresh
|
|
49
|
+
timeout budget.
|
|
50
|
+
"""
|
|
51
|
+
|
|
52
|
+
def __init__(self, timeout: float) -> None:
|
|
53
|
+
if timeout <= 0:
|
|
54
|
+
raise ValueError(f"timeout must be positive, got {timeout}")
|
|
55
|
+
self._timeout = timeout
|
|
56
|
+
|
|
57
|
+
def __call__(self, handler: StreamHandler) -> StreamHandler:
|
|
58
|
+
"""Wrap a StreamHandler with timeout enforcement."""
|
|
59
|
+
timeout = self._timeout
|
|
60
|
+
|
|
61
|
+
async def wrapper(task: TaskConfig) -> AsyncIterator[IREvent]:
|
|
62
|
+
start = time.monotonic()
|
|
63
|
+
aiter = handler(task).__aiter__()
|
|
64
|
+
|
|
65
|
+
while True:
|
|
66
|
+
elapsed = time.monotonic() - start
|
|
67
|
+
remaining = timeout - elapsed
|
|
68
|
+
if remaining <= 0:
|
|
69
|
+
raise asyncio.TimeoutError(
|
|
70
|
+
f"Task exceeded timeout of {timeout}s (elapsed: {elapsed:.1f}s)"
|
|
71
|
+
)
|
|
72
|
+
try:
|
|
73
|
+
event = await asyncio.wait_for(aiter.__anext__(), timeout=remaining)
|
|
74
|
+
except StopAsyncIteration:
|
|
75
|
+
break
|
|
76
|
+
yield event
|
|
77
|
+
|
|
78
|
+
return wrapper
|
|
79
|
+
|
|
80
|
+
def __repr__(self) -> str:
|
|
81
|
+
return f"TimeoutMiddleware(timeout={self._timeout})"
|
|
82
|
+
|
|
83
|
+
|
|
84
|
+
__all__ = ["TimeoutMiddleware"]
|
|
@@ -0,0 +1,134 @@
|
|
|
1
|
+
"""
|
|
2
|
+
agentabi - Usage Meter Middleware
|
|
3
|
+
|
|
4
|
+
Tracks cumulative token usage and cost across multiple task executions.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from __future__ import annotations
|
|
8
|
+
|
|
9
|
+
import threading
|
|
10
|
+
from collections.abc import AsyncIterator
|
|
11
|
+
from typing import Callable
|
|
12
|
+
|
|
13
|
+
from typing_extensions import TypedDict
|
|
14
|
+
|
|
15
|
+
from ..types.ir.events import IREvent
|
|
16
|
+
from ..types.ir.task import TaskConfig
|
|
17
|
+
|
|
18
|
+
StreamHandler = Callable[[TaskConfig], AsyncIterator[IREvent]]
|
|
19
|
+
|
|
20
|
+
|
|
21
|
+
class UsageSummary(TypedDict, total=False):
|
|
22
|
+
"""Cumulative usage statistics across multiple task executions."""
|
|
23
|
+
|
|
24
|
+
call_count: int
|
|
25
|
+
total_input_tokens: int
|
|
26
|
+
total_output_tokens: int
|
|
27
|
+
total_cache_read_tokens: int
|
|
28
|
+
total_cache_creation_tokens: int
|
|
29
|
+
total_tokens: int
|
|
30
|
+
total_cost_usd: float
|
|
31
|
+
|
|
32
|
+
|
|
33
|
+
class UsageMeterMiddleware:
|
|
34
|
+
"""Middleware that accumulates token usage and cost across calls.
|
|
35
|
+
|
|
36
|
+
Watches for ``usage`` events in the stream and accumulates token
|
|
37
|
+
counts and cost. The accumulated totals are available via the
|
|
38
|
+
``summary`` property at any time.
|
|
39
|
+
|
|
40
|
+
Thread-safe: uses a lock for counter updates so the summary can
|
|
41
|
+
be read from any thread.
|
|
42
|
+
|
|
43
|
+
Example::
|
|
44
|
+
|
|
45
|
+
meter = UsageMeterMiddleware()
|
|
46
|
+
session = Session(agent="claude_code", middleware=[meter])
|
|
47
|
+
|
|
48
|
+
await session.run(prompt="Task 1")
|
|
49
|
+
await session.run(prompt="Task 2")
|
|
50
|
+
|
|
51
|
+
print(meter.summary)
|
|
52
|
+
# {'call_count': 2, 'total_input_tokens': 1500, ...}
|
|
53
|
+
|
|
54
|
+
meter.reset()
|
|
55
|
+
"""
|
|
56
|
+
|
|
57
|
+
def __init__(self) -> None:
|
|
58
|
+
self._lock = threading.Lock()
|
|
59
|
+
self._call_count = 0
|
|
60
|
+
self._input_tokens = 0
|
|
61
|
+
self._output_tokens = 0
|
|
62
|
+
self._cache_read_tokens = 0
|
|
63
|
+
self._cache_creation_tokens = 0
|
|
64
|
+
self._total_tokens = 0
|
|
65
|
+
self._cost_usd = 0.0
|
|
66
|
+
|
|
67
|
+
@property
|
|
68
|
+
def summary(self) -> UsageSummary:
|
|
69
|
+
"""Return cumulative usage statistics.
|
|
70
|
+
|
|
71
|
+
Returns:
|
|
72
|
+
A UsageSummary dict with all accumulated counters.
|
|
73
|
+
"""
|
|
74
|
+
with self._lock:
|
|
75
|
+
result: UsageSummary = {
|
|
76
|
+
"call_count": self._call_count,
|
|
77
|
+
"total_input_tokens": self._input_tokens,
|
|
78
|
+
"total_output_tokens": self._output_tokens,
|
|
79
|
+
"total_tokens": self._total_tokens,
|
|
80
|
+
"total_cost_usd": self._cost_usd,
|
|
81
|
+
}
|
|
82
|
+
if self._cache_read_tokens:
|
|
83
|
+
result["total_cache_read_tokens"] = self._cache_read_tokens
|
|
84
|
+
if self._cache_creation_tokens:
|
|
85
|
+
result["total_cache_creation_tokens"] = self._cache_creation_tokens
|
|
86
|
+
return result
|
|
87
|
+
|
|
88
|
+
def reset(self) -> None:
|
|
89
|
+
"""Reset all accumulated counters to zero."""
|
|
90
|
+
with self._lock:
|
|
91
|
+
self._call_count = 0
|
|
92
|
+
self._input_tokens = 0
|
|
93
|
+
self._output_tokens = 0
|
|
94
|
+
self._cache_read_tokens = 0
|
|
95
|
+
self._cache_creation_tokens = 0
|
|
96
|
+
self._total_tokens = 0
|
|
97
|
+
self._cost_usd = 0.0
|
|
98
|
+
|
|
99
|
+
def __call__(self, handler: StreamHandler) -> StreamHandler:
|
|
100
|
+
"""Wrap a StreamHandler with usage metering."""
|
|
101
|
+
|
|
102
|
+
async def wrapper(task: TaskConfig) -> AsyncIterator[IREvent]:
|
|
103
|
+
with self._lock:
|
|
104
|
+
self._call_count += 1
|
|
105
|
+
|
|
106
|
+
async for event in handler(task):
|
|
107
|
+
if event.get("type") == "usage":
|
|
108
|
+
usage = event.get("usage", {})
|
|
109
|
+
cost = event.get("cost_usd")
|
|
110
|
+
with self._lock:
|
|
111
|
+
self._input_tokens += usage.get("input_tokens", 0)
|
|
112
|
+
self._output_tokens += usage.get("output_tokens", 0)
|
|
113
|
+
self._cache_read_tokens += usage.get("cache_read_tokens", 0)
|
|
114
|
+
self._cache_creation_tokens += usage.get(
|
|
115
|
+
"cache_creation_tokens", 0
|
|
116
|
+
)
|
|
117
|
+
self._total_tokens += usage.get("total_tokens", 0)
|
|
118
|
+
if cost is not None:
|
|
119
|
+
self._cost_usd += cost
|
|
120
|
+
yield event
|
|
121
|
+
|
|
122
|
+
return wrapper
|
|
123
|
+
|
|
124
|
+
def __repr__(self) -> str:
|
|
125
|
+
return (
|
|
126
|
+
f"UsageMeterMiddleware(calls={self._call_count}, "
|
|
127
|
+
f"tokens={self._total_tokens}, cost=${self._cost_usd:.4f})"
|
|
128
|
+
)
|
|
129
|
+
|
|
130
|
+
|
|
131
|
+
__all__ = [
|
|
132
|
+
"UsageMeterMiddleware",
|
|
133
|
+
"UsageSummary",
|
|
134
|
+
]
|
|
@@ -7,13 +7,14 @@ abstraction layer between Session and agent CLIs/SDKs.
|
|
|
7
7
|
|
|
8
8
|
from __future__ import annotations
|
|
9
9
|
|
|
10
|
+
import asyncio
|
|
10
11
|
from collections.abc import AsyncIterator
|
|
11
12
|
from typing import Any
|
|
12
13
|
|
|
13
14
|
from typing_extensions import Protocol, runtime_checkable
|
|
14
15
|
|
|
15
16
|
from ..types.ir.capabilities import AgentCapabilities
|
|
16
|
-
from ..types.ir.events import IREvent, UsageInfo
|
|
17
|
+
from ..types.ir.events import ErrorEvent, IREvent, UsageInfo
|
|
17
18
|
from ..types.ir.session import SessionResult, SessionStatus
|
|
18
19
|
from ..types.ir.task import TaskConfig
|
|
19
20
|
|
|
@@ -71,6 +72,54 @@ class Provider(Protocol):
|
|
|
71
72
|
...
|
|
72
73
|
|
|
73
74
|
|
|
75
|
+
async def collect_subprocess_errors(
|
|
76
|
+
proc: asyncio.subprocess.Process,
|
|
77
|
+
*,
|
|
78
|
+
timed_out: bool = False,
|
|
79
|
+
timeout_seconds: float | None = None,
|
|
80
|
+
) -> list[ErrorEvent]:
|
|
81
|
+
"""Collect stderr and exit-code errors after a subprocess finishes.
|
|
82
|
+
|
|
83
|
+
Call after the stdout reading loop ends and ``await proc.wait()``
|
|
84
|
+
has been called so that ``proc.returncode`` is set.
|
|
85
|
+
|
|
86
|
+
Args:
|
|
87
|
+
proc: The completed subprocess.
|
|
88
|
+
timed_out: Whether the process was killed due to timeout.
|
|
89
|
+
timeout_seconds: The configured timeout value (for the error message).
|
|
90
|
+
|
|
91
|
+
Returns:
|
|
92
|
+
A (possibly empty) list of ErrorEvents to yield.
|
|
93
|
+
"""
|
|
94
|
+
events: list[ErrorEvent] = []
|
|
95
|
+
stderr_text = ""
|
|
96
|
+
|
|
97
|
+
if proc.stderr:
|
|
98
|
+
stderr_data = await proc.stderr.read()
|
|
99
|
+
if stderr_data:
|
|
100
|
+
stderr_text = stderr_data.decode().strip()
|
|
101
|
+
if stderr_text:
|
|
102
|
+
events.append(ErrorEvent(type="error", error=stderr_text))
|
|
103
|
+
|
|
104
|
+
if timed_out:
|
|
105
|
+
events.append(
|
|
106
|
+
ErrorEvent(
|
|
107
|
+
type="error",
|
|
108
|
+
error=f"Process timed out after {timeout_seconds}s",
|
|
109
|
+
is_fatal=True,
|
|
110
|
+
)
|
|
111
|
+
)
|
|
112
|
+
elif proc.returncode:
|
|
113
|
+
events.append(
|
|
114
|
+
ErrorEvent(
|
|
115
|
+
type="error",
|
|
116
|
+
error=f"Process exited with code {proc.returncode}",
|
|
117
|
+
)
|
|
118
|
+
)
|
|
119
|
+
|
|
120
|
+
return events
|
|
121
|
+
|
|
122
|
+
|
|
74
123
|
class _RunState:
|
|
75
124
|
"""Mutable accumulator for default_run event aggregation."""
|
|
76
125
|
|
|
@@ -163,5 +212,6 @@ async def default_run(provider: Any, task: TaskConfig) -> SessionResult:
|
|
|
163
212
|
|
|
164
213
|
__all__ = [
|
|
165
214
|
"Provider",
|
|
215
|
+
"collect_subprocess_errors",
|
|
166
216
|
"default_run",
|
|
167
217
|
]
|