retrace-sdk 0.1.3__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.
- retrace_sdk-0.1.3/.gitignore +16 -0
- retrace_sdk-0.1.3/PKG-INFO +56 -0
- retrace_sdk-0.1.3/README.md +39 -0
- retrace_sdk-0.1.3/pyproject.toml +30 -0
- retrace_sdk-0.1.3/src/retrace/__init__.py +12 -0
- retrace_sdk-0.1.3/src/retrace/config.py +52 -0
- retrace_sdk-0.1.3/src/retrace/interceptors/__init__.py +3 -0
- retrace_sdk-0.1.3/src/retrace/interceptors/gemini.py +92 -0
- retrace_sdk-0.1.3/src/retrace/recorder.py +277 -0
- retrace_sdk-0.1.3/src/retrace/trace.py +115 -0
- retrace_sdk-0.1.3/src/retrace/transport.py +155 -0
- retrace_sdk-0.1.3/src/retrace/utils.py +25 -0
- retrace_sdk-0.1.3/tests/test_config.py +21 -0
|
@@ -0,0 +1,56 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: retrace-sdk
|
|
3
|
+
Version: 0.1.3
|
|
4
|
+
Summary: Record, replay, fork & share AI agent executions
|
|
5
|
+
Project-URL: Homepage, https://retrace.yashbogam.me
|
|
6
|
+
Project-URL: Repository, https://github.com/yash1511-bogam/retrace
|
|
7
|
+
Author: Yash Bogam
|
|
8
|
+
License: MIT
|
|
9
|
+
Requires-Python: >=3.10
|
|
10
|
+
Requires-Dist: requests>=2.32.0
|
|
11
|
+
Requires-Dist: websocket-client>=1.9.0
|
|
12
|
+
Provides-Extra: dev
|
|
13
|
+
Requires-Dist: pytest>=8.0; extra == 'dev'
|
|
14
|
+
Provides-Extra: gemini
|
|
15
|
+
Requires-Dist: google-genai>=1.52.0; extra == 'gemini'
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# retrace-sdk
|
|
19
|
+
|
|
20
|
+
The execution replay engine for AI agents. Record every LLM call, tool invocation, and error your AI agent makes. Replay step-by-step. Fork from any point. Share interactive traces via URL.
|
|
21
|
+
|
|
22
|
+
## Install
|
|
23
|
+
|
|
24
|
+
```bash
|
|
25
|
+
pip install retrace-sdk
|
|
26
|
+
```
|
|
27
|
+
|
|
28
|
+
## Quick Start
|
|
29
|
+
|
|
30
|
+
```python
|
|
31
|
+
from retrace_sdk import configure, record
|
|
32
|
+
|
|
33
|
+
configure(api_key="rt_live_...")
|
|
34
|
+
|
|
35
|
+
@record(name="my-agent")
|
|
36
|
+
def run_agent(prompt: str):
|
|
37
|
+
response = client.chat.completions.create(
|
|
38
|
+
model="gpt-4o",
|
|
39
|
+
messages=[{"role": "user", "content": prompt}]
|
|
40
|
+
)
|
|
41
|
+
return response.choices[0].message.content
|
|
42
|
+
|
|
43
|
+
run_agent("What is quantum computing?")
|
|
44
|
+
```
|
|
45
|
+
|
|
46
|
+
## Features
|
|
47
|
+
|
|
48
|
+
- **Record** — One decorator captures every LLM call, tool call, and error
|
|
49
|
+
- **Replay** — Step through executions with play/pause/speed controls
|
|
50
|
+
- **Fork** — Branch from any step, modify input, watch a new path diverge
|
|
51
|
+
- **Share** — Publish traces as shareable "tapes" with interactive playback
|
|
52
|
+
|
|
53
|
+
## Links
|
|
54
|
+
|
|
55
|
+
- [Documentation](https://retrace.yashbogam.me/docs)
|
|
56
|
+
- [GitHub](https://github.com/yash1511-bogam/retrace)
|
|
@@ -0,0 +1,39 @@
|
|
|
1
|
+
# retrace-sdk
|
|
2
|
+
|
|
3
|
+
The execution replay engine for AI agents. Record every LLM call, tool invocation, and error your AI agent makes. Replay step-by-step. Fork from any point. Share interactive traces via URL.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```bash
|
|
8
|
+
pip install retrace-sdk
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
## Quick Start
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
from retrace_sdk import configure, record
|
|
15
|
+
|
|
16
|
+
configure(api_key="rt_live_...")
|
|
17
|
+
|
|
18
|
+
@record(name="my-agent")
|
|
19
|
+
def run_agent(prompt: str):
|
|
20
|
+
response = client.chat.completions.create(
|
|
21
|
+
model="gpt-4o",
|
|
22
|
+
messages=[{"role": "user", "content": prompt}]
|
|
23
|
+
)
|
|
24
|
+
return response.choices[0].message.content
|
|
25
|
+
|
|
26
|
+
run_agent("What is quantum computing?")
|
|
27
|
+
```
|
|
28
|
+
|
|
29
|
+
## Features
|
|
30
|
+
|
|
31
|
+
- **Record** — One decorator captures every LLM call, tool call, and error
|
|
32
|
+
- **Replay** — Step through executions with play/pause/speed controls
|
|
33
|
+
- **Fork** — Branch from any step, modify input, watch a new path diverge
|
|
34
|
+
- **Share** — Publish traces as shareable "tapes" with interactive playback
|
|
35
|
+
|
|
36
|
+
## Links
|
|
37
|
+
|
|
38
|
+
- [Documentation](https://retrace.yashbogam.me/docs)
|
|
39
|
+
- [GitHub](https://github.com/yash1511-bogam/retrace)
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "retrace-sdk"
|
|
3
|
+
version = "0.1.3"
|
|
4
|
+
description = "Record, replay, fork & share AI agent executions"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
requires-python = ">=3.10"
|
|
7
|
+
license = {text = "MIT"}
|
|
8
|
+
authors = [{name = "Yash Bogam"}]
|
|
9
|
+
dependencies = [
|
|
10
|
+
"websocket-client>=1.9.0",
|
|
11
|
+
"requests>=2.32.0",
|
|
12
|
+
]
|
|
13
|
+
|
|
14
|
+
[project.urls]
|
|
15
|
+
Homepage = "https://retrace.yashbogam.me"
|
|
16
|
+
Repository = "https://github.com/yash1511-bogam/retrace"
|
|
17
|
+
|
|
18
|
+
[project.optional-dependencies]
|
|
19
|
+
gemini = ["google-genai>=1.52.0"]
|
|
20
|
+
dev = ["pytest>=8.0"]
|
|
21
|
+
|
|
22
|
+
[build-system]
|
|
23
|
+
requires = ["hatchling"]
|
|
24
|
+
build-backend = "hatchling.build"
|
|
25
|
+
|
|
26
|
+
[tool.pytest.ini_options]
|
|
27
|
+
testpaths = ["tests"]
|
|
28
|
+
|
|
29
|
+
[tool.hatch.build.targets.wheel]
|
|
30
|
+
packages = ["src/retrace"]
|
|
@@ -0,0 +1,12 @@
|
|
|
1
|
+
from .config import configure, get_config
|
|
2
|
+
from .recorder import record, TraceRecorder
|
|
3
|
+
from .trace import Span, Trace, SpanType, TraceStatus
|
|
4
|
+
from .interceptors.gemini import install_gemini_interceptor, uninstall_gemini_interceptor
|
|
5
|
+
|
|
6
|
+
__version__ = "0.1.0"
|
|
7
|
+
__all__ = [
|
|
8
|
+
"configure", "get_config",
|
|
9
|
+
"record", "TraceRecorder",
|
|
10
|
+
"Span", "Trace", "SpanType", "TraceStatus",
|
|
11
|
+
"install_gemini_interceptor", "uninstall_gemini_interceptor",
|
|
12
|
+
]
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
import logging
|
|
2
|
+
import os
|
|
3
|
+
from dataclasses import dataclass, field
|
|
4
|
+
|
|
5
|
+
logger = logging.getLogger("retrace")
|
|
6
|
+
|
|
7
|
+
|
|
8
|
+
@dataclass
|
|
9
|
+
class RetraceConfig:
|
|
10
|
+
api_key: str = ""
|
|
11
|
+
base_url: str = ""
|
|
12
|
+
project_id: str | None = None
|
|
13
|
+
ws_url: str = ""
|
|
14
|
+
flush_interval: float = 2.0
|
|
15
|
+
enabled: bool = True
|
|
16
|
+
|
|
17
|
+
def __post_init__(self):
|
|
18
|
+
if not self.api_key:
|
|
19
|
+
self.api_key = os.environ.get("RETRACE_API_KEY", "")
|
|
20
|
+
if not self.base_url:
|
|
21
|
+
self.base_url = os.environ.get("RETRACE_BASE_URL", "http://localhost:3001")
|
|
22
|
+
if not self.project_id:
|
|
23
|
+
self.project_id = os.environ.get("RETRACE_PROJECT_ID") or None
|
|
24
|
+
if not self.ws_url:
|
|
25
|
+
self.ws_url = self.base_url.replace("https://", "wss://").replace("http://", "ws://")
|
|
26
|
+
enabled_env = os.environ.get("RETRACE_ENABLED", "true").lower()
|
|
27
|
+
if enabled_env in ("false", "0", "no"):
|
|
28
|
+
self.enabled = False
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
_config: RetraceConfig | None = None
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def get_config() -> RetraceConfig:
|
|
35
|
+
global _config
|
|
36
|
+
if _config is None:
|
|
37
|
+
_config = RetraceConfig()
|
|
38
|
+
return _config
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def configure(**kwargs) -> RetraceConfig:
|
|
42
|
+
global _config
|
|
43
|
+
if _config is None:
|
|
44
|
+
_config = RetraceConfig(**kwargs)
|
|
45
|
+
else:
|
|
46
|
+
for k, v in kwargs.items():
|
|
47
|
+
setattr(_config, k, v)
|
|
48
|
+
if "base_url" in kwargs and "ws_url" not in kwargs:
|
|
49
|
+
_config.ws_url = _config.base_url.replace("https://", "wss://").replace("http://", "ws://")
|
|
50
|
+
if _config.api_key and not _config.api_key.startswith("rt_live_"):
|
|
51
|
+
logger.warning("API key does not start with 'rt_live_'. This may be invalid.")
|
|
52
|
+
return _config
|
|
@@ -0,0 +1,92 @@
|
|
|
1
|
+
"""Gemini interceptor for Retrace Python SDK."""
|
|
2
|
+
import time
|
|
3
|
+
import uuid
|
|
4
|
+
from typing import Callable
|
|
5
|
+
|
|
6
|
+
_original_generate = None
|
|
7
|
+
_installed = False
|
|
8
|
+
_on_span = None
|
|
9
|
+
|
|
10
|
+
PRICING = {
|
|
11
|
+
"gemini-3.1-pro-preview": (2.0, 12.0),
|
|
12
|
+
"gemini-2.5-pro": (1.25, 10.0),
|
|
13
|
+
"gemini-2.5-flash": (0.15, 0.60),
|
|
14
|
+
"gemini-2.0-flash": (0.10, 0.40),
|
|
15
|
+
}
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def _calc_cost(model: str, input_tokens: int, output_tokens: int) -> float:
|
|
19
|
+
p = PRICING.get(model, (0, 0))
|
|
20
|
+
return (input_tokens * p[0] + output_tokens * p[1]) / 1_000_000
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def install_gemini_interceptor(on_span=None):
|
|
24
|
+
global _original_generate, _installed, _on_span
|
|
25
|
+
if _installed:
|
|
26
|
+
if on_span:
|
|
27
|
+
_on_span = on_span
|
|
28
|
+
return
|
|
29
|
+
|
|
30
|
+
try:
|
|
31
|
+
from google import genai
|
|
32
|
+
except ImportError:
|
|
33
|
+
return
|
|
34
|
+
|
|
35
|
+
_on_span = on_span
|
|
36
|
+
_original_generate = genai.models.Models.generate_content
|
|
37
|
+
|
|
38
|
+
def patched_generate(self, *args, **kwargs):
|
|
39
|
+
model = kwargs.get("model", args[0] if args else "unknown")
|
|
40
|
+
contents = kwargs.get("contents", args[1] if len(args) > 1 else None)
|
|
41
|
+
span_id = str(uuid.uuid4())
|
|
42
|
+
start = time.time()
|
|
43
|
+
|
|
44
|
+
try:
|
|
45
|
+
result = _original_generate(self, *args, **kwargs)
|
|
46
|
+
duration_ms = int((time.time() - start) * 1000)
|
|
47
|
+
input_tokens = getattr(getattr(result, "usage_metadata", None), "prompt_token_count", 0) or 0
|
|
48
|
+
output_tokens = getattr(getattr(result, "usage_metadata", None), "candidates_token_count", 0) or 0
|
|
49
|
+
|
|
50
|
+
if _on_span:
|
|
51
|
+
_on_span({
|
|
52
|
+
"id": span_id,
|
|
53
|
+
"span_type": "llm_call",
|
|
54
|
+
"name": "retrace.ai.generate",
|
|
55
|
+
"model": model,
|
|
56
|
+
"input": str(contents)[:2000] if contents else None,
|
|
57
|
+
"output": getattr(result, "text", "")[:2000],
|
|
58
|
+
"input_tokens": input_tokens,
|
|
59
|
+
"output_tokens": output_tokens,
|
|
60
|
+
"cost": _calc_cost(model, input_tokens, output_tokens),
|
|
61
|
+
"duration_ms": duration_ms,
|
|
62
|
+
})
|
|
63
|
+
return result
|
|
64
|
+
except Exception as e:
|
|
65
|
+
duration_ms = int((time.time() - start) * 1000)
|
|
66
|
+
if _on_span:
|
|
67
|
+
_on_span({
|
|
68
|
+
"id": span_id,
|
|
69
|
+
"span_type": "llm_call",
|
|
70
|
+
"name": "retrace.ai.generate",
|
|
71
|
+
"model": model,
|
|
72
|
+
"input": str(contents)[:2000] if contents else None,
|
|
73
|
+
"duration_ms": duration_ms,
|
|
74
|
+
"error": str(e),
|
|
75
|
+
})
|
|
76
|
+
raise
|
|
77
|
+
|
|
78
|
+
genai.models.Models.generate_content = patched_generate
|
|
79
|
+
_installed = True
|
|
80
|
+
|
|
81
|
+
|
|
82
|
+
def uninstall_gemini_interceptor():
|
|
83
|
+
global _installed, _on_span, _original_generate
|
|
84
|
+
if not _installed or not _original_generate:
|
|
85
|
+
return
|
|
86
|
+
try:
|
|
87
|
+
from google import genai
|
|
88
|
+
genai.models.Models.generate_content = _original_generate
|
|
89
|
+
except ImportError:
|
|
90
|
+
pass
|
|
91
|
+
_installed = False
|
|
92
|
+
_on_span = None
|
|
@@ -0,0 +1,277 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import functools
|
|
4
|
+
import logging
|
|
5
|
+
import threading
|
|
6
|
+
import time
|
|
7
|
+
from typing import Any, Callable
|
|
8
|
+
|
|
9
|
+
logger = logging.getLogger("retrace")
|
|
10
|
+
|
|
11
|
+
from .config import get_config
|
|
12
|
+
from .trace import Trace, Span, SpanType, TraceStatus
|
|
13
|
+
from .transport import create_transport, WSTransport, HTTPTransport
|
|
14
|
+
from .utils import gen_id, utcnow
|
|
15
|
+
|
|
16
|
+
|
|
17
|
+
class TraceRecorder:
|
|
18
|
+
"""Manages recording of a single trace and its spans."""
|
|
19
|
+
|
|
20
|
+
def __init__(self, name: str | None = None, input: Any = None, metadata: dict | None = None):
|
|
21
|
+
self._trace = Trace(
|
|
22
|
+
name=name,
|
|
23
|
+
input=input,
|
|
24
|
+
metadata=metadata or {},
|
|
25
|
+
project_id=get_config().project_id,
|
|
26
|
+
)
|
|
27
|
+
self._lock = threading.Lock()
|
|
28
|
+
self._transport = create_transport()
|
|
29
|
+
self._interceptors_installed = False
|
|
30
|
+
|
|
31
|
+
@property
|
|
32
|
+
def trace(self) -> Trace:
|
|
33
|
+
return self._trace
|
|
34
|
+
|
|
35
|
+
@property
|
|
36
|
+
def output(self):
|
|
37
|
+
return self._trace.output
|
|
38
|
+
|
|
39
|
+
@output.setter
|
|
40
|
+
def output(self, value):
|
|
41
|
+
self._trace.output = value
|
|
42
|
+
|
|
43
|
+
def _install_interceptors(self):
|
|
44
|
+
if self._interceptors_installed:
|
|
45
|
+
return
|
|
46
|
+
try:
|
|
47
|
+
from .interceptors.gemini import install_gemini_interceptor
|
|
48
|
+
install_gemini_interceptor(lambda span_data: self._handle_intercepted_span(span_data))
|
|
49
|
+
except Exception as e:
|
|
50
|
+
logger.debug(f"Failed to install interceptors: {e}")
|
|
51
|
+
self._interceptors_installed = True
|
|
52
|
+
|
|
53
|
+
def _handle_intercepted_span(self, span_data: dict):
|
|
54
|
+
span = Span(
|
|
55
|
+
trace_id=self._trace.id,
|
|
56
|
+
span_type=SpanType(span_data.get("span_type", "llm_call")),
|
|
57
|
+
name=span_data.get("name", ""),
|
|
58
|
+
model=span_data.get("model"),
|
|
59
|
+
input=span_data.get("input"),
|
|
60
|
+
output=span_data.get("output"),
|
|
61
|
+
input_tokens=span_data.get("input_tokens"),
|
|
62
|
+
output_tokens=span_data.get("output_tokens"),
|
|
63
|
+
cost=span_data.get("cost"),
|
|
64
|
+
duration_ms=span_data.get("duration_ms"),
|
|
65
|
+
error=span_data.get("error"),
|
|
66
|
+
)
|
|
67
|
+
if span_data.get("id"):
|
|
68
|
+
span.id = span_data["id"]
|
|
69
|
+
span.ended_at = utcnow()
|
|
70
|
+
self.add_span(span)
|
|
71
|
+
|
|
72
|
+
def start_trace(self, name: str | None = None, input: Any = None, metadata: dict | None = None):
|
|
73
|
+
if name:
|
|
74
|
+
self._trace.name = name
|
|
75
|
+
if input is not None:
|
|
76
|
+
self._trace.input = input
|
|
77
|
+
if metadata:
|
|
78
|
+
self._trace.metadata.update(metadata)
|
|
79
|
+
self._install_interceptors()
|
|
80
|
+
self._send("trace_started", self._trace.to_dict())
|
|
81
|
+
|
|
82
|
+
def end_trace(self, output: Any = None, status: TraceStatus = TraceStatus.COMPLETED):
|
|
83
|
+
self._trace.output = output if output is not None else self._trace.output
|
|
84
|
+
self._trace.status = status
|
|
85
|
+
self._trace.ended_at = utcnow()
|
|
86
|
+
if self._trace.started_at:
|
|
87
|
+
self._trace.total_duration_ms = int(
|
|
88
|
+
(self._trace.ended_at - self._trace.started_at).total_seconds() * 1000
|
|
89
|
+
)
|
|
90
|
+
self._send("trace_ended", {
|
|
91
|
+
"id": self._trace.id,
|
|
92
|
+
"ended_at": self._trace.ended_at.isoformat().replace("+00:00", "Z"),
|
|
93
|
+
"output": self._trace.output,
|
|
94
|
+
"status": status.value,
|
|
95
|
+
"total_tokens": self._trace.total_tokens,
|
|
96
|
+
"total_cost": self._trace.total_cost,
|
|
97
|
+
})
|
|
98
|
+
self._transport.close()
|
|
99
|
+
|
|
100
|
+
def add_span(self, span: Span):
|
|
101
|
+
span.trace_id = self._trace.id
|
|
102
|
+
with self._lock:
|
|
103
|
+
self._trace.spans.append(span)
|
|
104
|
+
self._trace.total_tokens += (span.input_tokens or 0) + (span.output_tokens or 0)
|
|
105
|
+
self._trace.total_cost += span.cost or 0.0
|
|
106
|
+
|
|
107
|
+
if span.ended_at:
|
|
108
|
+
# Span is complete — send both started and ended
|
|
109
|
+
self._send("span_started", span.to_dict())
|
|
110
|
+
self._send("span_ended", {
|
|
111
|
+
"id": span.id,
|
|
112
|
+
"ended_at": span.ended_at.isoformat().replace("+00:00", "Z"),
|
|
113
|
+
"output": span.output,
|
|
114
|
+
"output_tokens": span.output_tokens,
|
|
115
|
+
"cost": span.cost,
|
|
116
|
+
"error": span.error,
|
|
117
|
+
})
|
|
118
|
+
else:
|
|
119
|
+
self._send("span_started", span.to_dict())
|
|
120
|
+
|
|
121
|
+
def start_span(
|
|
122
|
+
self,
|
|
123
|
+
name: str,
|
|
124
|
+
span_type: SpanType = SpanType.LLM_CALL,
|
|
125
|
+
input: Any = None,
|
|
126
|
+
model: str | None = None,
|
|
127
|
+
parent_id: str | None = None,
|
|
128
|
+
) -> Span:
|
|
129
|
+
span = Span(
|
|
130
|
+
trace_id=self._trace.id,
|
|
131
|
+
span_type=span_type,
|
|
132
|
+
name=name,
|
|
133
|
+
input=input,
|
|
134
|
+
model=model,
|
|
135
|
+
parent_id=parent_id,
|
|
136
|
+
)
|
|
137
|
+
with self._lock:
|
|
138
|
+
self._trace.spans.append(span)
|
|
139
|
+
self._send("span_started", span.to_dict())
|
|
140
|
+
return span
|
|
141
|
+
|
|
142
|
+
def end_span(self, span_id: str, output: Any = None, error: str | None = None):
|
|
143
|
+
with self._lock:
|
|
144
|
+
span = next((s for s in self._trace.spans if s.id == span_id), None)
|
|
145
|
+
if not span:
|
|
146
|
+
return
|
|
147
|
+
span.output = output
|
|
148
|
+
span.error = error
|
|
149
|
+
span.ended_at = utcnow()
|
|
150
|
+
if span.started_at:
|
|
151
|
+
span.duration_ms = int((span.ended_at - span.started_at).total_seconds() * 1000)
|
|
152
|
+
self._send("span_ended", {
|
|
153
|
+
"id": span.id,
|
|
154
|
+
"ended_at": span.ended_at.isoformat().replace("+00:00", "Z"),
|
|
155
|
+
"output": output,
|
|
156
|
+
"error": error,
|
|
157
|
+
})
|
|
158
|
+
|
|
159
|
+
def _send(self, event_type: str, data: dict):
|
|
160
|
+
try:
|
|
161
|
+
self._transport.send(event_type, data)
|
|
162
|
+
except Exception as e:
|
|
163
|
+
logger.debug(f"Failed to send {event_type}: {e}")
|
|
164
|
+
|
|
165
|
+
# Context manager support
|
|
166
|
+
def __enter__(self):
|
|
167
|
+
self.start_trace()
|
|
168
|
+
return self
|
|
169
|
+
|
|
170
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
171
|
+
if exc_type:
|
|
172
|
+
self.end_trace(status=TraceStatus.FAILED)
|
|
173
|
+
else:
|
|
174
|
+
self.end_trace(status=TraceStatus.COMPLETED)
|
|
175
|
+
return False
|
|
176
|
+
|
|
177
|
+
|
|
178
|
+
def record(name: str | None = None, input: Any = None, metadata: dict | None = None):
|
|
179
|
+
"""Decorator and context manager for recording agent executions.
|
|
180
|
+
|
|
181
|
+
Usage as decorator:
|
|
182
|
+
@retrace.record(name="my-agent")
|
|
183
|
+
def my_agent(prompt):
|
|
184
|
+
...
|
|
185
|
+
|
|
186
|
+
Usage as context manager:
|
|
187
|
+
with retrace.record(name="my-agent", input={"prompt": "hi"}) as t:
|
|
188
|
+
result = agent.run("hi")
|
|
189
|
+
t.output = result
|
|
190
|
+
"""
|
|
191
|
+
cfg = get_config()
|
|
192
|
+
|
|
193
|
+
# If called with a function directly: @record without parens
|
|
194
|
+
if callable(name):
|
|
195
|
+
fn = name
|
|
196
|
+
if not cfg.enabled:
|
|
197
|
+
return fn
|
|
198
|
+
|
|
199
|
+
@functools.wraps(fn)
|
|
200
|
+
def wrapper(*args, **kwargs):
|
|
201
|
+
recorder = TraceRecorder(name=fn.__name__, input={"args": list(args), "kwargs": kwargs})
|
|
202
|
+
recorder.start_trace()
|
|
203
|
+
try:
|
|
204
|
+
result = fn(*args, **kwargs)
|
|
205
|
+
recorder.end_trace(output=result, status=TraceStatus.COMPLETED)
|
|
206
|
+
return result
|
|
207
|
+
except Exception as e:
|
|
208
|
+
recorder.end_trace(status=TraceStatus.FAILED)
|
|
209
|
+
raise
|
|
210
|
+
|
|
211
|
+
return wrapper
|
|
212
|
+
|
|
213
|
+
# Called with arguments: @record(name="...") or as context manager
|
|
214
|
+
def decorator(fn: Callable | None = None):
|
|
215
|
+
if fn is None:
|
|
216
|
+
# Context manager usage
|
|
217
|
+
return TraceRecorder(name=name, input=input, metadata=metadata)
|
|
218
|
+
|
|
219
|
+
if not cfg.enabled:
|
|
220
|
+
return fn
|
|
221
|
+
|
|
222
|
+
@functools.wraps(fn)
|
|
223
|
+
def wrapper(*args, **kwargs):
|
|
224
|
+
recorder = TraceRecorder(
|
|
225
|
+
name=name or fn.__name__,
|
|
226
|
+
input=input if input is not None else {"args": list(args), "kwargs": kwargs},
|
|
227
|
+
metadata=metadata,
|
|
228
|
+
)
|
|
229
|
+
recorder.start_trace()
|
|
230
|
+
try:
|
|
231
|
+
result = fn(*args, **kwargs)
|
|
232
|
+
recorder.end_trace(output=result, status=TraceStatus.COMPLETED)
|
|
233
|
+
return result
|
|
234
|
+
except Exception as e:
|
|
235
|
+
recorder.end_trace(status=TraceStatus.FAILED)
|
|
236
|
+
raise
|
|
237
|
+
|
|
238
|
+
return wrapper
|
|
239
|
+
|
|
240
|
+
# If no function passed, could be context manager or decorator
|
|
241
|
+
if not cfg.enabled:
|
|
242
|
+
# Return a no-op context manager
|
|
243
|
+
class _NoOp:
|
|
244
|
+
output = None
|
|
245
|
+
def __enter__(self): return self
|
|
246
|
+
def __exit__(self, *a): return False
|
|
247
|
+
def __call__(self, fn): return fn
|
|
248
|
+
return _NoOp()
|
|
249
|
+
|
|
250
|
+
# Return something that works as both decorator and context manager
|
|
251
|
+
class _RecordProxy:
|
|
252
|
+
def __init__(self):
|
|
253
|
+
self._recorder = TraceRecorder(name=name, input=input, metadata=metadata)
|
|
254
|
+
|
|
255
|
+
@property
|
|
256
|
+
def output(self):
|
|
257
|
+
return self._recorder.output
|
|
258
|
+
|
|
259
|
+
@output.setter
|
|
260
|
+
def output(self, value):
|
|
261
|
+
self._recorder.output = value
|
|
262
|
+
|
|
263
|
+
def __call__(self, fn):
|
|
264
|
+
return decorator(fn)
|
|
265
|
+
|
|
266
|
+
def __enter__(self):
|
|
267
|
+
self._recorder.start_trace()
|
|
268
|
+
return self
|
|
269
|
+
|
|
270
|
+
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
271
|
+
if exc_type:
|
|
272
|
+
self._recorder.end_trace(status=TraceStatus.FAILED)
|
|
273
|
+
else:
|
|
274
|
+
self._recorder.end_trace(status=TraceStatus.COMPLETED)
|
|
275
|
+
return False
|
|
276
|
+
|
|
277
|
+
return _RecordProxy()
|
|
@@ -0,0 +1,115 @@
|
|
|
1
|
+
from dataclasses import dataclass, field
|
|
2
|
+
from datetime import datetime
|
|
3
|
+
from typing import Any, Optional
|
|
4
|
+
from enum import Enum
|
|
5
|
+
|
|
6
|
+
from .utils import gen_id, utcnow
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class SpanType(str, Enum):
|
|
10
|
+
LLM_CALL = "llm_call"
|
|
11
|
+
TOOL_CALL = "tool_call"
|
|
12
|
+
TOOL_RESULT = "tool_result"
|
|
13
|
+
REASONING = "reasoning"
|
|
14
|
+
ACTION = "action"
|
|
15
|
+
ERROR = "error"
|
|
16
|
+
FORK_POINT = "fork_point"
|
|
17
|
+
|
|
18
|
+
|
|
19
|
+
class TraceStatus(str, Enum):
|
|
20
|
+
RUNNING = "running"
|
|
21
|
+
COMPLETED = "completed"
|
|
22
|
+
FAILED = "failed"
|
|
23
|
+
|
|
24
|
+
|
|
25
|
+
@dataclass
|
|
26
|
+
class Span:
|
|
27
|
+
id: str = field(default_factory=gen_id)
|
|
28
|
+
trace_id: str = ""
|
|
29
|
+
span_type: SpanType = SpanType.LLM_CALL
|
|
30
|
+
name: str = ""
|
|
31
|
+
parent_id: Optional[str] = None
|
|
32
|
+
model: Optional[str] = None
|
|
33
|
+
input: Any = None
|
|
34
|
+
output: Any = None
|
|
35
|
+
input_tokens: Optional[int] = None
|
|
36
|
+
output_tokens: Optional[int] = None
|
|
37
|
+
cost: Optional[float] = None
|
|
38
|
+
duration_ms: Optional[int] = None
|
|
39
|
+
metadata: dict = field(default_factory=dict)
|
|
40
|
+
started_at: Optional[datetime] = field(default_factory=utcnow)
|
|
41
|
+
ended_at: Optional[datetime] = None
|
|
42
|
+
error: Optional[str] = None
|
|
43
|
+
|
|
44
|
+
def to_dict(self) -> dict:
|
|
45
|
+
d: dict[str, Any] = {
|
|
46
|
+
"id": self.id,
|
|
47
|
+
"trace_id": self.trace_id,
|
|
48
|
+
"parent_id": self.parent_id,
|
|
49
|
+
"span_type": self.span_type.value,
|
|
50
|
+
"name": self.name,
|
|
51
|
+
"started_at": self.started_at.isoformat().replace("+00:00", "Z") if self.started_at else None,
|
|
52
|
+
}
|
|
53
|
+
if self.model:
|
|
54
|
+
d["model"] = self.model
|
|
55
|
+
if self.input is not None:
|
|
56
|
+
d["input"] = self.input
|
|
57
|
+
if self.output is not None:
|
|
58
|
+
d["output"] = self.output
|
|
59
|
+
if self.input_tokens is not None:
|
|
60
|
+
d["input_tokens"] = self.input_tokens
|
|
61
|
+
if self.output_tokens is not None:
|
|
62
|
+
d["output_tokens"] = self.output_tokens
|
|
63
|
+
if self.cost is not None:
|
|
64
|
+
d["cost"] = self.cost
|
|
65
|
+
if self.duration_ms is not None:
|
|
66
|
+
d["duration_ms"] = self.duration_ms
|
|
67
|
+
if self.metadata:
|
|
68
|
+
d["metadata"] = self.metadata
|
|
69
|
+
if self.ended_at:
|
|
70
|
+
d["ended_at"] = self.ended_at.isoformat().replace("+00:00", "Z")
|
|
71
|
+
if self.error:
|
|
72
|
+
d["error"] = self.error
|
|
73
|
+
return d
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
@dataclass
|
|
77
|
+
class Trace:
|
|
78
|
+
id: str = field(default_factory=gen_id)
|
|
79
|
+
name: Optional[str] = None
|
|
80
|
+
input: Any = None
|
|
81
|
+
output: Any = None
|
|
82
|
+
status: TraceStatus = TraceStatus.RUNNING
|
|
83
|
+
total_tokens: int = 0
|
|
84
|
+
total_cost: float = 0.0
|
|
85
|
+
total_duration_ms: int = 0
|
|
86
|
+
metadata: dict = field(default_factory=dict)
|
|
87
|
+
started_at: Optional[datetime] = field(default_factory=utcnow)
|
|
88
|
+
ended_at: Optional[datetime] = None
|
|
89
|
+
spans: list = field(default_factory=list)
|
|
90
|
+
project_id: Optional[str] = None
|
|
91
|
+
|
|
92
|
+
def to_dict(self) -> dict:
|
|
93
|
+
d: dict[str, Any] = {
|
|
94
|
+
"id": self.id,
|
|
95
|
+
"status": self.status.value,
|
|
96
|
+
"total_tokens": self.total_tokens,
|
|
97
|
+
"total_cost": self.total_cost,
|
|
98
|
+
"total_duration_ms": self.total_duration_ms,
|
|
99
|
+
"started_at": self.started_at.isoformat().replace("+00:00", "Z") if self.started_at else None,
|
|
100
|
+
}
|
|
101
|
+
if self.name:
|
|
102
|
+
d["name"] = self.name
|
|
103
|
+
if self.input is not None:
|
|
104
|
+
d["input"] = self.input
|
|
105
|
+
if self.output is not None:
|
|
106
|
+
d["output"] = self.output
|
|
107
|
+
if self.metadata:
|
|
108
|
+
d["metadata"] = self.metadata
|
|
109
|
+
if self.ended_at:
|
|
110
|
+
d["ended_at"] = self.ended_at.isoformat().replace("+00:00", "Z")
|
|
111
|
+
if self.project_id:
|
|
112
|
+
d["project_id"] = self.project_id
|
|
113
|
+
if self.spans:
|
|
114
|
+
d["spans"] = [s.to_dict() for s in self.spans]
|
|
115
|
+
return d
|
|
@@ -0,0 +1,155 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import logging
|
|
5
|
+
import threading
|
|
6
|
+
import time
|
|
7
|
+
from typing import Any
|
|
8
|
+
|
|
9
|
+
logger = logging.getLogger("retrace")
|
|
10
|
+
|
|
11
|
+
from .config import get_config
|
|
12
|
+
from .trace import Trace
|
|
13
|
+
|
|
14
|
+
|
|
15
|
+
class WSTransport:
|
|
16
|
+
"""WebSocket transport using websocket-client (sync) for span streaming."""
|
|
17
|
+
|
|
18
|
+
def __init__(self):
|
|
19
|
+
self._ws = None
|
|
20
|
+
self._lock = threading.Lock()
|
|
21
|
+
self._connected = False
|
|
22
|
+
self._backoff = 1.0
|
|
23
|
+
|
|
24
|
+
def connect(self):
|
|
25
|
+
import websocket
|
|
26
|
+
|
|
27
|
+
cfg = get_config()
|
|
28
|
+
url = f"{cfg.ws_url}/ws/v1/stream"
|
|
29
|
+
try:
|
|
30
|
+
self._ws = websocket.create_connection(url, timeout=10)
|
|
31
|
+
# Auth
|
|
32
|
+
self._ws.send(json.dumps({"type": "auth", "api_key": cfg.api_key}))
|
|
33
|
+
resp = json.loads(self._ws.recv())
|
|
34
|
+
if resp.get("type") == "auth_ok":
|
|
35
|
+
self._connected = True
|
|
36
|
+
self._backoff = 1.0
|
|
37
|
+
else:
|
|
38
|
+
self._ws.close()
|
|
39
|
+
self._ws = None
|
|
40
|
+
except Exception as e:
|
|
41
|
+
logger.debug(f"WebSocket connection failed: {e}")
|
|
42
|
+
self._ws = None
|
|
43
|
+
self._connected = False
|
|
44
|
+
|
|
45
|
+
def _ensure_connected(self):
|
|
46
|
+
if not self._connected or self._ws is None:
|
|
47
|
+
self.connect()
|
|
48
|
+
|
|
49
|
+
def send(self, event_type: str, data: dict[str, Any]):
|
|
50
|
+
with self._lock:
|
|
51
|
+
self._ensure_connected()
|
|
52
|
+
if not self._ws:
|
|
53
|
+
return
|
|
54
|
+
try:
|
|
55
|
+
self._ws.send(json.dumps({"type": event_type, "data": data}))
|
|
56
|
+
# Handle ping
|
|
57
|
+
self._ws.settimeout(0.1)
|
|
58
|
+
try:
|
|
59
|
+
msg = self._ws.recv()
|
|
60
|
+
parsed = json.loads(msg)
|
|
61
|
+
if parsed.get("type") == "ping":
|
|
62
|
+
self._ws.send(json.dumps({"type": "pong"}))
|
|
63
|
+
except Exception as e:
|
|
64
|
+
logger.debug(f"Ping handling error: {e}")
|
|
65
|
+
self._ws.settimeout(10)
|
|
66
|
+
except Exception as e:
|
|
67
|
+
logger.debug(f"WebSocket send failed: {e}")
|
|
68
|
+
self._connected = False
|
|
69
|
+
self._ws = None
|
|
70
|
+
# Retry with backoff
|
|
71
|
+
time.sleep(min(self._backoff, 30.0))
|
|
72
|
+
self._backoff *= 2
|
|
73
|
+
|
|
74
|
+
def close(self):
|
|
75
|
+
with self._lock:
|
|
76
|
+
if self._ws:
|
|
77
|
+
try:
|
|
78
|
+
self._ws.close()
|
|
79
|
+
except Exception as e:
|
|
80
|
+
logger.debug(f"WebSocket close error: {e}")
|
|
81
|
+
self._ws = None
|
|
82
|
+
self._connected = False
|
|
83
|
+
|
|
84
|
+
|
|
85
|
+
class HTTPTransport:
|
|
86
|
+
"""HTTP fallback transport using requests."""
|
|
87
|
+
|
|
88
|
+
def __init__(self):
|
|
89
|
+
self._trace_data: dict | None = None
|
|
90
|
+
self._spans: list[dict] = []
|
|
91
|
+
self._lock = threading.Lock()
|
|
92
|
+
|
|
93
|
+
def send_trace(self, trace: Trace):
|
|
94
|
+
import requests
|
|
95
|
+
|
|
96
|
+
cfg = get_config()
|
|
97
|
+
url = f"{cfg.base_url}/api/v1/traces"
|
|
98
|
+
headers = {"x-retrace-key": cfg.api_key, "Content-Type": "application/json"}
|
|
99
|
+
try:
|
|
100
|
+
requests.post(url, json=trace.to_dict(), headers=headers, timeout=10)
|
|
101
|
+
except Exception as e:
|
|
102
|
+
logger.debug(f"HTTP send_trace failed: {e}")
|
|
103
|
+
|
|
104
|
+
def send(self, event_type: str, data: dict[str, Any]):
|
|
105
|
+
with self._lock:
|
|
106
|
+
if event_type == "trace_started":
|
|
107
|
+
self._trace_data = dict(data)
|
|
108
|
+
elif event_type in ("span_started", "span_ended"):
|
|
109
|
+
self._spans.append({"_event": event_type, **data})
|
|
110
|
+
elif event_type == "trace_ended":
|
|
111
|
+
if self._trace_data:
|
|
112
|
+
self._trace_data.update(data)
|
|
113
|
+
self.flush()
|
|
114
|
+
|
|
115
|
+
def flush(self):
|
|
116
|
+
with self._lock:
|
|
117
|
+
if not self._trace_data:
|
|
118
|
+
return
|
|
119
|
+
# Merge span_started and span_ended events into complete spans
|
|
120
|
+
merged: dict[str, dict] = {}
|
|
121
|
+
for ev in self._spans:
|
|
122
|
+
event_type = ev.pop("_event", None)
|
|
123
|
+
span_id = ev.get("id", "")
|
|
124
|
+
if event_type == "span_started":
|
|
125
|
+
merged[span_id] = dict(ev)
|
|
126
|
+
elif event_type == "span_ended" and span_id in merged:
|
|
127
|
+
merged[span_id].update(ev)
|
|
128
|
+
self._trace_data["spans"] = list(merged.values())
|
|
129
|
+
|
|
130
|
+
import requests
|
|
131
|
+
cfg = get_config()
|
|
132
|
+
url = f"{cfg.base_url}/api/v1/traces"
|
|
133
|
+
headers = {"x-retrace-key": cfg.api_key, "Content-Type": "application/json"}
|
|
134
|
+
try:
|
|
135
|
+
requests.post(url, json=self._trace_data, headers=headers, timeout=10)
|
|
136
|
+
except Exception as e:
|
|
137
|
+
logger.debug(f"HTTP flush failed: {e}")
|
|
138
|
+
self._trace_data = None
|
|
139
|
+
self._spans = []
|
|
140
|
+
|
|
141
|
+
def close(self):
|
|
142
|
+
self.flush()
|
|
143
|
+
|
|
144
|
+
|
|
145
|
+
def create_transport(mode: str = "auto") -> WSTransport | HTTPTransport:
|
|
146
|
+
if mode == "http":
|
|
147
|
+
return HTTPTransport()
|
|
148
|
+
if mode == "ws":
|
|
149
|
+
return WSTransport()
|
|
150
|
+
# Auto: try WS
|
|
151
|
+
try:
|
|
152
|
+
import websocket # noqa: F401
|
|
153
|
+
return WSTransport()
|
|
154
|
+
except ImportError:
|
|
155
|
+
return HTTPTransport()
|
|
@@ -0,0 +1,25 @@
|
|
|
1
|
+
import json
|
|
2
|
+
from datetime import datetime, timezone
|
|
3
|
+
from uuid import uuid4
|
|
4
|
+
|
|
5
|
+
|
|
6
|
+
def gen_id() -> str:
|
|
7
|
+
return str(uuid4())
|
|
8
|
+
|
|
9
|
+
|
|
10
|
+
def now_iso() -> str:
|
|
11
|
+
return utcnow().isoformat().replace("+00:00", "Z")
|
|
12
|
+
|
|
13
|
+
|
|
14
|
+
def utcnow() -> datetime:
|
|
15
|
+
return datetime.now(timezone.utc)
|
|
16
|
+
|
|
17
|
+
|
|
18
|
+
def truncate_json(obj, max_bytes: int = 10240):
|
|
19
|
+
try:
|
|
20
|
+
s = json.dumps(obj)
|
|
21
|
+
if len(s.encode()) <= max_bytes:
|
|
22
|
+
return obj
|
|
23
|
+
return json.loads(s.encode()[:max_bytes].decode(errors="ignore"))
|
|
24
|
+
except (TypeError, ValueError):
|
|
25
|
+
return str(obj)[:max_bytes]
|
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
import os
|
|
2
|
+
|
|
3
|
+
def test_default_config():
|
|
4
|
+
# Clear env and reset
|
|
5
|
+
os.environ.pop("RETRACE_API_KEY", None)
|
|
6
|
+
os.environ.pop("RETRACE_BASE_URL", None)
|
|
7
|
+
import retrace.config as cfg
|
|
8
|
+
cfg._config = None
|
|
9
|
+
from retrace.config import get_config
|
|
10
|
+
config = get_config()
|
|
11
|
+
assert config.base_url == "http://localhost:3001"
|
|
12
|
+
assert config.enabled == True
|
|
13
|
+
|
|
14
|
+
def test_configure():
|
|
15
|
+
import retrace.config as cfg
|
|
16
|
+
cfg._config = None
|
|
17
|
+
from retrace.config import configure, get_config
|
|
18
|
+
configure(api_key="rt_live_test", base_url="http://custom:3001")
|
|
19
|
+
config = get_config()
|
|
20
|
+
assert config.api_key == "rt_live_test"
|
|
21
|
+
assert config.base_url == "http://custom:3001"
|