agentpoke 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,38 @@
1
+ name: CI
2
+ on:
3
+ push:
4
+ branches: [main]
5
+ pull_request:
6
+ branches: [main]
7
+
8
+ jobs:
9
+ test:
10
+ runs-on: ${{ matrix.os }}
11
+ strategy:
12
+ fail-fast: false
13
+ matrix:
14
+ os: [ubuntu-latest, windows-latest]
15
+ python-version: ["3.10", "3.11", "3.12", "3.13"]
16
+ steps:
17
+ - uses: actions/checkout@v6
18
+ - uses: actions/setup-python@v6
19
+ with:
20
+ python-version: ${{ matrix.python-version }}
21
+ cache: pip
22
+ - run: python -m pip install -U pip
23
+ - run: python -m pip install -e ".[dev]"
24
+ - run: python -m ruff check .
25
+ - run: python -m pytest tests/ -v
26
+ - run: python -m compileall -q agentprobe tests
27
+
28
+ package:
29
+ runs-on: ubuntu-latest
30
+ steps:
31
+ - uses: actions/checkout@v6
32
+ - uses: actions/setup-python@v6
33
+ with:
34
+ python-version: "3.13"
35
+ cache: pip
36
+ - run: python -m pip install -U pip build twine
37
+ - run: python -m build
38
+ - run: python -m twine check dist/*
@@ -0,0 +1,12 @@
1
+ __pycache__/
2
+ *.py[cod]
3
+ *.egg-info/
4
+ dist/
5
+ build/
6
+ # Note: .agentprobe/snapshots/ should be committed in user projects
7
+ # so CI can compare against baselines. We ignore it here because
8
+ # this is the AgentProbe source repo, not a user project.
9
+ .pytest_cache/
10
+ .ruff_cache/
11
+ *.egg
12
+ .venv/
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Yufeng He
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,395 @@
1
+ Metadata-Version: 2.4
2
+ Name: agentpoke
3
+ Version: 0.1.0
4
+ Summary: Drop-in pytest plugin for regression-testing AI agents — snapshot baselines, semantic comparison, mock LLMs
5
+ Project-URL: Homepage, https://github.com/he-yufeng/AgentProbe
6
+ Project-URL: Repository, https://github.com/he-yufeng/AgentProbe
7
+ Author-email: Yufeng He <40085740+he-yufeng@users.noreply.github.com>
8
+ License-Expression: MIT
9
+ License-File: LICENSE
10
+ Keywords: agents,ai,ci,llm,pytest,regression,snapshot,testing
11
+ Classifier: Development Status :: 4 - Beta
12
+ Classifier: Framework :: Pytest
13
+ Classifier: Intended Audience :: Developers
14
+ Classifier: License :: OSI Approved :: MIT License
15
+ Classifier: Programming Language :: Python :: 3
16
+ Classifier: Programming Language :: Python :: 3.10
17
+ Classifier: Programming Language :: Python :: 3.11
18
+ Classifier: Programming Language :: Python :: 3.12
19
+ Classifier: Topic :: Software Development :: Testing
20
+ Requires-Python: >=3.10
21
+ Requires-Dist: click>=8.0
22
+ Requires-Dist: pydantic>=2.0
23
+ Requires-Dist: pytest>=7.0
24
+ Provides-Extra: all
25
+ Requires-Dist: openai>=1.0; extra == 'all'
26
+ Requires-Dist: sentence-transformers>=2.0; extra == 'all'
27
+ Provides-Extra: dev
28
+ Requires-Dist: pytest>=7.0; extra == 'dev'
29
+ Requires-Dist: ruff>=0.4; extra == 'dev'
30
+ Provides-Extra: openai
31
+ Requires-Dist: openai>=1.0; extra == 'openai'
32
+ Provides-Extra: semantic
33
+ Requires-Dist: sentence-transformers>=2.0; extra == 'semantic'
34
+ Description-Content-Type: text/markdown
35
+
36
+ <div align="center">
37
+
38
+ <img src="docs/banner.png" alt="AgentProbe — regression-testing for AI agents" width="100%">
39
+
40
+ Capture your agent's outputs, store them as baselines, and catch regressions in CI — with one decorator.
41
+
42
+ [![License: MIT](https://img.shields.io/badge/License-MIT-green.svg)](LICENSE)
43
+ [![Python 3.10+](https://img.shields.io/badge/python-3.10+-blue.svg)](https://www.python.org/downloads/)
44
+ [![CI](https://github.com/he-yufeng/AgentProbe/actions/workflows/ci.yml/badge.svg)](https://github.com/he-yufeng/AgentProbe/actions)
45
+
46
+ **[English](README.md) · [中文](README_CN.md)** &nbsp;·&nbsp; [Quick Start](#quick-start) · [How It Works](#how-it-works) · [How It Compares](#how-it-compares)
47
+
48
+ </div>
49
+
50
+ ---
51
+
52
+ ## The Problem
53
+
54
+ You ship an AI agent. It works great. Two weeks later, you update a prompt, swap a model, or bump a dependency — and something breaks. But you don't notice until a user complains, because **there's no test that catches agent behavior regressions**.
55
+
56
+ Traditional unit tests don't work for agents. The outputs are non-deterministic. They're natural language, not exact values. You can't just `assertEqual`. And even if you could, you'd spend more time writing test fixtures than writing the agent itself.
57
+
58
+ **AgentProbe** fixes this. One decorator captures your agent's output and saves it as a baseline snapshot. On the next run, it compares the new output against the baseline — using exact match or semantic similarity. If something changed, the test fails. Run it in CI, and you'll catch regressions before they hit production.
59
+
60
+ ## How It Works
61
+
62
+ ![AgentProbe snapshot flow](docs/architecture.png)
63
+
64
+ ## Quick Start
65
+
66
+ ```bash
67
+ pip install agentpoke
68
+ ```
69
+
70
+ ### 1. Snapshot Testing
71
+
72
+ Capture agent outputs and compare them across runs:
73
+
74
+ ```python
75
+ from agentprobe import snapshot
76
+
77
+ @snapshot("summarize_article")
78
+ def test_summarize():
79
+ # Your agent code here
80
+ result = my_agent.summarize("The quick brown fox jumps over the lazy dog.")
81
+ return result
82
+ ```
83
+
84
+ First run: creates a baseline in `.agentprobe/snapshots/summarize_article.json`.
85
+ Next runs: compares the output against the baseline. Fails if they differ.
86
+
87
+ Async agents work the same way:
88
+
89
+ ```python
90
+ @snapshot("async_summarize")
91
+ async def test_async_summarize():
92
+ result = await my_agent.summarize_async("The quick brown fox jumps over the lazy dog.")
93
+ return result
94
+ ```
95
+
96
+ When the output carries non-deterministic fields like timestamps or request ids, list them in `redact` so they're masked before comparison and don't cause spurious mismatches:
97
+
98
+ ```python
99
+ @snapshot("summarize_article", redact=["timestamp", "request_id"])
100
+ def test_summarize():
101
+ return my_agent.summarize("...") # {"summary": "...", "timestamp": 1718...}
102
+ ```
103
+
104
+ The named keys are replaced with `"<redacted>"` at any depth before the snapshot is saved and compared. Real changes to other fields still fail the snapshot.
105
+
106
+ ### 2. Mock LLM
107
+
108
+ Test agent logic without hitting any API:
109
+
110
+ ```python
111
+ from agentprobe import MockLLM
112
+
113
+ def test_agent_with_mock():
114
+ mock = MockLLM(responses=[
115
+ "The document discusses three main topics.",
116
+ "Based on my analysis, the sentiment is positive."
117
+ ])
118
+
119
+ # Use mock.chat.completions.create as a drop-in for openai
120
+ result = mock.chat.completions.create(
121
+ messages=[{"role": "user", "content": "Summarize this doc"}]
122
+ )
123
+ assert "three main topics" in result.choices[0].message.content
124
+ assert mock.call_count == 1
125
+ ```
126
+
127
+ ### 3. Tool Call Assertions
128
+
129
+ Verify your agent calls the right tools:
130
+
131
+ ```python
132
+ from agentprobe import assert_no_tool_called, assert_tool_called, assert_tool_sequence
133
+
134
+ def test_agent_uses_search():
135
+ tool_calls = [
136
+ {"name": "web_search", "arguments": {"query": "latest news"}},
137
+ {"name": "summarize", "arguments": {"text": "..."}},
138
+ ]
139
+ assert_tool_called(tool_calls, "web_search", times=1)
140
+ assert_tool_called(tool_calls, "web_search", with_args={"query": "latest news"})
141
+ assert_tool_sequence(tool_calls, ["web_search", "summarize"])
142
+ assert_no_tool_called(tool_calls, "delete_file")
143
+ ```
144
+
145
+ For multi-step agents, `assert_tool_sequence(..., contiguous=True)` catches accidental planner reorderings where a tool must immediately follow another tool.
146
+
147
+ When the exact call count is non-deterministic, use `min_times`/`max_times` instead of `times` — for example, assert a flaky API was retried at most three times, or a search ran at least twice:
148
+
149
+ ```python
150
+ assert_tool_called(tool_calls, "api_call", max_times=3) # retried, but bounded
151
+ assert_tool_called(tool_calls, "web_search", min_times=2) # at least two searches
152
+ ```
153
+
154
+ For overall efficiency, `assert_max_tool_calls` bounds the whole run rather than one tool — an agent can avoid repeating any single call and still be wastefully chatty. Unlike `max_times`, the budget may be met with zero calls:
155
+
156
+ ```python
157
+ assert_max_tool_calls(tool_calls, 10) # solve it in at most 10 calls
158
+ assert_max_tool_calls(tool_calls, 3, tool_name="api_call") # at most 3 api_calls (zero is fine)
159
+ ```
160
+
161
+ `with_args` is a nested subset match and also accepts OpenAI-style JSON string arguments:
162
+
163
+ ```python
164
+ assert_tool_called(
165
+ tool_calls,
166
+ "write_file",
167
+ with_args={"metadata": {"mode": "safe"}},
168
+ )
169
+ ```
170
+
171
+ For safety checks, `assert_tool_not_called_with` is the negative counterpart: it allows the tool but fails if any call carried a forbidden argument subset — handy when a tool is fine to use except in a dangerous mode:
172
+
173
+ ```python
174
+ # the agent may run shell commands, but never with sudo
175
+ assert_tool_not_called_with(tool_calls, "run", {"sudo": True})
176
+ # and may delete files, but never the filesystem root
177
+ assert_tool_not_called_with(tool_calls, "delete_file", {"path": "/"})
178
+ ```
179
+
180
+ ### 4. Schema Validation
181
+
182
+ Assert that agent outputs conform to a structure:
183
+
184
+ ```python
185
+ from pydantic import BaseModel
186
+ from agentprobe import assert_schema
187
+
188
+ class AgentResponse(BaseModel):
189
+ answer: str
190
+ confidence: float
191
+ sources: list[str]
192
+
193
+ def test_output_structure():
194
+ output = my_agent.run("What is the capital of France?")
195
+ result = assert_schema(output, AgentResponse)
196
+ assert result.confidence > 0.8
197
+ ```
198
+
199
+ ### 5. Multi-Step Tracing
200
+
201
+ Record what an agent did step by step, then assert over the trace or snapshot it. `trace.tool_calls` drops straight into the assertion helpers:
202
+
203
+ ```python
204
+ from agentprobe import Trace, assert_tool_sequence
205
+
206
+ def test_research_flow():
207
+ trace = Trace()
208
+ # record steps as your agent runs (tool calls, LLM turns, custom events)
209
+ trace.record_llm("planning the search")
210
+ trace.record_tool_call("search", {"query": "rainfall 2023"})
211
+ trace.record_event("retry", attempt=2)
212
+ trace.record_tool_call("fetch", {"url": "https://example.com"})
213
+
214
+ assert_tool_sequence(trace.tool_calls, ["search", "fetch"])
215
+ assert trace.names == ["llm", "search", "retry", "fetch"]
216
+ # trace.to_dict() is snapshot-friendly for full-run regression tests
217
+ ```
218
+
219
+ ### 6. Cost Tracking
220
+
221
+ Record token usage on a trace and assert the run stayed under a USD budget — catching regressions that quietly burn more money (longer prompts, extra turns, a pricier model). Pricing comes from a dict, a callable, or [TokenTracker](https://github.com/he-yufeng/TokenTracker)'s price table when it's installed:
222
+
223
+ ```python
224
+ from agentprobe import Trace, assert_cost_under
225
+
226
+ def test_run_stays_under_budget():
227
+ trace = Trace()
228
+ trace.record_llm("plan", model="gpt-4o", input_tokens=1200, output_tokens=300)
229
+ trace.record_llm("answer", model="gpt-4o", input_tokens=800, output_tokens=500)
230
+
231
+ # pricing dict: {model: (input_per_1k_usd, output_per_1k_usd)}
232
+ assert_cost_under(trace, 0.05, pricing={"gpt-4o": (0.005, 0.015)})
233
+ # or pricing=None to use TokenTracker's table (pip install toktally)
234
+ ```
235
+
236
+ ## Pytest Integration
237
+
238
+ AgentProbe registers as a pytest plugin automatically. Use the `agentprobe` fixture:
239
+
240
+ ```python
241
+ def test_with_fixture(agentprobe):
242
+ output = my_agent.run("Hello")
243
+ result = agentprobe.capture("greeting_test", output)
244
+ assert result.passed
245
+ ```
246
+
247
+ ### CLI Flags
248
+
249
+ ```bash
250
+ # Run tests normally
251
+ pytest tests/
252
+
253
+ # Update all snapshots (regenerate baselines)
254
+ pytest tests/ --agentprobe-update
255
+
256
+ # Use semantic comparison instead of exact match
257
+ pytest tests/ --agentprobe-mode=semantic --agentprobe-threshold=0.85
258
+ ```
259
+
260
+ When a snapshot changes, AgentProbe prints a unified diff between the stored JSON snapshot
261
+ and the current output, so CI logs show the exact field or sentence that drifted.
262
+
263
+ ### AgentProbe CLI
264
+
265
+ ```bash
266
+ # Run tests
267
+ agentprobe run
268
+
269
+ # Run with semantic comparison
270
+ agentprobe run --mode semantic --threshold 0.9
271
+
272
+ # Update all snapshots
273
+ agentprobe update
274
+ ```
275
+
276
+ ## Comparison Modes
277
+
278
+ | Mode | How it works | When to use |
279
+ |------|-------------|-------------|
280
+ | `exact` (default) | String equality after serialization | Deterministic agents, structured outputs |
281
+ | `semantic` | Cosine similarity via sentence-transformers | Non-deterministic LLM outputs |
282
+
283
+ For semantic mode, install the optional dependency:
284
+
285
+ ```bash
286
+ pip install agentpoke[semantic]
287
+ ```
288
+
289
+ ## MockLLM Features
290
+
291
+ `MockLLM` is a drop-in replacement for `openai.Client` that returns scripted responses:
292
+
293
+ ```python
294
+ from agentprobe import MockLLM
295
+
296
+ # Scripted responses (consumed in order)
297
+ mock = MockLLM(responses=["First response", "Second response"])
298
+
299
+ # Falls back to default after scripted responses are exhausted
300
+ mock = MockLLM(responses=["Only one"], default_response="I don't know")
301
+
302
+ # Simulate tool calls
303
+ mock = MockLLM(responses=[
304
+ {"tool_calls": [{"id": "1", "function": {"name": "search", "arguments": '{"q": "test"}'}}]}
305
+ ])
306
+
307
+ # Check what was called
308
+ mock.create(messages=[{"role": "user", "content": "Hi"}])
309
+ print(mock.calls) # all recorded calls
310
+ print(mock.call_count) # number of calls
311
+
312
+ # Reset for reuse
313
+ mock.reset()
314
+ ```
315
+
316
+ ## How It Compares
317
+
318
+ | Feature | AgentProbe | DeepEval | Promptfoo |
319
+ |---------|-----------|----------|-----------|
320
+ | pytest native | Yes (plugin) | Separate runner | CLI only |
321
+ | Snapshot baselines | Yes | No | No |
322
+ | Semantic comparison | Yes | Yes | Yes |
323
+ | Mock LLM | Yes (built-in) | No | Partial |
324
+ | Tool call assertions | Yes | No | No |
325
+ | Schema validation | Yes (Pydantic) | Partial | No |
326
+ | Cloud required | No | Optional | No |
327
+ | Config format | Python code | Python code | YAML |
328
+
329
+ ## GitHub Actions
330
+
331
+ Add this to your CI pipeline:
332
+
333
+ ```yaml
334
+ - name: Run agent tests
335
+ run: |
336
+ pip install agentpoke
337
+ pytest tests/ -v
338
+ ```
339
+
340
+ Snapshot files (`.agentprobe/snapshots/`) should be committed to your repo so CI can compare against them.
341
+
342
+ ## FAQ
343
+
344
+ **Do I need an API key to use AgentProbe?**
345
+ No. Use `MockLLM` for deterministic tests without any API calls. If you want to test against a real LLM, you'll need the appropriate API key, but that's your agent's dependency, not AgentProbe's.
346
+
347
+ **How does semantic comparison work?**
348
+ It uses sentence-transformers to embed both the baseline and current output, then computes cosine similarity. If the score is above the threshold (default 0.85), the test passes. This handles cases where the wording changes but the meaning stays the same.
349
+
350
+ **Can I use this with LangChain / CrewAI / AutoGen?**
351
+ Yes. AgentProbe doesn't care what framework you use. It tests the output of your agent, not the internals. Just call your agent inside the test function and return the result.
352
+
353
+ **What about flaky tests from non-deterministic outputs?**
354
+ Use semantic mode with an appropriate threshold. If your agent's outputs vary significantly between runs, lower the threshold. If they should be consistent, raise it. You can also use `MockLLM` to make the underlying LLM deterministic.
355
+
356
+ **How do snapshot files work?**
357
+ Snapshots are stored as JSON in `.agentprobe/snapshots/`. The first time you run a test, it creates the baseline. Subsequent runs compare against it. Use `--agentprobe-update` to regenerate baselines after intentional changes.
358
+
359
+ ## Roadmap
360
+
361
+ **Shipped:** async agent tests (`async def`), multi-step tracing of intermediate steps, cost tracking via TokenTracker, an in-terminal visual diff for snapshot mismatches, and `pytest-xdist` parallel runs with atomic snapshot writes.
362
+
363
+ **Planned:**
364
+
365
+ - **Interactive snapshot review** — an `--agentprobe-review` mode that walks each changed snapshot and lets you accept or reject it one at a time, instead of regenerating every baseline at once.
366
+ - **Tool-call sequence assertions** — assert an agent called tools in an expected order, not only that the final output matches, since the order is often where a regression actually hides.
367
+ - **Framework adapters** — first-class step capture for LangChain, LlamaIndex, and the OpenAI Assistants API, so tracing a multi-step run needs no hand-written glue.
368
+ - **Offline semantic mode** — a local embedding backend for semantic comparison, so threshold checks run without an API call per assertion.
369
+
370
+ ## Contributing
371
+
372
+ Contributions welcome. If you're testing AI agents in production and have ideas for what's missing, open an issue.
373
+
374
+ ## Related Projects
375
+
376
+ AgentProbe is part of a small family of agent-testing tools I maintain. A few related ones:
377
+
378
+ - **[CoreCoder](https://github.com/he-yufeng/CoreCoder)** — want to understand how a coding agent really works? Read the whole ~1k-line engine end to end, not a black box.
379
+ - **[RepoWiki](https://github.com/he-yufeng/RepoWiki)** — dropped into an unfamiliar codebase? It gives you a guided wiki and a where-to-start reading path, a self-hostable DeepWiki alternative.
380
+ - **[LiteBench](https://github.com/he-yufeng/LiteBench)** — benchmark any LLM in one command: HumanEval, GSM8K and MMLU built in, plus your own tasks.
381
+ - **[agentcikit](https://github.com/he-yufeng/agentcikit)** — the CI safety layer for LLM agents: replay runs, fence tool calls, and triage failures before they ship.
382
+
383
+ ## License
384
+
385
+ [MIT](LICENSE)
386
+
387
+ ---
388
+
389
+ <div align="center">
390
+
391
+ **Stop shipping untested agents.**
392
+
393
+ [Report a Bug](https://github.com/he-yufeng/AgentProbe/issues) · [Request a Feature](https://github.com/he-yufeng/AgentProbe/issues)
394
+
395
+ </div>