trace-agent-sdk 0.0.3__tar.gz → 0.2.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.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: trace-agent-sdk
3
- Version: 0.0.3
3
+ Version: 0.2.0
4
4
  Summary: TraceAgent SDK for instrumenting AI agents.
5
5
  Author-email: Enrique Javier Villar Cea <villarcea.enriquejavier@gmail.com>, Javier Rodriguez Castellano <javirodriguezcastellano@gmail.com>
6
6
  License: MIT
@@ -29,6 +29,7 @@ Requires-Dist: pytest<9.0.0,>=8.3.2; extra == "dev"
29
29
  ![GitHub Stars](https://img.shields.io/github/stars/LixusSoftware/TraceAgent)
30
30
  ![GitHub License](https://img.shields.io/github/license/LixusSoftware/TraceAgent)
31
31
  ![Tests](https://github.com/LixusSoftware/TraceAgent/actions/workflows/test.yml/badge.svg)
32
+ [![Website](https://img.shields.io/badge/website-traceagent.vercel.app-blue)](https://traceagent.vercel.app)
32
33
 
33
34
  Python SDK for instrumenting AI agents and streaming events to a TraceAgent backend.
34
35
 
@@ -4,6 +4,7 @@
4
4
  ![GitHub Stars](https://img.shields.io/github/stars/LixusSoftware/TraceAgent)
5
5
  ![GitHub License](https://img.shields.io/github/license/LixusSoftware/TraceAgent)
6
6
  ![Tests](https://github.com/LixusSoftware/TraceAgent/actions/workflows/test.yml/badge.svg)
7
+ [![Website](https://img.shields.io/badge/website-traceagent.vercel.app-blue)](https://traceagent.vercel.app)
7
8
 
8
9
  Python SDK for instrumenting AI agents and streaming events to a TraceAgent backend.
9
10
 
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
4
4
 
5
5
  [project]
6
6
  name = "trace-agent-sdk"
7
- version = "0.0.3"
7
+ version = "0.2.0"
8
8
  description = "TraceAgent SDK for instrumenting AI agents."
9
9
  readme = "README.md"
10
10
  authors = [
@@ -8,7 +8,7 @@ import subprocess
8
8
  import time
9
9
  from dataclasses import dataclass, field
10
10
  from pathlib import Path
11
- from typing import Any, Callable
11
+ from typing import Any, Callable, get_type_hints
12
12
  from uuid import uuid4
13
13
 
14
14
  import httpx
@@ -469,6 +469,34 @@ class RunSession:
469
469
  record_retriever_events=record_retriever_events,
470
470
  )
471
471
 
472
+ def as_pydantic_ai_capability(
473
+ self,
474
+ *,
475
+ record_model_events: bool = True,
476
+ record_tool_events: bool = True,
477
+ record_node_events: bool = False,
478
+ redact_fields: list[str] | None = None,
479
+ failure_cooldown: float = 30.0,
480
+ flush_timeout: float = 5.0,
481
+ ) -> Any:
482
+ """Return a Pydantic AI capability bound to this run.
483
+
484
+ Usage:
485
+ agent = Agent("openai:gpt-4o", capabilities=[run.as_pydantic_ai_capability()])
486
+ agent.run_sync("...")
487
+ """
488
+ from trace_agent_pydantic_ai import TraceAgentCapability
489
+
490
+ return TraceAgentCapability(
491
+ self,
492
+ record_model_events=record_model_events,
493
+ record_tool_events=record_tool_events,
494
+ record_node_events=record_node_events,
495
+ redact_fields=redact_fields or [],
496
+ failure_cooldown=failure_cooldown,
497
+ flush_timeout=flush_timeout,
498
+ )
499
+
472
500
  def _execute_tool(self, call: dict[str, Any]) -> tuple[dict[str, Any], dict[str, Any]]:
473
501
  binding = self.tools.get(call["name"])
474
502
  if binding is None:
@@ -520,10 +548,15 @@ class RunSession:
520
548
 
521
549
  def _build_input_schema(self, func: Callable[..., Any]) -> dict[str, Any]:
522
550
  signature = inspect.signature(func)
551
+ try:
552
+ # Resolves postponed annotations ("int" under `from __future__ import annotations`).
553
+ hints = get_type_hints(func)
554
+ except Exception:
555
+ hints = {}
523
556
  properties: dict[str, Any] = {}
524
557
  required: list[str] = []
525
558
  for parameter_name, parameter in signature.parameters.items():
526
- annotation = parameter.annotation if parameter.annotation is not inspect._empty else str
559
+ annotation = hints.get(parameter_name, parameter.annotation if parameter.annotation is not inspect._empty else str)
527
560
  properties[parameter_name] = {"type": TYPE_MAP.get(annotation, "string")}
528
561
  if parameter.default is inspect._empty:
529
562
  required.append(parameter_name)
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: trace-agent-sdk
3
- Version: 0.0.3
3
+ Version: 0.2.0
4
4
  Summary: TraceAgent SDK for instrumenting AI agents.
5
5
  Author-email: Enrique Javier Villar Cea <villarcea.enriquejavier@gmail.com>, Javier Rodriguez Castellano <javirodriguezcastellano@gmail.com>
6
6
  License: MIT
@@ -29,6 +29,7 @@ Requires-Dist: pytest<9.0.0,>=8.3.2; extra == "dev"
29
29
  ![GitHub Stars](https://img.shields.io/github/stars/LixusSoftware/TraceAgent)
30
30
  ![GitHub License](https://img.shields.io/github/license/LixusSoftware/TraceAgent)
31
31
  ![Tests](https://github.com/LixusSoftware/TraceAgent/actions/workflows/test.yml/badge.svg)
32
+ [![Website](https://img.shields.io/badge/website-traceagent.vercel.app-blue)](https://traceagent.vercel.app)
32
33
 
33
34
  Python SDK for instrumenting AI agents and streaming events to a TraceAgent backend.
34
35
 
@@ -7,4 +7,6 @@ src/trace_agent_sdk.egg-info/PKG-INFO
7
7
  src/trace_agent_sdk.egg-info/SOURCES.txt
8
8
  src/trace_agent_sdk.egg-info/dependency_links.txt
9
9
  src/trace_agent_sdk.egg-info/requires.txt
10
- src/trace_agent_sdk.egg-info/top_level.txt
10
+ src/trace_agent_sdk.egg-info/top_level.txt
11
+ tests/test_client.py
12
+ tests/test_redaction.py
@@ -0,0 +1,241 @@
1
+ from __future__ import annotations
2
+
3
+ import json
4
+ from typing import Any
5
+
6
+ import httpx
7
+ import pytest
8
+
9
+ from trace_agent_sdk import RunSession, TraceAgentClient
10
+
11
+
12
+ class Backend:
13
+ """In-memory stand-in for the TraceAgent API that records every request."""
14
+
15
+ def __init__(self) -> None:
16
+ self.requests: list[tuple[str, str, dict[str, Any]]] = []
17
+ self.responses: dict[str, list[httpx.Response]] = {}
18
+
19
+ def queue(self, path_suffix: str, *responses: httpx.Response) -> None:
20
+ self.responses.setdefault(path_suffix, []).extend(responses)
21
+
22
+ def __call__(self, request: httpx.Request) -> httpx.Response:
23
+ body = json.loads(request.content) if request.content else {}
24
+ self.requests.append((request.method, request.url.path, body))
25
+ for suffix, queued in self.responses.items():
26
+ if request.url.path.endswith(suffix) and queued:
27
+ return queued.pop(0)
28
+ if request.url.path == "/api/runs":
29
+ return httpx.Response(200, json={"run_id": "run-1"})
30
+ return httpx.Response(200, json={"status": "ok"})
31
+
32
+ def observations(self, kind: str | None = None) -> list[dict[str, Any]]:
33
+ found = [
34
+ obs
35
+ for _, path, body in self.requests
36
+ if path.endswith("/observations")
37
+ for obs in body["observations"]
38
+ ]
39
+ return [obs for obs in found if kind is None or obs["kind"] == kind]
40
+
41
+
42
+ @pytest.fixture
43
+ def backend() -> Backend:
44
+ return Backend()
45
+
46
+
47
+ @pytest.fixture
48
+ def client(backend) -> TraceAgentClient:
49
+ http = httpx.Client(transport=httpx.MockTransport(backend), base_url="http://trace.test")
50
+ return TraceAgentClient(base_url="http://trace.test", http_client=http)
51
+
52
+
53
+ @pytest.fixture
54
+ def run(client) -> RunSession:
55
+ return client.start_run("agent", "goal", metadata={"env": "test"})
56
+
57
+
58
+ class TestClient:
59
+ def test_start_run_posts_and_returns_a_bound_session(self, client, backend) -> None:
60
+ session = client.start_run("agent", "goal", metadata={"env": "test"})
61
+
62
+ assert session.run_id == "run-1"
63
+ assert (session.agent_name, session.goal, session.metadata) == ("agent", "goal", {"env": "test"})
64
+ assert backend.requests[0] == ("POST", "/api/runs", {"agent_name": "agent", "goal": "goal", "metadata": {"env": "test"}})
65
+
66
+ def test_error_responses_raise_with_the_server_detail(self, client, backend) -> None:
67
+ backend.queue("/api/runs", httpx.Response(404, json={"detail": "Run not found."}))
68
+
69
+ with pytest.raises(httpx.HTTPStatusError, match="404 for /api/runs: Run not found."):
70
+ client.start_run("agent", "goal")
71
+
72
+ def test_error_responses_without_json_fall_back_to_the_body_text(self, client, backend) -> None:
73
+ backend.queue("/api/runs", httpx.Response(502, text="bad gateway"))
74
+
75
+ with pytest.raises(httpx.HTTPStatusError, match="bad gateway"):
76
+ client.start_run("agent", "goal")
77
+
78
+ def test_api_key_and_base_url_configure_the_http_client(self) -> None:
79
+ client = TraceAgentClient(base_url="http://trace.test/", api_key="secret")
80
+
81
+ assert client.http.headers["Authorization"] == "Bearer secret"
82
+ assert str(client.http.base_url).rstrip("/") == "http://trace.test"
83
+
84
+
85
+ class TestRecording:
86
+ def test_command_renders_lists_and_summarizes_output(self, run, backend) -> None:
87
+ run.record_command(["python", "-m", "pytest"], status="failed", output={"stdout": "x"}, exit_code=1, cwd=".")
88
+
89
+ obs = backend.observations("command")[0]
90
+ assert obs["command"] == "python -m pytest"
91
+ assert obs["status"] == "failed"
92
+ assert obs["exit_code"] == 1
93
+ assert obs["step_id"].startswith("command-")
94
+ assert json.loads(obs["output_summary"]) == {"stdout": "x"}
95
+
96
+ def test_file_read_hashes_the_content(self, run, backend) -> None:
97
+ run.record_file_read("src/app.py", content="print('x')")
98
+
99
+ obs = backend.observations("file_read")[0]
100
+ assert obs["path"] == "src/app.py"
101
+ assert obs["extension"] == ".py"
102
+ assert obs["size_bytes"] == len("print('x')")
103
+ assert len(obs["content_hash"]) == 64
104
+
105
+ def test_file_write_counts_added_and_removed_lines(self, run, backend) -> None:
106
+ run.record_file_write("a.txt", change_type="update", before_content="one\ntwo\n", after_content="one\nthree\nfour\n")
107
+
108
+ obs = backend.observations("file_write")[0]
109
+ assert (obs["line_additions"], obs["line_deletions"]) == (2, 1)
110
+ assert obs["before_hash"] != obs["after_hash"]
111
+ assert obs["change_type"] == "update"
112
+
113
+ def test_file_write_without_changes_reports_no_diff(self, run, backend) -> None:
114
+ run.record_file_write("a.txt", before_content="same", after_content="same")
115
+
116
+ obs = backend.observations("file_write")[0]
117
+ assert obs["diff_summary"] == "No textual diff."
118
+ assert (obs["line_additions"], obs["line_deletions"]) == (0, 0)
119
+
120
+ def test_patch_prefers_the_given_diff_text(self, run, backend) -> None:
121
+ run.record_patch("a.py", diff_text="--- custom diff", before_content="a", after_content="b")
122
+
123
+ assert backend.observations("patch")[0]["diff_summary"] == "--- custom diff"
124
+
125
+ def test_artifact_carries_a_content_hash(self, run, backend) -> None:
126
+ run.record_artifact("report", "weekly", "summary", content={"k": 1}, metadata={"source": "test"})
127
+
128
+ artifact = backend.observations("artifact")[0]["artifact"]
129
+ assert (artifact["kind"], artifact["label"]) == ("report", "weekly")
130
+ assert artifact["metadata"]["source"] == "test"
131
+ assert len(artifact["metadata"]["content_hash"]) == 64
132
+
133
+ def test_finish_and_fail_payloads(self, run, backend) -> None:
134
+ run.finish("done", artifacts=[{"kind": "k"}])
135
+ run.fail("ValueError", "boom")
136
+
137
+ assert backend.requests[-2] == ("POST", "/api/runs/run-1/finish", {"final_output": "done", "artifacts": [{"kind": "k"}]})
138
+ assert backend.requests[-1] == ("POST", "/api/runs/run-1/fail", {"error_code": "ValueError", "error_summary": "boom"})
139
+
140
+
141
+ class TestToolsAndTurns:
142
+ def test_tool_decorator_builds_the_input_schema(self, run) -> None:
143
+ @run.tool(description="Look things up", redact_fields=["api_key"])
144
+ def search(query: str, limit: int = 5, exact: bool = False) -> str:
145
+ return query
146
+
147
+ binding = run.tools["search"]
148
+ assert binding.description == "Look things up"
149
+ assert binding.redact_fields == ["api_key"]
150
+ assert binding.input_schema == {
151
+ "type": "object",
152
+ "properties": {"query": {"type": "string"}, "limit": {"type": "integer"}, "exact": {"type": "boolean"}},
153
+ "required": ["query"],
154
+ }
155
+
156
+ def test_tool_description_defaults_to_the_docstring(self, run) -> None:
157
+ @run.tool()
158
+ def lookup(item: str) -> str:
159
+ """Find an item."""
160
+ return item
161
+
162
+ assert run.tools["lookup"].description == "Find an item."
163
+
164
+ def test_model_turn_runs_requested_tools_and_returns_the_answer(self, run, backend) -> None:
165
+ @run.tool()
166
+ def get_weather(city: str) -> str:
167
+ return f"Sunny in {city}"
168
+
169
+ backend.queue(
170
+ "/turns",
171
+ httpx.Response(
172
+ 200,
173
+ json={
174
+ "status": "requires_action",
175
+ "tool_calls": [{"id": "call-1", "name": "get_weather", "arguments": {"city": "Madrid"}, "step_id": "step-1"}],
176
+ },
177
+ ),
178
+ httpx.Response(200, json={"status": "message", "assistant_message": {"role": "assistant", "content": "It is sunny."}}),
179
+ )
180
+
181
+ result = run.create_model_turn([{"role": "user", "content": "Weather?"}], model="m")
182
+
183
+ assert result.assistant_message["content"] == "It is sunny."
184
+ assert [t["name"] for t in result.executed_tools] == ["get_weather"]
185
+ posted = next(body for _, path, body in backend.requests if path.endswith("/tool-results"))["results"][0]
186
+ assert posted["status"] == "succeeded"
187
+ assert posted["args"] == {"city": "Madrid"}
188
+ assert posted["step_id"] == "step-1"
189
+ roles = [m["role"] for m in run.messages]
190
+ assert roles == ["user", "assistant", "tool", "assistant"]
191
+
192
+ def test_a_failing_tool_is_reported_and_the_error_goes_back_to_the_model(self, run, backend) -> None:
193
+ @run.tool()
194
+ def fetch(url: str) -> str:
195
+ raise ConnectionError("down")
196
+
197
+ backend.queue(
198
+ "/turns",
199
+ httpx.Response(
200
+ 200,
201
+ json={"status": "requires_action", "tool_calls": [{"id": "c", "name": "fetch", "arguments": {"url": "u"}, "step_id": "s"}]},
202
+ ),
203
+ httpx.Response(200, json={"status": "message", "assistant_message": {"role": "assistant", "content": "failed"}}),
204
+ )
205
+
206
+ run.create_model_turn([{"role": "user", "content": "go"}], model="m")
207
+
208
+ posted = next(body for _, path, body in backend.requests if path.endswith("/tool-results"))["results"][0]
209
+ assert (posted["status"], posted["error_code"], posted["error_summary"]) == ("failed", "ConnectionError", "down")
210
+ tool_message = next(m for m in run.messages if m["role"] == "tool")
211
+ assert json.loads(tool_message["content"]) == {"error": "down"}
212
+
213
+ def test_calling_an_unregistered_tool_raises(self, run) -> None:
214
+ with pytest.raises(KeyError, match="not registered"):
215
+ run._execute_tool({"id": "c", "name": "ghost", "arguments": {}, "step_id": "s"})
216
+
217
+ def test_redacted_tool_output_is_masked_in_the_summary(self, run) -> None:
218
+ @run.tool(redact_fields=["token"])
219
+ def login(user: str) -> dict:
220
+ return {"user": user, "token": "abc123"}
221
+
222
+ payload, _ = run._execute_tool({"id": "c", "name": "login", "arguments": {"user": "ann"}, "step_id": "s"})
223
+
224
+ assert "abc123" not in payload["output_summary"]
225
+ assert "[redacted]" in payload["output_summary"]
226
+
227
+
228
+ class TestIntegrationsFactories:
229
+ def test_langchain_callback_is_bound_to_the_run(self, run) -> None:
230
+ pytest.importorskip("trace_agent_langchain")
231
+ callback = run.as_langchain_callback(record_chain_events=False)
232
+
233
+ assert callback.run_session is run
234
+ assert callback.record_chain_events is False
235
+
236
+ def test_pydantic_ai_capability_is_bound_to_the_run(self, run) -> None:
237
+ pytest.importorskip("trace_agent_pydantic_ai")
238
+ capability = run.as_pydantic_ai_capability(record_node_events=True, failure_cooldown=5, flush_timeout=1)
239
+
240
+ assert capability.run_session is run
241
+ assert (capability.record_node_events, capability.failure_cooldown, capability.flush_timeout) == (True, 5, 1)
@@ -0,0 +1,54 @@
1
+ from __future__ import annotations
2
+
3
+ from trace_agent_sdk.redaction import redact_payload, stable_hash, summarize_value
4
+
5
+
6
+ class TestStableHash:
7
+ def test_is_independent_of_key_order(self) -> None:
8
+ assert stable_hash({"a": 1, "b": [1, 2]}) == stable_hash({"b": [1, 2], "a": 1})
9
+
10
+ def test_differs_for_different_values(self) -> None:
11
+ assert stable_hash({"a": 1}) != stable_hash({"a": 2})
12
+
13
+ def test_handles_values_json_cannot_encode(self) -> None:
14
+ assert stable_hash({"when": object}) == stable_hash({"when": object})
15
+
16
+
17
+ class TestRedactPayload:
18
+ def test_redacts_listed_fields_at_any_depth(self) -> None:
19
+ payload = {"token": "abc", "nested": {"token": "def", "keep": 1}, "items": [{"token": "ghi"}]}
20
+
21
+ redacted = redact_payload(payload, ["token"])
22
+
23
+ assert redacted == {"token": "[redacted]", "nested": {"token": "[redacted]", "keep": 1}, "items": [{"token": "[redacted]"}]}
24
+
25
+ def test_does_not_mutate_the_input(self) -> None:
26
+ payload = {"token": "abc"}
27
+ redact_payload(payload, ["token"])
28
+ assert payload == {"token": "abc"}
29
+
30
+ def test_truncates_long_strings_but_not_short_ones(self) -> None:
31
+ assert redact_payload("short") == "short"
32
+ long_value = "x" * 500
33
+ redacted = redact_payload(long_value)
34
+ assert len(redacted) == 160
35
+ assert redacted.endswith("...")
36
+
37
+ def test_leaves_other_types_untouched(self) -> None:
38
+ assert redact_payload({"n": 3, "ok": True, "none": None}) == {"n": 3, "ok": True, "none": None}
39
+
40
+
41
+ class TestSummarizeValue:
42
+ def test_strings_are_returned_as_is(self) -> None:
43
+ assert summarize_value("hello") == "hello"
44
+
45
+ def test_dicts_are_rendered_as_sorted_json(self) -> None:
46
+ assert summarize_value({"b": 1, "a": 2}) == '{"a": 2, "b": 1}'
47
+
48
+ def test_output_is_capped_at_max_length(self) -> None:
49
+ summary = summarize_value({"items": list(range(500))}, max_length=50)
50
+ assert len(summary) == 50
51
+ assert summary.endswith("...")
52
+
53
+ def test_applies_redaction(self) -> None:
54
+ assert "secret" not in summarize_value({"password": "secret"}, ["password"])