a3s-code 0.4.4__tar.gz → 0.4.5__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.
- {a3s_code-0.4.4 → a3s_code-0.4.5}/PKG-INFO +1 -1
- {a3s_code-0.4.4 → a3s_code-0.4.5}/a3s_code/__init__.py +30 -1
- {a3s_code-0.4.4 → a3s_code-0.4.5}/a3s_code/client.py +67 -1
- a3s_code-0.4.5/a3s_code/provider.py +63 -0
- a3s_code-0.4.5/a3s_code/session.py +476 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/pyproject.toml +1 -1
- {a3s_code-0.4.4 → a3s_code-0.4.5}/.gitignore +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/LICENSE +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/MANIFEST.in +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/README.md +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/a3s_code/py.typed +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/a3s_code/types.py +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/justfile +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/proto/code_agent.proto +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/pytest.ini +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/tests/conftest.py +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/tests/test_client.py +0 -0
- {a3s_code-0.4.4 → a3s_code-0.4.5}/tests/test_e2e.py +0 -0
|
@@ -5,6 +5,20 @@ A Python client for the A3S Code Agent gRPC service.
|
|
|
5
5
|
"""
|
|
6
6
|
|
|
7
7
|
from .client import A3sClient
|
|
8
|
+
from .provider import ModelRef, create_provider
|
|
9
|
+
from .session import (
|
|
10
|
+
CodeSession,
|
|
11
|
+
SendResult,
|
|
12
|
+
StepResult,
|
|
13
|
+
SessionStats,
|
|
14
|
+
AgentLoopEvent,
|
|
15
|
+
TextEvent,
|
|
16
|
+
ToolCallEvent,
|
|
17
|
+
ToolResultEvent,
|
|
18
|
+
StepFinishEvent,
|
|
19
|
+
ErrorEvent,
|
|
20
|
+
DoneEvent,
|
|
21
|
+
)
|
|
8
22
|
from .types import (
|
|
9
23
|
# Enums
|
|
10
24
|
HealthStatus,
|
|
@@ -87,9 +101,24 @@ from .types import (
|
|
|
87
101
|
QueueStats,
|
|
88
102
|
)
|
|
89
103
|
|
|
90
|
-
__version__ = "0.4.
|
|
104
|
+
__version__ = "0.4.4"
|
|
91
105
|
__all__ = [
|
|
92
106
|
"A3sClient",
|
|
107
|
+
# Provider
|
|
108
|
+
"ModelRef",
|
|
109
|
+
"create_provider",
|
|
110
|
+
# Session (high-level)
|
|
111
|
+
"CodeSession",
|
|
112
|
+
"SendResult",
|
|
113
|
+
"StepResult",
|
|
114
|
+
"SessionStats",
|
|
115
|
+
"AgentLoopEvent",
|
|
116
|
+
"TextEvent",
|
|
117
|
+
"ToolCallEvent",
|
|
118
|
+
"ToolResultEvent",
|
|
119
|
+
"StepFinishEvent",
|
|
120
|
+
"ErrorEvent",
|
|
121
|
+
"DoneEvent",
|
|
93
122
|
# Enums
|
|
94
123
|
"HealthStatus",
|
|
95
124
|
"HealthStatusCode",
|
|
@@ -8,8 +8,10 @@ import grpc
|
|
|
8
8
|
import json
|
|
9
9
|
import os
|
|
10
10
|
from pathlib import Path
|
|
11
|
-
from typing import Optional, List, Dict, Iterator, Any, Union
|
|
11
|
+
from typing import Optional, List, Dict, Iterator, Any, Union, overload
|
|
12
12
|
|
|
13
|
+
from .provider import ModelRef
|
|
14
|
+
from .session import CodeSession
|
|
13
15
|
from .types import (
|
|
14
16
|
ProviderInfo,
|
|
15
17
|
ModelInfo,
|
|
@@ -234,6 +236,70 @@ class A3sClient:
|
|
|
234
236
|
"session": response.session,
|
|
235
237
|
}
|
|
236
238
|
|
|
239
|
+
async def session(
|
|
240
|
+
self,
|
|
241
|
+
model: ModelRef,
|
|
242
|
+
*,
|
|
243
|
+
workspace: str = "",
|
|
244
|
+
system: Optional[str] = None,
|
|
245
|
+
session_id: Optional[str] = None,
|
|
246
|
+
auto_compact: Optional[bool] = None,
|
|
247
|
+
auto_compact_threshold: Optional[float] = None,
|
|
248
|
+
) -> CodeSession:
|
|
249
|
+
"""Create a session with a model reference (high-level API).
|
|
250
|
+
|
|
251
|
+
Returns a CodeSession object with send(), stream(), delegate() methods.
|
|
252
|
+
Supports ``async with`` for automatic cleanup.
|
|
253
|
+
|
|
254
|
+
Args:
|
|
255
|
+
model: Model reference from create_provider()
|
|
256
|
+
workspace: Working directory for tool sandboxing
|
|
257
|
+
system: System prompt
|
|
258
|
+
session_id: Optional session ID (server generates one if omitted)
|
|
259
|
+
auto_compact: Enable automatic context compaction
|
|
260
|
+
auto_compact_threshold: Auto-compact threshold (0.0-1.0)
|
|
261
|
+
|
|
262
|
+
Returns:
|
|
263
|
+
CodeSession with high-level API
|
|
264
|
+
|
|
265
|
+
Example:
|
|
266
|
+
```python
|
|
267
|
+
from a3s_code import A3sClient, create_provider
|
|
268
|
+
|
|
269
|
+
anthropic = create_provider(name="anthropic", api_key="sk-ant-xxx")
|
|
270
|
+
|
|
271
|
+
async with A3sClient() as client:
|
|
272
|
+
async with await client.session(
|
|
273
|
+
model=anthropic("claude-sonnet-4-20250514"),
|
|
274
|
+
workspace="/project",
|
|
275
|
+
system="You are a senior engineer.",
|
|
276
|
+
) as session:
|
|
277
|
+
result = await session.send("Refactor the auth module")
|
|
278
|
+
print(result.text)
|
|
279
|
+
```
|
|
280
|
+
"""
|
|
281
|
+
from .types import LLMConfig, SessionConfig
|
|
282
|
+
|
|
283
|
+
llm = LLMConfig(
|
|
284
|
+
provider=model.provider,
|
|
285
|
+
model=model.model,
|
|
286
|
+
api_key=model.api_key,
|
|
287
|
+
base_url=model.base_url,
|
|
288
|
+
)
|
|
289
|
+
config = SessionConfig(
|
|
290
|
+
name=f"session-{session_id or 'auto'}",
|
|
291
|
+
workspace=workspace,
|
|
292
|
+
llm=llm,
|
|
293
|
+
system_prompt=system,
|
|
294
|
+
auto_compact=auto_compact,
|
|
295
|
+
)
|
|
296
|
+
request = {
|
|
297
|
+
"session_id": session_id,
|
|
298
|
+
"config": self._session_config_to_proto(config),
|
|
299
|
+
}
|
|
300
|
+
response = await self._stub.CreateSession(request)
|
|
301
|
+
return CodeSession(self, response.session_id)
|
|
302
|
+
|
|
237
303
|
async def destroy_session(self, session_id: str) -> bool:
|
|
238
304
|
"""Destroy a session."""
|
|
239
305
|
response = await self._stub.DestroySession({"session_id": session_id})
|
|
@@ -0,0 +1,63 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Provider Factory
|
|
3
|
+
|
|
4
|
+
Vercel AI SDK-style provider abstraction for A3S Code.
|
|
5
|
+
|
|
6
|
+
Example:
|
|
7
|
+
```python
|
|
8
|
+
from a3s_code import create_provider
|
|
9
|
+
|
|
10
|
+
openai = create_provider(name="openai", api_key="sk-xxx")
|
|
11
|
+
kimi = create_provider(name="kimi", api_key="sk-xxx", base_url="http://xxx/v1")
|
|
12
|
+
|
|
13
|
+
# Use as model selector
|
|
14
|
+
model = openai("gpt-4o")
|
|
15
|
+
model2 = kimi("k2.5")
|
|
16
|
+
```
|
|
17
|
+
"""
|
|
18
|
+
|
|
19
|
+
from dataclasses import dataclass
|
|
20
|
+
from typing import Optional, Callable
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
@dataclass
|
|
24
|
+
class ModelRef:
|
|
25
|
+
"""A resolved provider + model pair."""
|
|
26
|
+
provider: str
|
|
27
|
+
model: str
|
|
28
|
+
api_key: str
|
|
29
|
+
base_url: Optional[str] = None
|
|
30
|
+
|
|
31
|
+
|
|
32
|
+
ModelSelector = Callable[[str], ModelRef]
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def create_provider(
|
|
36
|
+
name: str,
|
|
37
|
+
api_key: str,
|
|
38
|
+
base_url: Optional[str] = None,
|
|
39
|
+
) -> ModelSelector:
|
|
40
|
+
"""Create a provider factory that returns model references.
|
|
41
|
+
|
|
42
|
+
Args:
|
|
43
|
+
name: Provider name (e.g., 'openai', 'anthropic', 'kimi')
|
|
44
|
+
api_key: API key
|
|
45
|
+
base_url: Base URL override
|
|
46
|
+
|
|
47
|
+
Returns:
|
|
48
|
+
A callable that takes a model ID and returns a ModelRef.
|
|
49
|
+
|
|
50
|
+
Example:
|
|
51
|
+
```python
|
|
52
|
+
openai = create_provider(name="openai", api_key="sk-xxx")
|
|
53
|
+
model = openai("gpt-4o")
|
|
54
|
+
```
|
|
55
|
+
"""
|
|
56
|
+
def selector(model_id: str) -> ModelRef:
|
|
57
|
+
return ModelRef(
|
|
58
|
+
provider=name,
|
|
59
|
+
model=model_id,
|
|
60
|
+
api_key=api_key,
|
|
61
|
+
base_url=base_url,
|
|
62
|
+
)
|
|
63
|
+
return selector
|
|
@@ -0,0 +1,476 @@
|
|
|
1
|
+
"""
|
|
2
|
+
Session — The core abstraction for A3S Code Python SDK.
|
|
3
|
+
|
|
4
|
+
A Session binds a workspace and model at creation time (immutable).
|
|
5
|
+
All agentic, generation, streaming, and context management calls are methods on the session.
|
|
6
|
+
|
|
7
|
+
Supports ``async with`` for automatic cleanup.
|
|
8
|
+
|
|
9
|
+
Example:
|
|
10
|
+
```python
|
|
11
|
+
from a3s_code import A3sClient, create_provider
|
|
12
|
+
|
|
13
|
+
anthropic = create_provider(name="anthropic", api_key="sk-ant-xxx")
|
|
14
|
+
|
|
15
|
+
async with A3sClient() as client:
|
|
16
|
+
async with await client.create_session(
|
|
17
|
+
model=anthropic("claude-sonnet-4-20250514"),
|
|
18
|
+
workspace="/project",
|
|
19
|
+
system="You are a senior engineer.",
|
|
20
|
+
) as session:
|
|
21
|
+
result = await session.send("Refactor the auth module")
|
|
22
|
+
print(result.text)
|
|
23
|
+
```
|
|
24
|
+
"""
|
|
25
|
+
|
|
26
|
+
from __future__ import annotations
|
|
27
|
+
|
|
28
|
+
from dataclasses import dataclass, field
|
|
29
|
+
from typing import (
|
|
30
|
+
Any,
|
|
31
|
+
AsyncIterator,
|
|
32
|
+
Dict,
|
|
33
|
+
List,
|
|
34
|
+
Literal,
|
|
35
|
+
Optional,
|
|
36
|
+
TYPE_CHECKING,
|
|
37
|
+
Union,
|
|
38
|
+
)
|
|
39
|
+
|
|
40
|
+
from .types import (
|
|
41
|
+
AgenticStrategy,
|
|
42
|
+
AgenticGenerateEvent,
|
|
43
|
+
Usage,
|
|
44
|
+
ToolCall,
|
|
45
|
+
ToolResult,
|
|
46
|
+
)
|
|
47
|
+
|
|
48
|
+
if TYPE_CHECKING:
|
|
49
|
+
from .client import A3sClient
|
|
50
|
+
|
|
51
|
+
|
|
52
|
+
# ============================================================================
|
|
53
|
+
# Types
|
|
54
|
+
# ============================================================================
|
|
55
|
+
|
|
56
|
+
@dataclass
|
|
57
|
+
class StepResult:
|
|
58
|
+
"""Step information from an agentic loop iteration."""
|
|
59
|
+
step_index: int
|
|
60
|
+
text: str
|
|
61
|
+
tool_calls: List[ToolCall] = field(default_factory=list)
|
|
62
|
+
tool_results: List[ToolResult] = field(default_factory=list)
|
|
63
|
+
usage: Optional[Usage] = None
|
|
64
|
+
finish_reason: Optional[str] = None
|
|
65
|
+
|
|
66
|
+
|
|
67
|
+
@dataclass
|
|
68
|
+
class SendResult:
|
|
69
|
+
"""Result from session.send()."""
|
|
70
|
+
text: str
|
|
71
|
+
steps: List[StepResult] = field(default_factory=list)
|
|
72
|
+
tool_calls: List[ToolCall] = field(default_factory=list)
|
|
73
|
+
usage: Optional[Usage] = None
|
|
74
|
+
finish_reason: str = "stop"
|
|
75
|
+
|
|
76
|
+
|
|
77
|
+
@dataclass
|
|
78
|
+
class SessionStats:
|
|
79
|
+
"""Session statistics."""
|
|
80
|
+
total_tokens: int = 0
|
|
81
|
+
prompt_tokens: int = 0
|
|
82
|
+
completion_tokens: int = 0
|
|
83
|
+
total_cost: float = 0.0
|
|
84
|
+
message_count: int = 0
|
|
85
|
+
tool_call_count: int = 0
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
# Agent loop event types
|
|
89
|
+
@dataclass
|
|
90
|
+
class TextEvent:
|
|
91
|
+
type: Literal["text"] = "text"
|
|
92
|
+
content: str = ""
|
|
93
|
+
|
|
94
|
+
@dataclass
|
|
95
|
+
class ToolCallEvent:
|
|
96
|
+
type: Literal["tool_call"] = "tool_call"
|
|
97
|
+
tool_name: str = ""
|
|
98
|
+
args: Dict[str, Any] = field(default_factory=dict)
|
|
99
|
+
tool_call_id: str = ""
|
|
100
|
+
|
|
101
|
+
@dataclass
|
|
102
|
+
class ToolResultEvent:
|
|
103
|
+
type: Literal["tool_result"] = "tool_result"
|
|
104
|
+
tool_call_id: str = ""
|
|
105
|
+
output: str = ""
|
|
106
|
+
success: bool = True
|
|
107
|
+
|
|
108
|
+
@dataclass
|
|
109
|
+
class StepFinishEvent:
|
|
110
|
+
type: Literal["step_finish"] = "step_finish"
|
|
111
|
+
step_index: int = 0
|
|
112
|
+
text: str = ""
|
|
113
|
+
|
|
114
|
+
@dataclass
|
|
115
|
+
class ErrorEvent:
|
|
116
|
+
type: Literal["error"] = "error"
|
|
117
|
+
message: str = ""
|
|
118
|
+
recoverable: bool = False
|
|
119
|
+
|
|
120
|
+
@dataclass
|
|
121
|
+
class DoneEvent:
|
|
122
|
+
type: Literal["done"] = "done"
|
|
123
|
+
finish_reason: str = "stop"
|
|
124
|
+
|
|
125
|
+
|
|
126
|
+
AgentLoopEvent = Union[
|
|
127
|
+
TextEvent, ToolCallEvent, ToolResultEvent,
|
|
128
|
+
StepFinishEvent, ErrorEvent, DoneEvent,
|
|
129
|
+
]
|
|
130
|
+
|
|
131
|
+
|
|
132
|
+
# ============================================================================
|
|
133
|
+
# Session Class
|
|
134
|
+
# ============================================================================
|
|
135
|
+
|
|
136
|
+
class CodeSession:
|
|
137
|
+
"""Session — The core object for interacting with A3S Code.
|
|
138
|
+
|
|
139
|
+
Created via ``client.create_session(model=...)``. Workspace and model are
|
|
140
|
+
immutable after creation. Supports ``async with`` for automatic cleanup.
|
|
141
|
+
"""
|
|
142
|
+
|
|
143
|
+
def __init__(self, client: "A3sClient", session_id: str) -> None:
|
|
144
|
+
self._client = client
|
|
145
|
+
self.id = session_id
|
|
146
|
+
self._closed = False
|
|
147
|
+
|
|
148
|
+
async def __aenter__(self) -> "CodeSession":
|
|
149
|
+
return self
|
|
150
|
+
|
|
151
|
+
async def __aexit__(self, exc_type: Any, exc_val: Any, exc_tb: Any) -> None:
|
|
152
|
+
await self.close()
|
|
153
|
+
|
|
154
|
+
# --------------------------------------------------------------------------
|
|
155
|
+
# AgenticLoop: send() / stream()
|
|
156
|
+
# --------------------------------------------------------------------------
|
|
157
|
+
|
|
158
|
+
async def send(
|
|
159
|
+
self,
|
|
160
|
+
prompt: str,
|
|
161
|
+
*,
|
|
162
|
+
max_steps: int = 50,
|
|
163
|
+
strategy: AgenticStrategy = AgenticStrategy.AUTO,
|
|
164
|
+
reflection: bool = True,
|
|
165
|
+
planning: bool = False,
|
|
166
|
+
) -> SendResult:
|
|
167
|
+
"""Send a message to the agent. Runs the server-side AgenticLoop.
|
|
168
|
+
|
|
169
|
+
Args:
|
|
170
|
+
prompt: The task prompt
|
|
171
|
+
max_steps: Maximum loop iterations (default 50)
|
|
172
|
+
strategy: Agentic strategy to use
|
|
173
|
+
reflection: Enable reflection after tool failures
|
|
174
|
+
planning: Enable planning before execution
|
|
175
|
+
|
|
176
|
+
Returns:
|
|
177
|
+
SendResult with text, steps, tool_calls, usage, finish_reason
|
|
178
|
+
|
|
179
|
+
Example:
|
|
180
|
+
```python
|
|
181
|
+
result = await session.send("Refactor the auth module")
|
|
182
|
+
print(result.text)
|
|
183
|
+
```
|
|
184
|
+
"""
|
|
185
|
+
self._ensure_open()
|
|
186
|
+
resp = await self._client.agentic_generate(
|
|
187
|
+
self.id, prompt,
|
|
188
|
+
strategy=strategy,
|
|
189
|
+
max_steps=max_steps,
|
|
190
|
+
reflection=reflection,
|
|
191
|
+
planning=planning,
|
|
192
|
+
)
|
|
193
|
+
steps = [
|
|
194
|
+
StepResult(
|
|
195
|
+
step_index=s.get("step_index", i),
|
|
196
|
+
text=s.get("text", ""),
|
|
197
|
+
tool_calls=s.get("tool_calls", []),
|
|
198
|
+
tool_results=s.get("tool_results", []),
|
|
199
|
+
usage=s.get("usage"),
|
|
200
|
+
finish_reason=s.get("finish_reason"),
|
|
201
|
+
)
|
|
202
|
+
for i, s in enumerate(resp.get("steps", []))
|
|
203
|
+
]
|
|
204
|
+
return SendResult(
|
|
205
|
+
text=resp.get("text", ""),
|
|
206
|
+
steps=steps,
|
|
207
|
+
tool_calls=resp.get("tool_calls", []),
|
|
208
|
+
usage=resp.get("usage"),
|
|
209
|
+
finish_reason=resp.get("finish_reason", "stop"),
|
|
210
|
+
)
|
|
211
|
+
|
|
212
|
+
async def stream(
|
|
213
|
+
self,
|
|
214
|
+
prompt: str,
|
|
215
|
+
*,
|
|
216
|
+
max_steps: int = 50,
|
|
217
|
+
strategy: AgenticStrategy = AgenticStrategy.AUTO,
|
|
218
|
+
reflection: bool = True,
|
|
219
|
+
planning: bool = False,
|
|
220
|
+
) -> AsyncIterator[AgentLoopEvent]:
|
|
221
|
+
"""Stream a message to the agent with real-time events.
|
|
222
|
+
|
|
223
|
+
Args:
|
|
224
|
+
prompt: The task prompt
|
|
225
|
+
max_steps: Maximum loop iterations
|
|
226
|
+
strategy: Agentic strategy
|
|
227
|
+
reflection: Enable reflection
|
|
228
|
+
planning: Enable planning
|
|
229
|
+
|
|
230
|
+
Yields:
|
|
231
|
+
AgentLoopEvent objects (TextEvent, ToolCallEvent, etc.)
|
|
232
|
+
|
|
233
|
+
Example:
|
|
234
|
+
```python
|
|
235
|
+
async for event in session.stream("Fix all TODOs"):
|
|
236
|
+
if isinstance(event, TextEvent):
|
|
237
|
+
print(event.content, end="", flush=True)
|
|
238
|
+
```
|
|
239
|
+
"""
|
|
240
|
+
self._ensure_open()
|
|
241
|
+
async for event in self._client.stream_agentic_generate(
|
|
242
|
+
self.id, prompt,
|
|
243
|
+
strategy=strategy,
|
|
244
|
+
max_steps=max_steps,
|
|
245
|
+
reflection=reflection,
|
|
246
|
+
planning=planning,
|
|
247
|
+
):
|
|
248
|
+
mapped = self._map_event(event)
|
|
249
|
+
if mapped is not None:
|
|
250
|
+
yield mapped
|
|
251
|
+
|
|
252
|
+
# --------------------------------------------------------------------------
|
|
253
|
+
# Delegation (Subagents)
|
|
254
|
+
# --------------------------------------------------------------------------
|
|
255
|
+
|
|
256
|
+
async def delegate(
|
|
257
|
+
self,
|
|
258
|
+
agent: str,
|
|
259
|
+
task: str,
|
|
260
|
+
*,
|
|
261
|
+
max_steps: int = 50,
|
|
262
|
+
) -> SendResult:
|
|
263
|
+
"""Delegate a task to a subagent.
|
|
264
|
+
|
|
265
|
+
Args:
|
|
266
|
+
agent: Agent type ('explore', 'plan', 'general')
|
|
267
|
+
task: The task prompt
|
|
268
|
+
max_steps: Maximum loop iterations
|
|
269
|
+
|
|
270
|
+
Example:
|
|
271
|
+
```python
|
|
272
|
+
result = await session.delegate("explore", "Find all API endpoints")
|
|
273
|
+
print(result.text)
|
|
274
|
+
```
|
|
275
|
+
"""
|
|
276
|
+
self._ensure_open()
|
|
277
|
+
resp = await self._client.delegate(self.id, agent, task)
|
|
278
|
+
return SendResult(
|
|
279
|
+
text=resp.get("text", ""),
|
|
280
|
+
tool_calls=resp.get("tool_calls", []),
|
|
281
|
+
usage=resp.get("usage"),
|
|
282
|
+
finish_reason=resp.get("finish_reason", "stop"),
|
|
283
|
+
)
|
|
284
|
+
|
|
285
|
+
async def delegate_stream(
|
|
286
|
+
self,
|
|
287
|
+
agent: str,
|
|
288
|
+
task: str,
|
|
289
|
+
*,
|
|
290
|
+
max_steps: int = 50,
|
|
291
|
+
) -> AsyncIterator[AgentLoopEvent]:
|
|
292
|
+
"""Delegate a task to a subagent with streaming.
|
|
293
|
+
|
|
294
|
+
Yields:
|
|
295
|
+
AgentLoopEvent objects
|
|
296
|
+
"""
|
|
297
|
+
self._ensure_open()
|
|
298
|
+
async for event in self._client.stream_delegate(self.id, agent, task):
|
|
299
|
+
mapped = self._map_event(event)
|
|
300
|
+
if mapped is not None:
|
|
301
|
+
yield mapped
|
|
302
|
+
|
|
303
|
+
# --------------------------------------------------------------------------
|
|
304
|
+
# Context Management
|
|
305
|
+
# --------------------------------------------------------------------------
|
|
306
|
+
|
|
307
|
+
async def get_context_usage(self) -> Dict[str, Any]:
|
|
308
|
+
"""Get context usage (token counts, message count)."""
|
|
309
|
+
self._ensure_open()
|
|
310
|
+
return await self._client.get_context_usage(self.id)
|
|
311
|
+
|
|
312
|
+
async def compact_context(self) -> None:
|
|
313
|
+
"""Compact the conversation context to save tokens."""
|
|
314
|
+
self._ensure_open()
|
|
315
|
+
await self._client.compact_context(self.id)
|
|
316
|
+
|
|
317
|
+
async def clear_context(self) -> None:
|
|
318
|
+
"""Clear conversation history."""
|
|
319
|
+
self._ensure_open()
|
|
320
|
+
await self._client.clear_context(self.id)
|
|
321
|
+
|
|
322
|
+
async def get_messages(
|
|
323
|
+
self, limit: Optional[int] = None, offset: Optional[int] = None
|
|
324
|
+
) -> Dict[str, Any]:
|
|
325
|
+
"""Get conversation messages."""
|
|
326
|
+
self._ensure_open()
|
|
327
|
+
return await self._client.get_messages(self.id, limit, offset)
|
|
328
|
+
|
|
329
|
+
# --------------------------------------------------------------------------
|
|
330
|
+
# Skills
|
|
331
|
+
# --------------------------------------------------------------------------
|
|
332
|
+
|
|
333
|
+
async def load_skill(self, name: str, content: Optional[str] = None) -> None:
|
|
334
|
+
"""Load a skill by name."""
|
|
335
|
+
self._ensure_open()
|
|
336
|
+
await self._client.load_skill(self.id, name, content or "")
|
|
337
|
+
|
|
338
|
+
async def load_skills(self, directory: str, recursive: bool = True) -> List[str]:
|
|
339
|
+
"""Load all skills from a directory."""
|
|
340
|
+
self._ensure_open()
|
|
341
|
+
resp = await self._client.load_skills_from_dir(self.id, directory, recursive)
|
|
342
|
+
return resp.get("loaded_skills", [])
|
|
343
|
+
|
|
344
|
+
async def unload_skill(self, name: str) -> None:
|
|
345
|
+
"""Unload a skill."""
|
|
346
|
+
self._ensure_open()
|
|
347
|
+
await self._client.unload_skill(self.id, name)
|
|
348
|
+
|
|
349
|
+
async def list_skills(self) -> List[Dict[str, Any]]:
|
|
350
|
+
"""List available skills."""
|
|
351
|
+
self._ensure_open()
|
|
352
|
+
return await self._client.list_skills(self.id)
|
|
353
|
+
|
|
354
|
+
# --------------------------------------------------------------------------
|
|
355
|
+
# HITL & Permissions
|
|
356
|
+
# --------------------------------------------------------------------------
|
|
357
|
+
|
|
358
|
+
async def set_confirmation(
|
|
359
|
+
self,
|
|
360
|
+
*,
|
|
361
|
+
require_confirmation: Optional[List[str]] = None,
|
|
362
|
+
auto_approve: Optional[List[str]] = None,
|
|
363
|
+
timeout: int = 30000,
|
|
364
|
+
timeout_action: str = "reject",
|
|
365
|
+
) -> None:
|
|
366
|
+
"""Set HITL confirmation policy."""
|
|
367
|
+
self._ensure_open()
|
|
368
|
+
from .types import ConfirmationPolicy
|
|
369
|
+
policy = ConfirmationPolicy(
|
|
370
|
+
enabled=True,
|
|
371
|
+
require_confirm_tools=require_confirmation or [],
|
|
372
|
+
auto_approve_tools=auto_approve or [],
|
|
373
|
+
default_timeout_ms=timeout,
|
|
374
|
+
timeout_action=timeout_action,
|
|
375
|
+
)
|
|
376
|
+
await self._client.set_confirmation_policy(self.id, policy)
|
|
377
|
+
|
|
378
|
+
async def set_permissions(
|
|
379
|
+
self,
|
|
380
|
+
*,
|
|
381
|
+
default_action: str = "allow",
|
|
382
|
+
allow: Optional[List[str]] = None,
|
|
383
|
+
deny: Optional[List[str]] = None,
|
|
384
|
+
ask: Optional[List[str]] = None,
|
|
385
|
+
) -> None:
|
|
386
|
+
"""Set tool permission policy."""
|
|
387
|
+
self._ensure_open()
|
|
388
|
+
from .types import PermissionPolicy, PermissionRule
|
|
389
|
+
policy = PermissionPolicy(
|
|
390
|
+
enabled=True,
|
|
391
|
+
allow=[PermissionRule(rule=r) for r in (allow or [])],
|
|
392
|
+
deny=[PermissionRule(rule=r) for r in (deny or [])],
|
|
393
|
+
ask=[PermissionRule(rule=r) for r in (ask or [])],
|
|
394
|
+
default_decision=default_action,
|
|
395
|
+
)
|
|
396
|
+
await self._client.set_permission_policy(self.id, policy)
|
|
397
|
+
|
|
398
|
+
async def confirm(self, confirmation_id: str, approved: bool, reason: str = "") -> None:
|
|
399
|
+
"""Respond to a confirmation request."""
|
|
400
|
+
self._ensure_open()
|
|
401
|
+
await self._client.confirm_tool_execution(self.id, confirmation_id, approved, reason)
|
|
402
|
+
|
|
403
|
+
# --------------------------------------------------------------------------
|
|
404
|
+
# Observability
|
|
405
|
+
# --------------------------------------------------------------------------
|
|
406
|
+
|
|
407
|
+
async def get_stats(self) -> SessionStats:
|
|
408
|
+
"""Get session statistics (tokens, cost, tool calls)."""
|
|
409
|
+
self._ensure_open()
|
|
410
|
+
cost = await self._client.get_cost_summary(self.id)
|
|
411
|
+
return SessionStats(
|
|
412
|
+
total_tokens=cost.get("total_tokens", 0),
|
|
413
|
+
prompt_tokens=cost.get("total_prompt_tokens", 0),
|
|
414
|
+
completion_tokens=cost.get("total_completion_tokens", 0),
|
|
415
|
+
total_cost=cost.get("total_cost_usd", 0.0),
|
|
416
|
+
message_count=cost.get("call_count", 0),
|
|
417
|
+
tool_call_count=cost.get("call_count", 0),
|
|
418
|
+
)
|
|
419
|
+
|
|
420
|
+
async def get_tool_metrics(self) -> Dict[str, Any]:
|
|
421
|
+
"""Get per-tool execution metrics."""
|
|
422
|
+
self._ensure_open()
|
|
423
|
+
return await self._client.get_tool_metrics(self.id)
|
|
424
|
+
|
|
425
|
+
# --------------------------------------------------------------------------
|
|
426
|
+
# Lifecycle
|
|
427
|
+
# --------------------------------------------------------------------------
|
|
428
|
+
|
|
429
|
+
async def close(self) -> None:
|
|
430
|
+
"""Close the session and release server resources."""
|
|
431
|
+
if self._closed:
|
|
432
|
+
return
|
|
433
|
+
self._closed = True
|
|
434
|
+
try:
|
|
435
|
+
await self._client.destroy_session(self.id)
|
|
436
|
+
except Exception:
|
|
437
|
+
pass
|
|
438
|
+
|
|
439
|
+
@property
|
|
440
|
+
def closed(self) -> bool:
|
|
441
|
+
"""Whether this session has been closed."""
|
|
442
|
+
return self._closed
|
|
443
|
+
|
|
444
|
+
def _ensure_open(self) -> None:
|
|
445
|
+
if self._closed:
|
|
446
|
+
raise RuntimeError(f"Session {self.id} has been closed")
|
|
447
|
+
|
|
448
|
+
@staticmethod
|
|
449
|
+
def _map_event(event: AgenticGenerateEvent) -> Optional[AgentLoopEvent]:
|
|
450
|
+
"""Map a proto AgenticGenerateEvent to an SDK AgentLoopEvent."""
|
|
451
|
+
etype = getattr(event, "type", None) or event.get("type", "")
|
|
452
|
+
if etype == "text":
|
|
453
|
+
content = getattr(event, "content", None) or event.get("text_delta", "") or event.get("content", "")
|
|
454
|
+
if content:
|
|
455
|
+
return TextEvent(content=content)
|
|
456
|
+
elif etype == "tool_call":
|
|
457
|
+
tc = getattr(event, "tool_call", None) or event.get("tool_call", {})
|
|
458
|
+
return ToolCallEvent(
|
|
459
|
+
tool_name=tc.get("name", "") if isinstance(tc, dict) else getattr(tc, "name", ""),
|
|
460
|
+
args={},
|
|
461
|
+
tool_call_id=tc.get("id", "") if isinstance(tc, dict) else getattr(tc, "id", ""),
|
|
462
|
+
)
|
|
463
|
+
elif etype == "tool_result":
|
|
464
|
+
tr = getattr(event, "tool_result", None) or event.get("tool_result", {})
|
|
465
|
+
return ToolResultEvent(
|
|
466
|
+
tool_call_id=event.get("tool_call_id", "") if isinstance(event, dict) else getattr(event, "tool_call_id", ""),
|
|
467
|
+
output=tr.get("output", "") if isinstance(tr, dict) else getattr(tr, "output", ""),
|
|
468
|
+
success=tr.get("success", True) if isinstance(tr, dict) else getattr(tr, "success", True),
|
|
469
|
+
)
|
|
470
|
+
elif etype == "error":
|
|
471
|
+
msg = getattr(event, "error_message", None) or event.get("error_message", "")
|
|
472
|
+
return ErrorEvent(message=msg)
|
|
473
|
+
elif etype == "done":
|
|
474
|
+
reason = getattr(event, "finish_reason", None) or event.get("finish_reason", "stop")
|
|
475
|
+
return DoneEvent(finish_reason=reason)
|
|
476
|
+
return None
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|
|
File without changes
|