antigravitystreaming 0.1.0__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.
@@ -0,0 +1,53 @@
1
+ Metadata-Version: 2.4
2
+ Name: antigravitystreaming
3
+ Version: 0.1.0
4
+ Summary: Stream text from the signed-in Google Antigravity CLI
5
+ Author: Eliran Wong
6
+ Requires-Python: >=3.10
7
+ Description-Content-Type: text/markdown
8
+
9
+ # antigravitystreaming
10
+
11
+ Stream text from the Google Antigravity CLI (`agy`). This library is a text-completion alternative to [agentmake](https://github.com/eliranwong/agentmake): you pass chat messages, and it yields response text as the CLI produces it.
12
+
13
+ Authentication is the signed-in session managed by `agy` (under `~/.gemini/antigravity-cli`). The child process does not receive `GEMINI_API_KEY` or `GOOGLE_API_KEY`. No API key is read or sent.
14
+
15
+ ## Install
16
+
17
+ ```bash
18
+ pip install antigravitystreaming
19
+ ```
20
+
21
+ Requires the `agy` executable on `PATH`, at `~/.local/bin/agy`, or at `~/.gemini/antigravity-cli/bin/agy`, and a completed `agy` sign-in.
22
+
23
+ ## Stream
24
+
25
+ ```python
26
+ import asyncio
27
+ from antigravitystreaming import stream_text
28
+
29
+ async def main():
30
+ messages = [
31
+ {"role": "system", "content": "You are a concise Bible study assistant."},
32
+ {"role": "user", "content": "What is John 3:16 about?"},
33
+ {"role": "assistant", "content": "It states God's love for the world."},
34
+ ]
35
+ async for chunk in stream_text(messages, follow_up="Quote it briefly."):
36
+ print(chunk, end="", flush=True)
37
+
38
+ asyncio.run(main())
39
+ ```
40
+
41
+ `messages` use `role` and `content`, the same shape agentmake accepts. `follow_up` is the newest user turn. An explicit `system` argument or system messages in `messages` are combined with the prompt under text-only guidance so instructions stay intact. Each call is a new headless process: the library does not resume a previous Antigravity session.
42
+
43
+ The process is text-only. It runs non-interactively with `--output-format stream-json`, `--dangerously-skip-permissions`, and `--disable-slash-commands`, in a temporary working directory outside the caller's project. Cancel an in-flight stream by setting the `asyncio.Event` passed as `cancel_event`.
44
+
45
+ `model` defaults to the CLI's configured model (e.g. `Gemini 3.8 Flash (High)`), so the library follows whatever `agy` is signed in and set up to run. Pass `model="gemini-3.8-flash-high"` or another model name to pin one. Pass `effort="high"` (or low, medium, xhigh, max) to configure reasoning effort.
46
+
47
+ `collect_text` returns the whole reply:
48
+
49
+ ```python
50
+ from antigravitystreaming import collect_text
51
+
52
+ text = await collect_text(follow_up="Reply with exactly: pong")
53
+ ```
@@ -0,0 +1,45 @@
1
+ # antigravitystreaming
2
+
3
+ Stream text from the Google Antigravity CLI (`agy`). This library is a text-completion alternative to [agentmake](https://github.com/eliranwong/agentmake): you pass chat messages, and it yields response text as the CLI produces it.
4
+
5
+ Authentication is the signed-in session managed by `agy` (under `~/.gemini/antigravity-cli`). The child process does not receive `GEMINI_API_KEY` or `GOOGLE_API_KEY`. No API key is read or sent.
6
+
7
+ ## Install
8
+
9
+ ```bash
10
+ pip install antigravitystreaming
11
+ ```
12
+
13
+ Requires the `agy` executable on `PATH`, at `~/.local/bin/agy`, or at `~/.gemini/antigravity-cli/bin/agy`, and a completed `agy` sign-in.
14
+
15
+ ## Stream
16
+
17
+ ```python
18
+ import asyncio
19
+ from antigravitystreaming import stream_text
20
+
21
+ async def main():
22
+ messages = [
23
+ {"role": "system", "content": "You are a concise Bible study assistant."},
24
+ {"role": "user", "content": "What is John 3:16 about?"},
25
+ {"role": "assistant", "content": "It states God's love for the world."},
26
+ ]
27
+ async for chunk in stream_text(messages, follow_up="Quote it briefly."):
28
+ print(chunk, end="", flush=True)
29
+
30
+ asyncio.run(main())
31
+ ```
32
+
33
+ `messages` use `role` and `content`, the same shape agentmake accepts. `follow_up` is the newest user turn. An explicit `system` argument or system messages in `messages` are combined with the prompt under text-only guidance so instructions stay intact. Each call is a new headless process: the library does not resume a previous Antigravity session.
34
+
35
+ The process is text-only. It runs non-interactively with `--output-format stream-json`, `--dangerously-skip-permissions`, and `--disable-slash-commands`, in a temporary working directory outside the caller's project. Cancel an in-flight stream by setting the `asyncio.Event` passed as `cancel_event`.
36
+
37
+ `model` defaults to the CLI's configured model (e.g. `Gemini 3.8 Flash (High)`), so the library follows whatever `agy` is signed in and set up to run. Pass `model="gemini-3.8-flash-high"` or another model name to pin one. Pass `effort="high"` (or low, medium, xhigh, max) to configure reasoning effort.
38
+
39
+ `collect_text` returns the whole reply:
40
+
41
+ ```python
42
+ from antigravitystreaming import collect_text
43
+
44
+ text = await collect_text(follow_up="Reply with exactly: pong")
45
+ ```
@@ -0,0 +1,41 @@
1
+ """Stream text from the signed-in Google Antigravity CLI.
2
+
3
+ The CLI session managed by Google Antigravity CLI (``agy``) is the credential.
4
+ This package does not read or send an API key.
5
+ """
6
+
7
+ from antigravitystreaming.stream import (
8
+ DEFAULT_MODEL,
9
+ AgyError,
10
+ AntigravityError,
11
+ build_command,
12
+ collect_text,
13
+ compose_prompt,
14
+ find_agy_bin,
15
+ find_antigravity_bin,
16
+ format_prompt,
17
+ parse_stream_line,
18
+ result_text_from_event,
19
+ stream_text,
20
+ subscription_env,
21
+ text_from_event,
22
+ )
23
+
24
+ __version__ = "0.1.0"
25
+
26
+ __all__ = [
27
+ "DEFAULT_MODEL",
28
+ "AgyError",
29
+ "AntigravityError",
30
+ "build_command",
31
+ "collect_text",
32
+ "compose_prompt",
33
+ "find_agy_bin",
34
+ "find_antigravity_bin",
35
+ "format_prompt",
36
+ "parse_stream_line",
37
+ "result_text_from_event",
38
+ "stream_text",
39
+ "subscription_env",
40
+ "text_from_event",
41
+ ]
@@ -0,0 +1,53 @@
1
+ Metadata-Version: 2.4
2
+ Name: antigravitystreaming
3
+ Version: 0.1.0
4
+ Summary: Stream text from the signed-in Google Antigravity CLI
5
+ Author: Eliran Wong
6
+ Requires-Python: >=3.10
7
+ Description-Content-Type: text/markdown
8
+
9
+ # antigravitystreaming
10
+
11
+ Stream text from the Google Antigravity CLI (`agy`). This library is a text-completion alternative to [agentmake](https://github.com/eliranwong/agentmake): you pass chat messages, and it yields response text as the CLI produces it.
12
+
13
+ Authentication is the signed-in session managed by `agy` (under `~/.gemini/antigravity-cli`). The child process does not receive `GEMINI_API_KEY` or `GOOGLE_API_KEY`. No API key is read or sent.
14
+
15
+ ## Install
16
+
17
+ ```bash
18
+ pip install antigravitystreaming
19
+ ```
20
+
21
+ Requires the `agy` executable on `PATH`, at `~/.local/bin/agy`, or at `~/.gemini/antigravity-cli/bin/agy`, and a completed `agy` sign-in.
22
+
23
+ ## Stream
24
+
25
+ ```python
26
+ import asyncio
27
+ from antigravitystreaming import stream_text
28
+
29
+ async def main():
30
+ messages = [
31
+ {"role": "system", "content": "You are a concise Bible study assistant."},
32
+ {"role": "user", "content": "What is John 3:16 about?"},
33
+ {"role": "assistant", "content": "It states God's love for the world."},
34
+ ]
35
+ async for chunk in stream_text(messages, follow_up="Quote it briefly."):
36
+ print(chunk, end="", flush=True)
37
+
38
+ asyncio.run(main())
39
+ ```
40
+
41
+ `messages` use `role` and `content`, the same shape agentmake accepts. `follow_up` is the newest user turn. An explicit `system` argument or system messages in `messages` are combined with the prompt under text-only guidance so instructions stay intact. Each call is a new headless process: the library does not resume a previous Antigravity session.
42
+
43
+ The process is text-only. It runs non-interactively with `--output-format stream-json`, `--dangerously-skip-permissions`, and `--disable-slash-commands`, in a temporary working directory outside the caller's project. Cancel an in-flight stream by setting the `asyncio.Event` passed as `cancel_event`.
44
+
45
+ `model` defaults to the CLI's configured model (e.g. `Gemini 3.8 Flash (High)`), so the library follows whatever `agy` is signed in and set up to run. Pass `model="gemini-3.8-flash-high"` or another model name to pin one. Pass `effort="high"` (or low, medium, xhigh, max) to configure reasoning effort.
46
+
47
+ `collect_text` returns the whole reply:
48
+
49
+ ```python
50
+ from antigravitystreaming import collect_text
51
+
52
+ text = await collect_text(follow_up="Reply with exactly: pong")
53
+ ```
@@ -0,0 +1,11 @@
1
+ README.md
2
+ __init__.py
3
+ pyproject.toml
4
+ stream.py
5
+ ./__init__.py
6
+ ./stream.py
7
+ antigravitystreaming.egg-info/PKG-INFO
8
+ antigravitystreaming.egg-info/SOURCES.txt
9
+ antigravitystreaming.egg-info/dependency_links.txt
10
+ antigravitystreaming.egg-info/top_level.txt
11
+ tests/test_stream.py
@@ -0,0 +1 @@
1
+ antigravitystreaming
@@ -0,0 +1,17 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "antigravitystreaming"
7
+ version = "0.1.0"
8
+ description = "Stream text from the signed-in Google Antigravity CLI"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ authors = [{ name = "Eliran Wong" }]
12
+
13
+ [tool.setuptools]
14
+ packages = ["antigravitystreaming"]
15
+
16
+ [tool.setuptools.package-dir]
17
+ antigravitystreaming = "."
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1,438 @@
1
+ """Headless Google Antigravity CLI text streaming.
2
+
3
+ Each call starts one ``agy`` process and yields text delta chunks from
4
+ ``--output-format stream-json``. The process uses the signed-in
5
+ subscription or personal account session. ``GEMINI_API_KEY`` and
6
+ ``GOOGLE_API_KEY`` are removed from the child environment so an exported
7
+ API key is not the credential.
8
+ """
9
+
10
+ from __future__ import annotations
11
+
12
+ import asyncio
13
+ import json
14
+ import os
15
+ import shutil
16
+ import signal
17
+ import tempfile
18
+ from collections.abc import Mapping, Sequence
19
+ from typing import Any
20
+
21
+ # None means the CLI's configured default model, so the library works with any
22
+ # sign-in. Pass model="gemini-3.8-flash-high" or another model name to pin one.
23
+ DEFAULT_MODEL: str | None = None
24
+ STREAM_LIMIT = 8 * 1024 * 1024
25
+
26
+ # API key variables that could bypass the signed-in session are removed.
27
+ _API_KEY_ENV = (
28
+ "GEMINI_API_KEY",
29
+ "GOOGLE_API_KEY",
30
+ "GOOGLE_GENAI_API_KEY",
31
+ "GOOGLE_AI_API_KEY",
32
+ )
33
+
34
+ # Parent agent session variables to prevent child session collisions.
35
+ _AGENT_SESSION_ENV = (
36
+ "ANTIGRAVITY_AGENT",
37
+ "ANTIGRAVITY_CONVERSATION_ID",
38
+ "ANTIGRAVITY_TRAJECTORY_ID",
39
+ "ANTIGRAVITY_SOURCE_METADATA",
40
+ )
41
+
42
+ # The system prompt owns the task. The user message is the input to that task.
43
+ _TEXT_ONLY_RULE = (
44
+ "Do not call tools, read or edit files, browse the web, or run commands. "
45
+ "Obey the system prompt. The user message is the input to that task. "
46
+ "When the system prompt asks for a specification or a formatted block instead of a direct answer, produce only that specification."
47
+ )
48
+
49
+
50
+ class AntigravityError(RuntimeError):
51
+ """The Google Antigravity CLI is missing, or a turn failed before any text."""
52
+
53
+
54
+ AgyError = AntigravityError
55
+
56
+
57
+ def find_agy_bin() -> str:
58
+ """Return the ``agy`` executable from ``PATH`` or a known install path."""
59
+ found = shutil.which("agy")
60
+ if found and os.path.isfile(found) and os.access(found, os.X_OK):
61
+ return found
62
+ for location in (
63
+ "~/.local/bin/agy",
64
+ "~/.gemini/antigravity-cli/bin/agy",
65
+ "~/.antigravity/bin/agy",
66
+ "~/.antigravity-ide/antigravity-ide/bin/agy",
67
+ "/usr/local/bin/agy",
68
+ ):
69
+ fallback = os.path.expanduser(location)
70
+ if os.path.isfile(fallback) and os.access(fallback, os.X_OK):
71
+ return fallback
72
+ raise AntigravityError(
73
+ "The agy CLI was not found. Install Google Antigravity CLI, then sign in with: agy"
74
+ )
75
+
76
+
77
+ find_antigravity_bin = find_agy_bin
78
+
79
+
80
+ def subscription_env(base: Mapping[str, str] | None = None) -> dict[str, str]:
81
+ """Environment for a child process that keeps the ``agy`` login session."""
82
+ env = dict(os.environ if base is None else base)
83
+ for name in _API_KEY_ENV:
84
+ env.pop(name, None)
85
+ for name in _AGENT_SESSION_ENV:
86
+ env.pop(name, None)
87
+ env["AGY_CLI_DISABLE_AUTO_UPDATE"] = "1"
88
+ env["DISABLE_AUTOUPDATER"] = "1"
89
+ return env
90
+
91
+
92
+ def _content_to_text(content: Any) -> str:
93
+ if content is None:
94
+ return ""
95
+ if isinstance(content, str):
96
+ return content.strip()
97
+ if isinstance(content, list):
98
+ parts: list[str] = []
99
+ for item in content:
100
+ if isinstance(item, str):
101
+ parts.append(item)
102
+ elif isinstance(item, dict):
103
+ text = item.get("text")
104
+ if not isinstance(text, str):
105
+ text = item.get("content")
106
+ if isinstance(text, str):
107
+ parts.append(text)
108
+ return "\n".join(part for part in parts if part).strip()
109
+ return str(content).strip()
110
+
111
+
112
+ def format_prompt(
113
+ messages: Sequence[Mapping[str, Any]] | None = None,
114
+ follow_up: str | None = None,
115
+ system: str | None = None,
116
+ ) -> tuple[str, str]:
117
+ """Split chat messages into a user prompt and a system prompt.
118
+
119
+ A lone user turn is left unlabeled so the model sees that text exactly.
120
+ Earlier turns are labeled when history is present. The explicit ``system``
121
+ argument is the system prompt. System messages already in ``messages`` stay
122
+ in the user prompt, so a generic identity does not outrank the task.
123
+ """
124
+ system_parts: list[str] = []
125
+ if isinstance(system, str) and system.strip():
126
+ system_parts.append(system.strip())
127
+ standing: list[str] = []
128
+ turns: list[tuple[str, str]] = []
129
+ for message in messages or ():
130
+ role = str(message.get("role") or "user").strip().lower()
131
+ content = _content_to_text(message.get("content"))
132
+ if not content:
133
+ continue
134
+ if role == "system":
135
+ standing.append(content)
136
+ elif role == "assistant":
137
+ turns.append(("Assistant", content))
138
+ else:
139
+ turns.append(("User", content))
140
+ if isinstance(follow_up, str) and follow_up.strip():
141
+ turns.append(("User", follow_up.strip()))
142
+ system_parts.append(_TEXT_ONLY_RULE)
143
+ if not turns and not standing:
144
+ raise AntigravityError("No prompt to send.")
145
+ if len(turns) == 1 and turns[0][0] == "User":
146
+ prompt = turns[0][1]
147
+ else:
148
+ prompt = "\n\n".join(f"{role}:\n{content}" for role, content in turns)
149
+ if standing:
150
+ context = "\n\n".join(standing)
151
+ prompt = f"{context}\n\n{prompt}" if prompt else context
152
+ return prompt, "\n\n".join(system_parts)
153
+
154
+
155
+ def compose_prompt(prompt: str, system_prompt: str) -> str:
156
+ """Combine user prompt and system prompt for the agy CLI."""
157
+ if system_prompt and system_prompt.strip():
158
+ return f"{system_prompt.strip()}\n\n{prompt}"
159
+ return prompt
160
+
161
+
162
+ def build_command(
163
+ agy_bin: str,
164
+ prompt: str,
165
+ model: str | None = None,
166
+ effort: str | None = None,
167
+ ) -> list[str]:
168
+ """Argv for one text-only headless turn. No API key and no tool approval."""
169
+ command = [
170
+ agy_bin,
171
+ "--output-format", "stream-json",
172
+ "--dangerously-skip-permissions",
173
+ "--disable-slash-commands",
174
+ ]
175
+ chosen_model = model.strip() if isinstance(model, str) and model.strip() else DEFAULT_MODEL
176
+ if chosen_model:
177
+ command.extend(("--model", chosen_model))
178
+ if isinstance(effort, str) and effort.strip():
179
+ command.extend(("--effort", effort.strip()))
180
+ command.extend(("--print", prompt))
181
+ return command
182
+
183
+
184
+ def parse_stream_line(line: str) -> dict[str, Any] | None:
185
+ """Parse one stdout line from ``--output-format stream-json``."""
186
+ text = line.strip()
187
+ if not text:
188
+ return None
189
+ try:
190
+ event = json.loads(text)
191
+ except json.JSONDecodeError:
192
+ return None
193
+ if isinstance(event, dict):
194
+ return event
195
+ return None
196
+
197
+
198
+ def text_from_event(event: Mapping[str, Any]) -> str | None:
199
+ """Return a streamed text chunk. Other event types contribute no answer text."""
200
+ if event.get("event") != "step_update":
201
+ return None
202
+ step_update = event.get("step_update")
203
+ if not isinstance(step_update, dict):
204
+ return None
205
+ if step_update.get("step_type") != "agent_response":
206
+ return None
207
+ delta = step_update.get("text_delta")
208
+ if isinstance(delta, str) and delta:
209
+ return delta
210
+ return None
211
+
212
+
213
+ def result_text_from_event(event: Mapping[str, Any]) -> str | None:
214
+ """Return text from a whole ``result`` event.
215
+
216
+ This is the fallback if no text delta was streamed.
217
+ """
218
+ if event.get("event") != "result":
219
+ return None
220
+ result = event.get("result")
221
+ if not isinstance(result, dict) or result.get("status") == "ERROR":
222
+ return None
223
+ response = result.get("response")
224
+ if isinstance(response, str) and response:
225
+ return response
226
+ return None
227
+
228
+
229
+ def _completion_workspace() -> str:
230
+ """Stable directory outside the caller's project, with no git root of its own."""
231
+ path = os.path.join(tempfile.gettempdir(), "antigravitystreaming-workspace")
232
+ os.makedirs(path, exist_ok=True)
233
+ return path
234
+
235
+
236
+ async def _cancel_task(task: asyncio.Task | None) -> None:
237
+ if task is None or task.done():
238
+ return
239
+ task.cancel()
240
+ try:
241
+ await task
242
+ except (asyncio.CancelledError, Exception):
243
+ return
244
+
245
+
246
+ async def _next_line(stream: asyncio.StreamReader, cancel_event: asyncio.Event | None) -> bytes | None:
247
+ """Read one stdout line. ``None`` means the caller cancelled."""
248
+ read_task = asyncio.create_task(stream.readline())
249
+ if cancel_event is None:
250
+ return await read_task
251
+ if cancel_event.is_set():
252
+ await _cancel_task(read_task)
253
+ return None
254
+ cancel_task = asyncio.create_task(cancel_event.wait())
255
+ try:
256
+ done, _pending = await asyncio.wait(
257
+ {read_task, cancel_task},
258
+ return_when=asyncio.FIRST_COMPLETED,
259
+ )
260
+ except asyncio.CancelledError:
261
+ await _cancel_task(read_task)
262
+ await _cancel_task(cancel_task)
263
+ raise
264
+ if read_task in done and not read_task.cancelled():
265
+ await _cancel_task(cancel_task)
266
+ return read_task.result()
267
+ await _cancel_task(read_task)
268
+ await _cancel_task(cancel_task)
269
+ return None
270
+
271
+
272
+ async def _read_stderr(stream: asyncio.StreamReader | None) -> str:
273
+ if stream is None:
274
+ return ""
275
+ chunks: list[str] = []
276
+ size = 0
277
+ try:
278
+ while True:
279
+ line = await stream.readline()
280
+ if not line:
281
+ break
282
+ text = line.decode("utf-8", errors="replace")
283
+ chunks.append(text)
284
+ size += len(text)
285
+ if size > 16000:
286
+ chunks = chunks[-20:]
287
+ size = sum(len(item) for item in chunks)
288
+ except asyncio.CancelledError:
289
+ return "".join(chunks)[-8000:]
290
+ return "".join(chunks)[-8000:]
291
+
292
+
293
+ async def _stop_process(process: asyncio.subprocess.Process | None) -> None:
294
+ if process is None or process.returncode is not None:
295
+ return
296
+ if os.name == "nt":
297
+ process.terminate()
298
+ else:
299
+ try:
300
+ os.killpg(process.pid, signal.SIGTERM)
301
+ except ProcessLookupError:
302
+ return
303
+ except Exception:
304
+ process.terminate()
305
+ try:
306
+ await asyncio.wait_for(process.wait(), timeout=3)
307
+ except asyncio.TimeoutError:
308
+ if os.name == "nt":
309
+ process.kill()
310
+ else:
311
+ try:
312
+ os.killpg(process.pid, signal.SIGKILL)
313
+ except ProcessLookupError:
314
+ return
315
+ except Exception:
316
+ process.kill()
317
+ try:
318
+ await process.wait()
319
+ except Exception:
320
+ return
321
+
322
+
323
+ async def stream_text(
324
+ messages: Sequence[Mapping[str, Any]] | None = None,
325
+ follow_up: str | None = None,
326
+ system: str | None = None,
327
+ model: str | None = None,
328
+ effort: str | None = None,
329
+ agy_bin: str | None = None,
330
+ antigravity_bin: str | None = None,
331
+ cancel_event: asyncio.Event | None = None,
332
+ ):
333
+ """Yield text chunks from one signed-in Google Antigravity CLI turn.
334
+
335
+ Closing the generator, or setting ``cancel_event``, stops the CLI process.
336
+ ``model`` defaults to :data:`DEFAULT_MODEL` (the CLI's configured model).
337
+ The working directory is a temporary folder, and this call does not resume
338
+ an earlier session.
339
+ """
340
+ prompt, system_prompt = format_prompt(messages, follow_up=follow_up, system=system)
341
+ full_prompt = compose_prompt(prompt, system_prompt)
342
+ executable = agy_bin or antigravity_bin or find_agy_bin()
343
+ work_dir = _completion_workspace()
344
+ process: asyncio.subprocess.Process | None = None
345
+ stderr_task: asyncio.Task | None = None
346
+ produced_text = False
347
+ try:
348
+ command = build_command(executable, full_prompt, model=model, effort=effort)
349
+ process = await asyncio.create_subprocess_exec(
350
+ *command,
351
+ stdin=asyncio.subprocess.DEVNULL,
352
+ stdout=asyncio.subprocess.PIPE,
353
+ stderr=asyncio.subprocess.PIPE,
354
+ cwd=work_dir,
355
+ env=subscription_env(),
356
+ start_new_session=True,
357
+ limit=STREAM_LIMIT,
358
+ )
359
+ if process.stdout is None:
360
+ raise AntigravityError("agy did not provide an output stream")
361
+ stderr_task = asyncio.create_task(_read_stderr(process.stderr))
362
+ while True:
363
+ if cancel_event is not None and cancel_event.is_set():
364
+ return
365
+ line = await _next_line(process.stdout, cancel_event)
366
+ if line is None:
367
+ return
368
+ if line == b"":
369
+ break
370
+ event = parse_stream_line(line.decode("utf-8", errors="replace"))
371
+ if not event:
372
+ continue
373
+ if event.get("event") == "result":
374
+ if cancel_event is not None and cancel_event.is_set():
375
+ return
376
+ res = event.get("result")
377
+ if isinstance(res, dict):
378
+ if res.get("status") == "ERROR" or res.get("error"):
379
+ message = res.get("error") or "Google Antigravity CLI reported an error."
380
+ raise AntigravityError(str(message))
381
+ if not produced_text:
382
+ text = res.get("response")
383
+ if isinstance(text, str) and text:
384
+ produced_text = True
385
+ yield text
386
+ continue
387
+ chunk = text_from_event(event)
388
+ if chunk:
389
+ produced_text = True
390
+ yield chunk
391
+ continue
392
+ if not produced_text:
393
+ chunk = result_text_from_event(event)
394
+ if chunk:
395
+ produced_text = True
396
+ yield chunk
397
+ returncode = await process.wait()
398
+ stderr_text = ""
399
+ if stderr_task is not None:
400
+ stderr_text = (await stderr_task).strip()
401
+ stderr_task = None
402
+ if returncode not in (0, None) and not produced_text:
403
+ detail = stderr_text or f"agy exited with code {returncode}"
404
+ raise AntigravityError(detail)
405
+ finally:
406
+ await _stop_process(process)
407
+ if stderr_task is not None and not stderr_task.done():
408
+ stderr_task.cancel()
409
+ try:
410
+ await stderr_task
411
+ except (asyncio.CancelledError, Exception):
412
+ pass
413
+
414
+
415
+ async def collect_text(
416
+ messages: Sequence[Mapping[str, Any]] | None = None,
417
+ follow_up: str | None = None,
418
+ system: str | None = None,
419
+ model: str | None = None,
420
+ effort: str | None = None,
421
+ agy_bin: str | None = None,
422
+ antigravity_bin: str | None = None,
423
+ cancel_event: asyncio.Event | None = None,
424
+ ) -> str:
425
+ """Run :func:`stream_text` and return the joined text."""
426
+ chunks: list[str] = []
427
+ async for chunk in stream_text(
428
+ messages,
429
+ follow_up=follow_up,
430
+ system=system,
431
+ model=model,
432
+ effort=effort,
433
+ agy_bin=agy_bin,
434
+ antigravity_bin=antigravity_bin,
435
+ cancel_event=cancel_event,
436
+ ):
437
+ chunks.append(chunk)
438
+ return "".join(chunks)
@@ -0,0 +1,175 @@
1
+ import unittest
2
+ from unittest import mock
3
+
4
+ from antigravitystreaming.stream import (
5
+ DEFAULT_MODEL,
6
+ AgyError,
7
+ AntigravityError,
8
+ build_command,
9
+ compose_prompt,
10
+ find_agy_bin,
11
+ format_prompt,
12
+ parse_stream_line,
13
+ result_text_from_event,
14
+ subscription_env,
15
+ text_from_event,
16
+ )
17
+
18
+
19
+ class FormatPromptTests(unittest.TestCase):
20
+ def test_single_user_turn_is_unlabeled(self):
21
+ prompt, system = format_prompt(
22
+ [{"role": "system", "content": "You are BibleMate AI."}],
23
+ follow_up="List two tools.",
24
+ )
25
+ self.assertIn("You are BibleMate AI.", prompt)
26
+ self.assertTrue(prompt.endswith("List two tools."))
27
+ self.assertNotIn("You are BibleMate AI.", system)
28
+ self.assertNotIn("List two tools.", system)
29
+ self.assertIn("Obey the system prompt.", system)
30
+ self.assertNotIn("User:", prompt)
31
+
32
+ def test_history_keeps_the_latest_request_intact(self):
33
+ prompt, system = format_prompt(
34
+ [
35
+ {"role": "system", "content": "Role."},
36
+ {"role": "user", "content": "First question"},
37
+ {"role": "assistant", "content": "First answer"},
38
+ ],
39
+ follow_up="Second question",
40
+ system="Extra role",
41
+ )
42
+ self.assertIn("Role.", prompt)
43
+ self.assertIn("User:\nFirst question", prompt)
44
+ self.assertIn("Assistant:\nFirst answer", prompt)
45
+ self.assertTrue(prompt.endswith("User:\nSecond question"))
46
+ self.assertIn("Extra role", system)
47
+ self.assertNotIn("Role.", system)
48
+ self.assertNotIn("Second question", system)
49
+
50
+ def test_content_blocks_join_text(self):
51
+ prompt, _system = format_prompt(
52
+ [{"role": "user", "content": [{"type": "text", "text": "Hello"}, {"text": "there"}]}],
53
+ )
54
+ self.assertEqual(prompt, "Hello\nthere")
55
+
56
+ def test_empty_prompt_raises(self):
57
+ with self.assertRaises(AntigravityError):
58
+ format_prompt([])
59
+
60
+ def test_compose_prompt(self):
61
+ composed = compose_prompt("User query", "System instruction")
62
+ self.assertEqual(composed, "System instruction\n\nUser query")
63
+ self.assertEqual(compose_prompt("User query", ""), "User query")
64
+
65
+
66
+ class EventTests(unittest.TestCase):
67
+ def test_text_delta_chunk(self):
68
+ event = parse_stream_line(
69
+ '{"event":"step_update","step_update":{"conversation_id":"c1",'
70
+ '"step_index":1,"state":"ACTIVE","step_type":"agent_response","text_delta":"pong"}}'
71
+ )
72
+ self.assertEqual(text_from_event(event), "pong")
73
+
74
+ def test_user_input_step_adds_no_answer_text(self):
75
+ event = parse_stream_line(
76
+ '{"event":"step_update","step_update":{"conversation_id":"c1",'
77
+ '"step_index":0,"state":"DONE","step_type":"user_input"}}'
78
+ )
79
+ self.assertIsNone(text_from_event(event))
80
+
81
+ def test_tool_event_adds_no_answer_text(self):
82
+ event = parse_stream_line(
83
+ '{"event":"step_update","step_update":{"conversation_id":"c1",'
84
+ '"step_index":2,"state":"ACTIVE","step_type":"tool","tool_name":"view_file"}}'
85
+ )
86
+ self.assertIsNone(text_from_event(event))
87
+
88
+ def test_result_event_fallback(self):
89
+ event = parse_stream_line(
90
+ '{"event":"result","result":{"conversation_id":"c1","status":"SUCCESS","response":"pong\\n"}}'
91
+ )
92
+ self.assertEqual(result_text_from_event(event), "pong\n")
93
+
94
+ def test_result_error_adds_no_answer_text(self):
95
+ event = parse_stream_line(
96
+ '{"event":"result","result":{"conversation_id":"c1","status":"ERROR","error":"bad model"}}'
97
+ )
98
+ self.assertIsNone(result_text_from_event(event))
99
+
100
+ def test_non_json_is_ignored(self):
101
+ self.assertIsNone(parse_stream_line("plain log line"))
102
+
103
+
104
+ class CommandTests(unittest.TestCase):
105
+ def test_command_is_text_only_and_has_no_api_key(self):
106
+ command = build_command("/usr/bin/agy", "Be brief.")
107
+ joined = " ".join(command)
108
+ self.assertNotIn("GEMINI_API_KEY", joined)
109
+ self.assertNotIn("GOOGLE_API_KEY", joined)
110
+ self.assertIn("--output-format", command)
111
+ self.assertIn("stream-json", command)
112
+ self.assertIn("--dangerously-skip-permissions", command)
113
+ self.assertIn("--disable-slash-commands", command)
114
+ self.assertIn("--print", command)
115
+ self.assertIn("Be brief.", command)
116
+ self.assertNotIn("--model", command)
117
+
118
+ def test_explicit_model_is_pinned(self):
119
+ command = build_command("/usr/bin/agy", "Be brief.", model="gemini-3.8-flash-high")
120
+ self.assertIn("--model", command)
121
+ self.assertIn("gemini-3.8-flash-high", command)
122
+
123
+ def test_explicit_effort_is_pinned(self):
124
+ command = build_command("/usr/bin/agy", "Be brief.", effort="high")
125
+ self.assertIn("--effort", command)
126
+ self.assertIn("high", command)
127
+
128
+ def test_default_model_is_none(self):
129
+ self.assertIsNone(DEFAULT_MODEL)
130
+
131
+ def test_agy_error_alias(self):
132
+ self.assertIs(AgyError, AntigravityError)
133
+
134
+ def test_subscription_env_drops_api_keys_and_agent_vars(self):
135
+ env = subscription_env(
136
+ {
137
+ "GEMINI_API_KEY": "sk-secret",
138
+ "GOOGLE_API_KEY": "sk-secret-2",
139
+ "ANTIGRAVITY_AGENT": "1",
140
+ "ANTIGRAVITY_CONVERSATION_ID": "conv-1",
141
+ "PATH": "/usr/bin",
142
+ "HOME": "/Users/me",
143
+ }
144
+ )
145
+ self.assertNotIn("GEMINI_API_KEY", env)
146
+ self.assertNotIn("GOOGLE_API_KEY", env)
147
+ self.assertNotIn("ANTIGRAVITY_AGENT", env)
148
+ self.assertNotIn("ANTIGRAVITY_CONVERSATION_ID", env)
149
+ self.assertEqual(env["PATH"], "/usr/bin")
150
+ self.assertEqual(env["AGY_CLI_DISABLE_AUTO_UPDATE"], "1")
151
+ self.assertEqual(env["DISABLE_AUTOUPDATER"], "1")
152
+
153
+ def test_find_agy_bin_uses_path(self):
154
+ with mock.patch("antigravitystreaming.stream.shutil.which", return_value="/opt/agy"):
155
+ with mock.patch("antigravitystreaming.stream.os.path.isfile", return_value=True):
156
+ with mock.patch("antigravitystreaming.stream.os.access", return_value=True):
157
+ self.assertEqual(find_agy_bin(), "/opt/agy")
158
+
159
+ def test_find_agy_bin_uses_known_install_fallback(self):
160
+ with mock.patch("antigravitystreaming.stream.shutil.which", return_value=None):
161
+ with mock.patch("antigravitystreaming.stream.os.path.expanduser", side_effect=lambda p: p):
162
+ with mock.patch("antigravitystreaming.stream.os.path.isfile", return_value=True):
163
+ with mock.patch("antigravitystreaming.stream.os.access", return_value=True):
164
+ self.assertEqual(find_agy_bin(), "~/.local/bin/agy")
165
+
166
+ def test_find_agy_bin_errors_when_missing(self):
167
+ with mock.patch("antigravitystreaming.stream.shutil.which", return_value=None):
168
+ with mock.patch("antigravitystreaming.stream.os.path.expanduser", side_effect=lambda p: p):
169
+ with mock.patch("antigravitystreaming.stream.os.path.isfile", return_value=False):
170
+ with self.assertRaises(AntigravityError):
171
+ find_agy_bin()
172
+
173
+
174
+ if __name__ == "__main__":
175
+ unittest.main()