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.
@@ -0,0 +1,16 @@
1
+ node_modules/
2
+ dist/
3
+ .env
4
+ .env.local
5
+ .env.staging
6
+ .env.production
7
+ *.pyc
8
+ __pycache__/
9
+ .turbo/
10
+ .next/
11
+ .firecrawl/
12
+ .vercel/
13
+ .fly/
14
+ .source/
15
+ coverage/
16
+ .husky/_
@@ -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,3 @@
1
+ from .gemini import install_gemini_interceptor, uninstall_gemini_interceptor
2
+
3
+ __all__ = ["install_gemini_interceptor", "uninstall_gemini_interceptor"]
@@ -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"