trace-ai-python 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.
- trace_ai_python-0.1.0/.gitignore +18 -0
- trace_ai_python-0.1.0/PKG-INFO +144 -0
- trace_ai_python-0.1.0/README.md +126 -0
- trace_ai_python-0.1.0/pyproject.toml +30 -0
- trace_ai_python-0.1.0/smoke_test.py +119 -0
- trace_ai_python-0.1.0/traceai/__init__.py +6 -0
- trace_ai_python-0.1.0/traceai/_cost.py +36 -0
- trace_ai_python-0.1.0/traceai/langchain.py +284 -0
- trace_ai_python-0.1.0/traceai/tracer.py +116 -0
|
@@ -0,0 +1,18 @@
|
|
|
1
|
+
# Python
|
|
2
|
+
.venv/
|
|
3
|
+
venv/
|
|
4
|
+
__pycache__/
|
|
5
|
+
*.py[cod]
|
|
6
|
+
*.egg-info/
|
|
7
|
+
.pytest_cache/
|
|
8
|
+
|
|
9
|
+
# Env
|
|
10
|
+
.env
|
|
11
|
+
.env.local
|
|
12
|
+
|
|
13
|
+
# OS
|
|
14
|
+
.DS_Store
|
|
15
|
+
|
|
16
|
+
# Database backups — full schema dumps, not versioned migrations
|
|
17
|
+
backend/migrations/trace_backup.sql
|
|
18
|
+
backend/migrations/*backup*.sql
|
|
@@ -0,0 +1,144 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: trace-ai-python
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Observability for LLM workflows — tokens, latency, cost, and anomaly detection
|
|
5
|
+
Project-URL: Homepage, https://use-trace-ai.vercel.app
|
|
6
|
+
Project-URL: Repository, https://github.com/joshuakim314/trace
|
|
7
|
+
Author-email: "trace.ai" <jjkk@umich.edu>
|
|
8
|
+
License: MIT
|
|
9
|
+
Keywords: anthropic,langchain,llm,observability,openai,tracing
|
|
10
|
+
Classifier: License :: OSI Approved :: MIT License
|
|
11
|
+
Classifier: Operating System :: OS Independent
|
|
12
|
+
Classifier: Programming Language :: Python :: 3
|
|
13
|
+
Classifier: Topic :: Software Development :: Libraries :: Python Modules
|
|
14
|
+
Requires-Python: >=3.9
|
|
15
|
+
Provides-Extra: langchain
|
|
16
|
+
Requires-Dist: langchain-core>=0.1.0; extra == 'langchain'
|
|
17
|
+
Description-Content-Type: text/markdown
|
|
18
|
+
|
|
19
|
+
# traceai
|
|
20
|
+
|
|
21
|
+
Python SDK for [trace.ai](https://use-trace-ai.vercel.app) — observability for LLM workflows.
|
|
22
|
+
|
|
23
|
+
Automatically captures tokens, latency, cost, and anomaly scores for every LLM call.
|
|
24
|
+
|
|
25
|
+
## Installation
|
|
26
|
+
|
|
27
|
+
```bash
|
|
28
|
+
pip install traceai # core — manual ingest()
|
|
29
|
+
pip install traceai[langchain] # + LangChain callback handler (Anthropic, OpenAI, etc.)
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
## LangChain (recommended)
|
|
33
|
+
|
|
34
|
+
Attach `TraceAICallbackHandler` to any LangChain LLM — every call is traced automatically:
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
from traceai import Tracer
|
|
38
|
+
from traceai.langchain import TraceAICallbackHandler
|
|
39
|
+
from langchain_anthropic import ChatAnthropic
|
|
40
|
+
from langchain_core.prompts import ChatPromptTemplate
|
|
41
|
+
from langchain_core.output_parsers import StrOutputParser
|
|
42
|
+
|
|
43
|
+
tracer = Tracer(api_key="trace_...")
|
|
44
|
+
handler = TraceAICallbackHandler(tracer)
|
|
45
|
+
|
|
46
|
+
llm = ChatAnthropic(model="claude-haiku-4-5-20251001", callbacks=[handler])
|
|
47
|
+
chain = ChatPromptTemplate.from_template("Summarize: {text}") | llm | StrOutputParser()
|
|
48
|
+
chain.invoke({"text": "..."})
|
|
49
|
+
# → shows up in your dashboard automatically
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
Works with any LangChain-compatible provider: Anthropic, OpenAI, Gemini, Cohere, and more.
|
|
53
|
+
|
|
54
|
+
## Step naming
|
|
55
|
+
|
|
56
|
+
Pass `step_name` in config metadata to label steps in the dashboard:
|
|
57
|
+
|
|
58
|
+
```python
|
|
59
|
+
chain.invoke(
|
|
60
|
+
{"text": "..."},
|
|
61
|
+
config={"metadata": {"step_name": "summarize"}}
|
|
62
|
+
)
|
|
63
|
+
```
|
|
64
|
+
|
|
65
|
+
Without a name, the step is labeled from the serialized model name (e.g. `ChatAnthropic`).
|
|
66
|
+
|
|
67
|
+
## Multi-step pipelines
|
|
68
|
+
|
|
69
|
+
Steps inside a single `chain.invoke()` are automatically grouped into one run in the dashboard. Use `RunnableLambda` to wrap multi-step workflows:
|
|
70
|
+
|
|
71
|
+
```python
|
|
72
|
+
from langchain_core.runnables import RunnableLambda
|
|
73
|
+
from langchain_core.messages import SystemMessage, HumanMessage
|
|
74
|
+
|
|
75
|
+
def pipeline(inputs, config):
|
|
76
|
+
intent = llm.invoke(
|
|
77
|
+
[SystemMessage(content="Classify as: billing, technical, general."),
|
|
78
|
+
HumanMessage(content=inputs["message"])],
|
|
79
|
+
config={**config, "metadata": {"step_name": "classify"}},
|
|
80
|
+
)
|
|
81
|
+
reply = llm.invoke(
|
|
82
|
+
[SystemMessage(content="You are a support agent. Be concise."),
|
|
83
|
+
HumanMessage(content=inputs["message"])],
|
|
84
|
+
config={**config, "metadata": {"step_name": "generate"}},
|
|
85
|
+
)
|
|
86
|
+
return reply.content
|
|
87
|
+
|
|
88
|
+
chain = RunnableLambda(pipeline)
|
|
89
|
+
chain.invoke({"message": "..."}, config={"callbacks": [handler]})
|
|
90
|
+
# → both steps appear under one run_id in the dashboard
|
|
91
|
+
```
|
|
92
|
+
|
|
93
|
+
## Manual ingest
|
|
94
|
+
|
|
95
|
+
For models outside LangChain, or to record any custom step:
|
|
96
|
+
|
|
97
|
+
```python
|
|
98
|
+
import time, json
|
|
99
|
+
|
|
100
|
+
start = time.monotonic()
|
|
101
|
+
response = my_model.generate(prompt)
|
|
102
|
+
latency = int((time.monotonic() - start) * 1000)
|
|
103
|
+
|
|
104
|
+
tracer.ingest(
|
|
105
|
+
run_id = "my-run-id",
|
|
106
|
+
step_name = "generate",
|
|
107
|
+
step_index = 0,
|
|
108
|
+
model = "my-model",
|
|
109
|
+
prompt = json.dumps({"messages": [{"role": "user", "content": prompt}]}),
|
|
110
|
+
input_tokens = response.input_tokens,
|
|
111
|
+
output_tokens = response.output_tokens,
|
|
112
|
+
total_tokens = response.total_tokens,
|
|
113
|
+
latency_ms = latency,
|
|
114
|
+
cost = 0.001,
|
|
115
|
+
status_success= True,
|
|
116
|
+
output_code = response.text,
|
|
117
|
+
)
|
|
118
|
+
```
|
|
119
|
+
|
|
120
|
+
`ingest()` fires in a background thread and never blocks your application.
|
|
121
|
+
|
|
122
|
+
## Configuration
|
|
123
|
+
|
|
124
|
+
```python
|
|
125
|
+
import os
|
|
126
|
+
from traceai import Tracer
|
|
127
|
+
|
|
128
|
+
tracer = Tracer(
|
|
129
|
+
api_key = os.environ["TRACE_API_KEY"],
|
|
130
|
+
api_url = os.environ.get("TRACE_API_URL", "https://trace-production-940c.up.railway.app"),
|
|
131
|
+
)
|
|
132
|
+
```
|
|
133
|
+
|
|
134
|
+
For local dev, add to `.env`:
|
|
135
|
+
```
|
|
136
|
+
TRACE_API_KEY=trace_...
|
|
137
|
+
TRACE_API_URL=http://localhost:8000
|
|
138
|
+
```
|
|
139
|
+
|
|
140
|
+
## Links
|
|
141
|
+
|
|
142
|
+
- [Dashboard](https://use-trace-ai.vercel.app)
|
|
143
|
+
- [Documentation](https://use-trace-ai.vercel.app/docs)
|
|
144
|
+
- [TypeScript SDK](../sdk/)
|
|
@@ -0,0 +1,126 @@
|
|
|
1
|
+
# traceai
|
|
2
|
+
|
|
3
|
+
Python SDK for [trace.ai](https://use-trace-ai.vercel.app) — observability for LLM workflows.
|
|
4
|
+
|
|
5
|
+
Automatically captures tokens, latency, cost, and anomaly scores for every LLM call.
|
|
6
|
+
|
|
7
|
+
## Installation
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install traceai # core — manual ingest()
|
|
11
|
+
pip install traceai[langchain] # + LangChain callback handler (Anthropic, OpenAI, etc.)
|
|
12
|
+
```
|
|
13
|
+
|
|
14
|
+
## LangChain (recommended)
|
|
15
|
+
|
|
16
|
+
Attach `TraceAICallbackHandler` to any LangChain LLM — every call is traced automatically:
|
|
17
|
+
|
|
18
|
+
```python
|
|
19
|
+
from traceai import Tracer
|
|
20
|
+
from traceai.langchain import TraceAICallbackHandler
|
|
21
|
+
from langchain_anthropic import ChatAnthropic
|
|
22
|
+
from langchain_core.prompts import ChatPromptTemplate
|
|
23
|
+
from langchain_core.output_parsers import StrOutputParser
|
|
24
|
+
|
|
25
|
+
tracer = Tracer(api_key="trace_...")
|
|
26
|
+
handler = TraceAICallbackHandler(tracer)
|
|
27
|
+
|
|
28
|
+
llm = ChatAnthropic(model="claude-haiku-4-5-20251001", callbacks=[handler])
|
|
29
|
+
chain = ChatPromptTemplate.from_template("Summarize: {text}") | llm | StrOutputParser()
|
|
30
|
+
chain.invoke({"text": "..."})
|
|
31
|
+
# → shows up in your dashboard automatically
|
|
32
|
+
```
|
|
33
|
+
|
|
34
|
+
Works with any LangChain-compatible provider: Anthropic, OpenAI, Gemini, Cohere, and more.
|
|
35
|
+
|
|
36
|
+
## Step naming
|
|
37
|
+
|
|
38
|
+
Pass `step_name` in config metadata to label steps in the dashboard:
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
chain.invoke(
|
|
42
|
+
{"text": "..."},
|
|
43
|
+
config={"metadata": {"step_name": "summarize"}}
|
|
44
|
+
)
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
Without a name, the step is labeled from the serialized model name (e.g. `ChatAnthropic`).
|
|
48
|
+
|
|
49
|
+
## Multi-step pipelines
|
|
50
|
+
|
|
51
|
+
Steps inside a single `chain.invoke()` are automatically grouped into one run in the dashboard. Use `RunnableLambda` to wrap multi-step workflows:
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
from langchain_core.runnables import RunnableLambda
|
|
55
|
+
from langchain_core.messages import SystemMessage, HumanMessage
|
|
56
|
+
|
|
57
|
+
def pipeline(inputs, config):
|
|
58
|
+
intent = llm.invoke(
|
|
59
|
+
[SystemMessage(content="Classify as: billing, technical, general."),
|
|
60
|
+
HumanMessage(content=inputs["message"])],
|
|
61
|
+
config={**config, "metadata": {"step_name": "classify"}},
|
|
62
|
+
)
|
|
63
|
+
reply = llm.invoke(
|
|
64
|
+
[SystemMessage(content="You are a support agent. Be concise."),
|
|
65
|
+
HumanMessage(content=inputs["message"])],
|
|
66
|
+
config={**config, "metadata": {"step_name": "generate"}},
|
|
67
|
+
)
|
|
68
|
+
return reply.content
|
|
69
|
+
|
|
70
|
+
chain = RunnableLambda(pipeline)
|
|
71
|
+
chain.invoke({"message": "..."}, config={"callbacks": [handler]})
|
|
72
|
+
# → both steps appear under one run_id in the dashboard
|
|
73
|
+
```
|
|
74
|
+
|
|
75
|
+
## Manual ingest
|
|
76
|
+
|
|
77
|
+
For models outside LangChain, or to record any custom step:
|
|
78
|
+
|
|
79
|
+
```python
|
|
80
|
+
import time, json
|
|
81
|
+
|
|
82
|
+
start = time.monotonic()
|
|
83
|
+
response = my_model.generate(prompt)
|
|
84
|
+
latency = int((time.monotonic() - start) * 1000)
|
|
85
|
+
|
|
86
|
+
tracer.ingest(
|
|
87
|
+
run_id = "my-run-id",
|
|
88
|
+
step_name = "generate",
|
|
89
|
+
step_index = 0,
|
|
90
|
+
model = "my-model",
|
|
91
|
+
prompt = json.dumps({"messages": [{"role": "user", "content": prompt}]}),
|
|
92
|
+
input_tokens = response.input_tokens,
|
|
93
|
+
output_tokens = response.output_tokens,
|
|
94
|
+
total_tokens = response.total_tokens,
|
|
95
|
+
latency_ms = latency,
|
|
96
|
+
cost = 0.001,
|
|
97
|
+
status_success= True,
|
|
98
|
+
output_code = response.text,
|
|
99
|
+
)
|
|
100
|
+
```
|
|
101
|
+
|
|
102
|
+
`ingest()` fires in a background thread and never blocks your application.
|
|
103
|
+
|
|
104
|
+
## Configuration
|
|
105
|
+
|
|
106
|
+
```python
|
|
107
|
+
import os
|
|
108
|
+
from traceai import Tracer
|
|
109
|
+
|
|
110
|
+
tracer = Tracer(
|
|
111
|
+
api_key = os.environ["TRACE_API_KEY"],
|
|
112
|
+
api_url = os.environ.get("TRACE_API_URL", "https://trace-production-940c.up.railway.app"),
|
|
113
|
+
)
|
|
114
|
+
```
|
|
115
|
+
|
|
116
|
+
For local dev, add to `.env`:
|
|
117
|
+
```
|
|
118
|
+
TRACE_API_KEY=trace_...
|
|
119
|
+
TRACE_API_URL=http://localhost:8000
|
|
120
|
+
```
|
|
121
|
+
|
|
122
|
+
## Links
|
|
123
|
+
|
|
124
|
+
- [Dashboard](https://use-trace-ai.vercel.app)
|
|
125
|
+
- [Documentation](https://use-trace-ai.vercel.app/docs)
|
|
126
|
+
- [TypeScript SDK](../sdk/)
|
|
@@ -0,0 +1,30 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "trace-ai-python"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Observability for LLM workflows — tokens, latency, cost, and anomaly detection"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.9"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [{ name = "trace.ai", email = "jjkk@umich.edu" }]
|
|
13
|
+
keywords = ["llm", "observability", "tracing", "langchain", "anthropic", "openai"]
|
|
14
|
+
classifiers = [
|
|
15
|
+
"Programming Language :: Python :: 3",
|
|
16
|
+
"License :: OSI Approved :: MIT License",
|
|
17
|
+
"Operating System :: OS Independent",
|
|
18
|
+
"Topic :: Software Development :: Libraries :: Python Modules",
|
|
19
|
+
]
|
|
20
|
+
dependencies = []
|
|
21
|
+
|
|
22
|
+
[project.optional-dependencies]
|
|
23
|
+
langchain = ["langchain-core>=0.1.0"]
|
|
24
|
+
|
|
25
|
+
[project.urls]
|
|
26
|
+
Homepage = "https://use-trace-ai.vercel.app"
|
|
27
|
+
Repository = "https://github.com/joshuakim314/trace"
|
|
28
|
+
|
|
29
|
+
[tool.hatch.build.targets.wheel]
|
|
30
|
+
packages = ["traceai"]
|
|
@@ -0,0 +1,119 @@
|
|
|
1
|
+
"""Smoke test — run with: python smoke_test.py
|
|
2
|
+
|
|
3
|
+
Tests:
|
|
4
|
+
1. Tracer.ingest() fires without error (fire-and-forget, expects a running backend)
|
|
5
|
+
2. Cost calculation
|
|
6
|
+
3. TraceAICallbackHandler with a mock LLM (no real API call)
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
import json
|
|
10
|
+
import sys
|
|
11
|
+
import uuid
|
|
12
|
+
from unittest.mock import MagicMock, patch
|
|
13
|
+
from uuid import UUID
|
|
14
|
+
|
|
15
|
+
sys.path.insert(0, ".")
|
|
16
|
+
|
|
17
|
+
from traceai import Tracer
|
|
18
|
+
from traceai._cost import get_cost
|
|
19
|
+
from traceai.langchain import (
|
|
20
|
+
TraceAICallbackHandler,
|
|
21
|
+
_extract_model,
|
|
22
|
+
_extract_tokens,
|
|
23
|
+
_serialize_messages,
|
|
24
|
+
)
|
|
25
|
+
|
|
26
|
+
# ── 1. Cost ──────────────────────────────────────────────────────────────────
|
|
27
|
+
|
|
28
|
+
assert get_cost("claude-haiku-4-5-20251001", 1_000_000, 1_000_000) == 4.8, "cost calc failed"
|
|
29
|
+
assert get_cost("gpt-4o", 500_000, 500_000) == 6.25, "cost calc failed"
|
|
30
|
+
assert get_cost("unknown-model", 1000, 1000) == 0.0, "unknown model should be 0"
|
|
31
|
+
print("✓ cost")
|
|
32
|
+
|
|
33
|
+
# ── 2. Token extraction ───────────────────────────────────────────────────────
|
|
34
|
+
|
|
35
|
+
anthropic_output = {"usage": {"input_tokens": 10, "output_tokens": 5}}
|
|
36
|
+
assert _extract_tokens(anthropic_output) == (10, 5), "anthropic token extraction failed"
|
|
37
|
+
|
|
38
|
+
openai_output = {"token_usage": {"prompt_tokens": 20, "completion_tokens": 8}}
|
|
39
|
+
assert _extract_tokens(openai_output) == (20, 8), "openai token extraction failed"
|
|
40
|
+
print("✓ token extraction")
|
|
41
|
+
|
|
42
|
+
# ── 3. Model extraction ───────────────────────────────────────────────────────
|
|
43
|
+
|
|
44
|
+
assert _extract_model({"model": "claude-haiku-4-5-20251001"}, {}) == "claude-haiku-4-5-20251001"
|
|
45
|
+
assert _extract_model({}, {"kwargs": {"model": "gpt-4o"}}) == "gpt-4o"
|
|
46
|
+
assert _extract_model({}, {"name": "ChatAnthropic"}) == "ChatAnthropic"
|
|
47
|
+
print("✓ model extraction")
|
|
48
|
+
|
|
49
|
+
# ── 4. Message serialization ──────────────────────────────────────────────────
|
|
50
|
+
|
|
51
|
+
msg1 = MagicMock()
|
|
52
|
+
msg1.type = "system"
|
|
53
|
+
msg1.content = "You are a helpful assistant."
|
|
54
|
+
|
|
55
|
+
msg2 = MagicMock()
|
|
56
|
+
msg2.type = "human"
|
|
57
|
+
msg2.content = "What is 2+2?"
|
|
58
|
+
|
|
59
|
+
serialized = json.loads(_serialize_messages([[msg1, msg2]]))
|
|
60
|
+
assert serialized["messages"][0] == {"role": "system", "content": "You are a helpful assistant."}
|
|
61
|
+
assert serialized["messages"][1] == {"role": "user", "content": "What is 2+2?"}
|
|
62
|
+
print("✓ message serialization")
|
|
63
|
+
|
|
64
|
+
# ── 5. Callback handler — ingest payload shape ────────────────────────────────
|
|
65
|
+
|
|
66
|
+
tracer = Tracer(api_key="trace_test", api_url="http://localhost:9999")
|
|
67
|
+
handler = TraceAICallbackHandler(tracer)
|
|
68
|
+
|
|
69
|
+
ingested = []
|
|
70
|
+
tracer.ingest = lambda **kw: ingested.append(kw) # capture instead of posting
|
|
71
|
+
|
|
72
|
+
run_id = UUID("11111111-1111-1111-1111-111111111111")
|
|
73
|
+
parent_id = UUID("22222222-2222-2222-2222-222222222222")
|
|
74
|
+
serialized_llm = {"name": "ChatAnthropic", "kwargs": {"model": "claude-haiku-4-5-20251001"}}
|
|
75
|
+
|
|
76
|
+
# Simulate on_chat_model_start
|
|
77
|
+
handler.on_chat_model_start(
|
|
78
|
+
serialized_llm, [[msg1, msg2]],
|
|
79
|
+
run_id=run_id, parent_run_id=parent_id, metadata={"step_name": "classify"},
|
|
80
|
+
)
|
|
81
|
+
|
|
82
|
+
# Simulate on_llm_end
|
|
83
|
+
from langchain_core.outputs import LLMResult, ChatGeneration
|
|
84
|
+
from langchain_core.messages import AIMessage
|
|
85
|
+
|
|
86
|
+
ai_msg = AIMessage(content="The answer is 4.")
|
|
87
|
+
gen = ChatGeneration(message=ai_msg)
|
|
88
|
+
result = LLMResult(
|
|
89
|
+
generations=[[gen]],
|
|
90
|
+
llm_output={"model": "claude-haiku-4-5-20251001", "usage": {"input_tokens": 12, "output_tokens": 7}},
|
|
91
|
+
)
|
|
92
|
+
handler.on_llm_end(result, run_id=run_id, parent_run_id=parent_id)
|
|
93
|
+
|
|
94
|
+
assert len(ingested) == 1, f"expected 1 ingest call, got {len(ingested)}"
|
|
95
|
+
payload = ingested[0]
|
|
96
|
+
assert payload["run_id"] == str(parent_id), f"run_id should be parent: {payload['run_id']}"
|
|
97
|
+
assert payload["step_name"] == "classify", f"step_name: {payload['step_name']}"
|
|
98
|
+
assert payload["step_index"] == 0
|
|
99
|
+
assert payload["model"] == "claude-haiku-4-5-20251001"
|
|
100
|
+
assert payload["input_tokens"] == 12
|
|
101
|
+
assert payload["output_tokens"] == 7
|
|
102
|
+
assert payload["total_tokens"] == 19
|
|
103
|
+
assert payload["status_success"] is True
|
|
104
|
+
assert payload["output_code"] == "The answer is 4."
|
|
105
|
+
print("✓ callback handler — ingest payload")
|
|
106
|
+
|
|
107
|
+
# ── 6. Error path ─────────────────────────────────────────────────────────────
|
|
108
|
+
|
|
109
|
+
run_id2 = UUID("33333333-3333-3333-3333-333333333333")
|
|
110
|
+
handler.on_chat_model_start(serialized_llm, [[msg1]], run_id=run_id2, parent_run_id=None)
|
|
111
|
+
handler.on_llm_error(ValueError("rate limit"), run_id=run_id2, parent_run_id=None)
|
|
112
|
+
|
|
113
|
+
err_payload = ingested[-1]
|
|
114
|
+
assert err_payload["run_id"] == str(run_id2) # no parent → use own run_id
|
|
115
|
+
assert err_payload["status_success"] is False
|
|
116
|
+
assert "rate limit" in err_payload["error"]
|
|
117
|
+
print("✓ callback handler — error path")
|
|
118
|
+
|
|
119
|
+
print("\nAll smoke tests passed.")
|
|
@@ -0,0 +1,36 @@
|
|
|
1
|
+
"""Pricing table for cost calculation — mirrors sdk/src/cost.ts."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
_PRICING: dict[str, tuple[float, float]] = {
|
|
6
|
+
# (input_per_1m_usd, output_per_1m_usd)
|
|
7
|
+
# Anthropic
|
|
8
|
+
"claude-opus-4-8": (15.0, 75.0),
|
|
9
|
+
"claude-opus-4-8-20251101": (15.0, 75.0),
|
|
10
|
+
"claude-sonnet-4-6": (3.0, 15.0),
|
|
11
|
+
"claude-sonnet-4-6-20251001": (3.0, 15.0),
|
|
12
|
+
"claude-haiku-4-5": (0.8, 4.0),
|
|
13
|
+
"claude-haiku-4-5-20251001": (0.8, 4.0),
|
|
14
|
+
"claude-3-5-sonnet-20241022": (3.0, 15.0),
|
|
15
|
+
"claude-3-5-haiku-20241022": (0.8, 4.0),
|
|
16
|
+
"claude-3-opus-20240229": (15.0, 75.0),
|
|
17
|
+
# OpenAI
|
|
18
|
+
"gpt-4o": (2.5, 10.0),
|
|
19
|
+
"gpt-4o-2024-11-20": (2.5, 10.0),
|
|
20
|
+
"gpt-4o-mini": (0.15, 0.6),
|
|
21
|
+
"gpt-4o-mini-2024-07-18": (0.15, 0.6),
|
|
22
|
+
"gpt-4-turbo": (10.0, 30.0),
|
|
23
|
+
"gpt-4": (30.0, 60.0),
|
|
24
|
+
"gpt-3.5-turbo": (0.5, 1.5),
|
|
25
|
+
"o1": (15.0, 60.0),
|
|
26
|
+
"o1-mini": (3.0, 12.0),
|
|
27
|
+
"o3-mini": (1.1, 4.4),
|
|
28
|
+
}
|
|
29
|
+
|
|
30
|
+
|
|
31
|
+
def get_cost(model: str, input_tokens: int, output_tokens: int) -> float:
|
|
32
|
+
pricing = _PRICING.get(model)
|
|
33
|
+
if not pricing:
|
|
34
|
+
return 0.0
|
|
35
|
+
input_per_1m, output_per_1m = pricing
|
|
36
|
+
return (input_tokens / 1_000_000) * input_per_1m + (output_tokens / 1_000_000) * output_per_1m
|
|
@@ -0,0 +1,284 @@
|
|
|
1
|
+
"""LangChain callback handler for trace.ai.
|
|
2
|
+
|
|
3
|
+
Attach to any LangChain LLM or chain — every LLM call is automatically traced:
|
|
4
|
+
|
|
5
|
+
from traceai import Tracer
|
|
6
|
+
from traceai.langchain import TraceAICallbackHandler
|
|
7
|
+
|
|
8
|
+
tracer = Tracer(api_key="trace_...")
|
|
9
|
+
handler = TraceAICallbackHandler(tracer)
|
|
10
|
+
|
|
11
|
+
llm = ChatAnthropic(model="claude-haiku-4-5-20251001", callbacks=[handler])
|
|
12
|
+
chain = prompt | llm | StrOutputParser()
|
|
13
|
+
chain.invoke({"topic": "AI safety"})
|
|
14
|
+
|
|
15
|
+
Run grouping
|
|
16
|
+
------------
|
|
17
|
+
LangChain passes a `run_id` (UUID) to each LLM call and a `parent_run_id` for
|
|
18
|
+
the chain that contains it. We use the immediate parent as the trace.ai run_id so
|
|
19
|
+
all LLM calls inside a single chain.invoke() share one run in the dashboard.
|
|
20
|
+
|
|
21
|
+
Step naming
|
|
22
|
+
-----------
|
|
23
|
+
Priority order:
|
|
24
|
+
1. metadata["step_name"] passed in invoke() / run_config
|
|
25
|
+
2. serialized["name"] (e.g. "ChatAnthropic", "ChatOpenAI")
|
|
26
|
+
3. "llm_call"
|
|
27
|
+
|
|
28
|
+
Thread safety
|
|
29
|
+
-------------
|
|
30
|
+
The handler can be shared across concurrent requests (threaded Flask, HTTPServer,
|
|
31
|
+
etc.). All per-call state is protected by a single RLock.
|
|
32
|
+
"""
|
|
33
|
+
|
|
34
|
+
from __future__ import annotations
|
|
35
|
+
|
|
36
|
+
import json
|
|
37
|
+
import threading
|
|
38
|
+
import time
|
|
39
|
+
from typing import Any
|
|
40
|
+
from uuid import UUID
|
|
41
|
+
|
|
42
|
+
try:
|
|
43
|
+
from langchain_core.callbacks import BaseCallbackHandler
|
|
44
|
+
from langchain_core.messages import BaseMessage
|
|
45
|
+
from langchain_core.outputs import LLMResult
|
|
46
|
+
except ImportError as e:
|
|
47
|
+
raise ImportError(
|
|
48
|
+
"langchain-core is required: pip install traceai[langchain]"
|
|
49
|
+
) from e
|
|
50
|
+
|
|
51
|
+
from ._cost import get_cost
|
|
52
|
+
from .tracer import Tracer
|
|
53
|
+
|
|
54
|
+
|
|
55
|
+
def _extract_tokens_anthropic(llm_output: dict) -> tuple[int, int]:
|
|
56
|
+
usage = llm_output.get("usage", {})
|
|
57
|
+
inp = usage.get("input_tokens") or usage.get("prompt_tokens") or 0
|
|
58
|
+
out = usage.get("output_tokens") or usage.get("completion_tokens") or 0
|
|
59
|
+
return int(inp), int(out)
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _extract_tokens_openai(llm_output: dict) -> tuple[int, int]:
|
|
63
|
+
usage = llm_output.get("token_usage", {})
|
|
64
|
+
inp = usage.get("prompt_tokens") or 0
|
|
65
|
+
out = usage.get("completion_tokens") or 0
|
|
66
|
+
return int(inp), int(out)
|
|
67
|
+
|
|
68
|
+
|
|
69
|
+
def _extract_tokens(llm_output: dict) -> tuple[int, int]:
|
|
70
|
+
inp, out = _extract_tokens_anthropic(llm_output)
|
|
71
|
+
if inp or out:
|
|
72
|
+
return inp, out
|
|
73
|
+
return _extract_tokens_openai(llm_output)
|
|
74
|
+
|
|
75
|
+
|
|
76
|
+
def _extract_model(llm_output: dict, serialized: dict) -> str:
|
|
77
|
+
return (
|
|
78
|
+
llm_output.get("model")
|
|
79
|
+
or llm_output.get("model_name")
|
|
80
|
+
or llm_output.get("model_id")
|
|
81
|
+
or (serialized.get("kwargs") or {}).get("model")
|
|
82
|
+
or (serialized.get("kwargs") or {}).get("model_name")
|
|
83
|
+
or serialized.get("name", "unknown")
|
|
84
|
+
)
|
|
85
|
+
|
|
86
|
+
|
|
87
|
+
def _serialize_messages(messages: list[list[BaseMessage]]) -> str:
|
|
88
|
+
out = []
|
|
89
|
+
for batch in messages:
|
|
90
|
+
for msg in batch:
|
|
91
|
+
role = getattr(msg, "type", "unknown")
|
|
92
|
+
role = {"human": "user", "ai": "assistant", "system": "system"}.get(role, role)
|
|
93
|
+
content = msg.content if isinstance(msg.content, str) else json.dumps(msg.content)
|
|
94
|
+
out.append({"role": role, "content": content})
|
|
95
|
+
return json.dumps({"messages": out})
|
|
96
|
+
|
|
97
|
+
|
|
98
|
+
def _extract_output(response: LLMResult) -> str | None:
|
|
99
|
+
try:
|
|
100
|
+
gen = response.generations[0][0]
|
|
101
|
+
if hasattr(gen, "message"):
|
|
102
|
+
content = gen.message.content
|
|
103
|
+
if isinstance(content, str):
|
|
104
|
+
return content
|
|
105
|
+
if isinstance(content, list):
|
|
106
|
+
return " ".join(
|
|
107
|
+
b.get("text", "") for b in content if isinstance(b, dict) and b.get("type") == "text"
|
|
108
|
+
)
|
|
109
|
+
return getattr(gen, "text", None)
|
|
110
|
+
except (IndexError, AttributeError):
|
|
111
|
+
return None
|
|
112
|
+
|
|
113
|
+
|
|
114
|
+
class TraceAICallbackHandler(BaseCallbackHandler):
|
|
115
|
+
"""Attach to any LangChain LLM or chain to automatically trace every call."""
|
|
116
|
+
|
|
117
|
+
def __init__(self, tracer: Tracer) -> None:
|
|
118
|
+
super().__init__()
|
|
119
|
+
self.tracer = tracer
|
|
120
|
+
self._lock = threading.RLock()
|
|
121
|
+
|
|
122
|
+
# run_id (LangChain UUID) → wall-clock start time
|
|
123
|
+
self._start_times: dict[UUID, float] = {}
|
|
124
|
+
# run_id → serialized dict (for model name extraction in on_llm_end)
|
|
125
|
+
self._serialized: dict[UUID, dict] = {}
|
|
126
|
+
# run_id → prompt string
|
|
127
|
+
self._prompts: dict[UUID, str] = {}
|
|
128
|
+
# run_id → step_name
|
|
129
|
+
self._step_names: dict[UUID, str] = {}
|
|
130
|
+
# trace_run_id (str) → step counter
|
|
131
|
+
self._step_counters: dict[str, int] = {}
|
|
132
|
+
|
|
133
|
+
# ── Helpers ───────────────────────────────────────────────────────────────
|
|
134
|
+
|
|
135
|
+
def _trace_run_id(self, lc_run_id: UUID, parent_run_id: UUID | None) -> str:
|
|
136
|
+
"""Map LangChain's run hierarchy to a trace.ai run_id.
|
|
137
|
+
|
|
138
|
+
The immediate parent (chain's run_id) becomes the trace.ai run_id so
|
|
139
|
+
all LLM calls inside one chain.invoke() share a single run.
|
|
140
|
+
If there's no parent (bare LLM call), the LLM's own run_id is used.
|
|
141
|
+
"""
|
|
142
|
+
return str(parent_run_id) if parent_run_id else str(lc_run_id)
|
|
143
|
+
|
|
144
|
+
def _next_step_index(self, trace_run_id: str) -> int:
|
|
145
|
+
with self._lock:
|
|
146
|
+
idx = self._step_counters.get(trace_run_id, 0)
|
|
147
|
+
self._step_counters[trace_run_id] = idx + 1
|
|
148
|
+
return idx
|
|
149
|
+
|
|
150
|
+
def _step_name(self, run_id: UUID, serialized: dict, metadata: dict | None) -> str:
|
|
151
|
+
if metadata and metadata.get("step_name"):
|
|
152
|
+
return str(metadata["step_name"])
|
|
153
|
+
return serialized.get("name") or "llm_call"
|
|
154
|
+
|
|
155
|
+
def _pop_start(self, run_id: UUID) -> float | None:
|
|
156
|
+
with self._lock:
|
|
157
|
+
return self._start_times.pop(run_id, None)
|
|
158
|
+
|
|
159
|
+
def _pop_state(self, run_id: UUID) -> tuple[dict, str, str]:
|
|
160
|
+
with self._lock:
|
|
161
|
+
serialized = self._serialized.pop(run_id, {})
|
|
162
|
+
prompt = self._prompts.pop(run_id, "")
|
|
163
|
+
step_name = self._step_names.pop(run_id, "llm_call")
|
|
164
|
+
return serialized, prompt, step_name
|
|
165
|
+
|
|
166
|
+
# ── LangChain callbacks ───────────────────────────────────────────────────
|
|
167
|
+
|
|
168
|
+
def on_chat_model_start(
|
|
169
|
+
self,
|
|
170
|
+
serialized: dict[str, Any],
|
|
171
|
+
messages: list[list[BaseMessage]],
|
|
172
|
+
*,
|
|
173
|
+
run_id: UUID,
|
|
174
|
+
parent_run_id: UUID | None = None,
|
|
175
|
+
metadata: dict[str, Any] | None = None,
|
|
176
|
+
**kwargs: Any,
|
|
177
|
+
) -> None:
|
|
178
|
+
with self._lock:
|
|
179
|
+
self._start_times[run_id] = time.monotonic()
|
|
180
|
+
self._serialized[run_id] = serialized
|
|
181
|
+
self._prompts[run_id] = _serialize_messages(messages)
|
|
182
|
+
self._step_names[run_id] = self._step_name(run_id, serialized, metadata)
|
|
183
|
+
|
|
184
|
+
def on_llm_start(
|
|
185
|
+
self,
|
|
186
|
+
serialized: dict[str, Any],
|
|
187
|
+
prompts: list[str],
|
|
188
|
+
*,
|
|
189
|
+
run_id: UUID,
|
|
190
|
+
parent_run_id: UUID | None = None,
|
|
191
|
+
metadata: dict[str, Any] | None = None,
|
|
192
|
+
**kwargs: Any,
|
|
193
|
+
) -> None:
|
|
194
|
+
with self._lock:
|
|
195
|
+
self._start_times[run_id] = time.monotonic()
|
|
196
|
+
self._serialized[run_id] = serialized
|
|
197
|
+
self._prompts[run_id] = json.dumps({"messages": [{"role": "user", "content": p} for p in prompts]})
|
|
198
|
+
self._step_names[run_id] = self._step_name(run_id, serialized, metadata)
|
|
199
|
+
|
|
200
|
+
def on_llm_end(
|
|
201
|
+
self,
|
|
202
|
+
response: LLMResult,
|
|
203
|
+
*,
|
|
204
|
+
run_id: UUID,
|
|
205
|
+
parent_run_id: UUID | None = None,
|
|
206
|
+
**kwargs: Any,
|
|
207
|
+
) -> None:
|
|
208
|
+
start = self._pop_start(run_id)
|
|
209
|
+
latency_ms = int((time.monotonic() - start) * 1000) if start is not None else 0
|
|
210
|
+
serialized, prompt, step_name = self._pop_state(run_id)
|
|
211
|
+
|
|
212
|
+
llm_output = response.llm_output or {}
|
|
213
|
+
input_tok, output_tok = _extract_tokens(llm_output)
|
|
214
|
+
total_tok = input_tok + output_tok
|
|
215
|
+
model = _extract_model(llm_output, serialized)
|
|
216
|
+
cost = get_cost(model, input_tok, output_tok)
|
|
217
|
+
output = _extract_output(response)
|
|
218
|
+
|
|
219
|
+
trace_run_id = self._trace_run_id(run_id, parent_run_id)
|
|
220
|
+
step_index = self._next_step_index(trace_run_id)
|
|
221
|
+
span_id = str(run_id)
|
|
222
|
+
parent_span_id = str(parent_run_id) if parent_run_id and str(parent_run_id) != trace_run_id else None
|
|
223
|
+
|
|
224
|
+
self.tracer.ingest(
|
|
225
|
+
run_id=trace_run_id,
|
|
226
|
+
step_name=step_name,
|
|
227
|
+
step_index=step_index,
|
|
228
|
+
model=model,
|
|
229
|
+
prompt=prompt,
|
|
230
|
+
input_tokens=input_tok,
|
|
231
|
+
output_tokens=output_tok,
|
|
232
|
+
total_tokens=total_tok,
|
|
233
|
+
latency_ms=latency_ms,
|
|
234
|
+
cost=cost,
|
|
235
|
+
status_success=True,
|
|
236
|
+
output_code=output,
|
|
237
|
+
span_id=span_id,
|
|
238
|
+
parent_span_id=parent_span_id,
|
|
239
|
+
)
|
|
240
|
+
|
|
241
|
+
def on_llm_error(
|
|
242
|
+
self,
|
|
243
|
+
error: BaseException,
|
|
244
|
+
*,
|
|
245
|
+
run_id: UUID,
|
|
246
|
+
parent_run_id: UUID | None = None,
|
|
247
|
+
**kwargs: Any,
|
|
248
|
+
) -> None:
|
|
249
|
+
start = self._pop_start(run_id)
|
|
250
|
+
latency_ms = int((time.monotonic() - start) * 1000) if start is not None else 0
|
|
251
|
+
serialized, prompt, step_name = self._pop_state(run_id)
|
|
252
|
+
model = _extract_model({}, serialized)
|
|
253
|
+
|
|
254
|
+
trace_run_id = self._trace_run_id(run_id, parent_run_id)
|
|
255
|
+
step_index = self._next_step_index(trace_run_id)
|
|
256
|
+
span_id = str(run_id)
|
|
257
|
+
parent_span_id = str(parent_run_id) if parent_run_id and str(parent_run_id) != trace_run_id else None
|
|
258
|
+
|
|
259
|
+
self.tracer.ingest(
|
|
260
|
+
run_id=trace_run_id,
|
|
261
|
+
step_name=step_name,
|
|
262
|
+
step_index=step_index,
|
|
263
|
+
model=model,
|
|
264
|
+
prompt=prompt,
|
|
265
|
+
input_tokens=0,
|
|
266
|
+
output_tokens=0,
|
|
267
|
+
total_tokens=0,
|
|
268
|
+
latency_ms=latency_ms,
|
|
269
|
+
cost=0.0,
|
|
270
|
+
status_success=False,
|
|
271
|
+
error=str(error),
|
|
272
|
+
span_id=span_id,
|
|
273
|
+
parent_span_id=parent_span_id,
|
|
274
|
+
)
|
|
275
|
+
|
|
276
|
+
def on_chain_end(
|
|
277
|
+
self,
|
|
278
|
+
outputs: dict[str, Any],
|
|
279
|
+
*,
|
|
280
|
+
run_id: UUID,
|
|
281
|
+
**kwargs: Any,
|
|
282
|
+
) -> None:
|
|
283
|
+
with self._lock:
|
|
284
|
+
self._step_counters.pop(str(run_id), None)
|
|
@@ -0,0 +1,116 @@
|
|
|
1
|
+
"""Core Tracer — fire-and-forget ingest + run context management."""
|
|
2
|
+
|
|
3
|
+
from __future__ import annotations
|
|
4
|
+
|
|
5
|
+
import json
|
|
6
|
+
import threading
|
|
7
|
+
import uuid as _uuid
|
|
8
|
+
from contextlib import contextmanager
|
|
9
|
+
from contextvars import ContextVar
|
|
10
|
+
from typing import Any, Generator
|
|
11
|
+
from urllib import request as _urllib_request
|
|
12
|
+
|
|
13
|
+
_DEFAULT_URL = "https://trace-production-940c.up.railway.app"
|
|
14
|
+
|
|
15
|
+
# ContextVar so run_id propagates automatically across async/threaded code
|
|
16
|
+
_active_run_id: ContextVar[str | None] = ContextVar("traceai_run_id", default=None)
|
|
17
|
+
_active_step_index: ContextVar[int] = ContextVar("traceai_step_index", default=0)
|
|
18
|
+
|
|
19
|
+
|
|
20
|
+
def _new_uuid() -> str:
|
|
21
|
+
return str(_uuid.uuid4())
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class Tracer:
|
|
25
|
+
"""
|
|
26
|
+
trace.ai Python client.
|
|
27
|
+
|
|
28
|
+
Usage::
|
|
29
|
+
|
|
30
|
+
tracer = Tracer(api_key="trace_...")
|
|
31
|
+
|
|
32
|
+
# Manual ingest (any framework)
|
|
33
|
+
tracer.ingest(
|
|
34
|
+
run_id="my-run",
|
|
35
|
+
step_name="classify",
|
|
36
|
+
step_index=0,
|
|
37
|
+
model="claude-haiku-4-5-20251001",
|
|
38
|
+
prompt=json.dumps({"messages": [...]}),
|
|
39
|
+
input_tokens=12,
|
|
40
|
+
output_tokens=4,
|
|
41
|
+
total_tokens=16,
|
|
42
|
+
latency_ms=84,
|
|
43
|
+
status_success=True,
|
|
44
|
+
output_code="billing",
|
|
45
|
+
)
|
|
46
|
+
|
|
47
|
+
# LangChain — see traceai.langchain.TraceAICallbackHandler
|
|
48
|
+
"""
|
|
49
|
+
|
|
50
|
+
def __init__(self, api_key: str, api_url: str = "") -> None:
|
|
51
|
+
self.api_key = api_key
|
|
52
|
+
# Empty string falls back to default so that os.environ.get("TRACE_API_URL", "")
|
|
53
|
+
# behaves the same as not passing api_url at all.
|
|
54
|
+
self.api_url = (api_url or _DEFAULT_URL).rstrip("/")
|
|
55
|
+
|
|
56
|
+
# ── Ingest ────────────────────────────────────────────────────────────────
|
|
57
|
+
|
|
58
|
+
def ingest(self, **fields: Any) -> None:
|
|
59
|
+
"""Fire-and-forget POST to /ingest. Never raises — failures are silent."""
|
|
60
|
+
threading.Thread(target=self._post, args=(fields,), daemon=True).start()
|
|
61
|
+
|
|
62
|
+
def _post(self, payload: dict[str, Any]) -> None:
|
|
63
|
+
try:
|
|
64
|
+
data = json.dumps(payload).encode()
|
|
65
|
+
req = _urllib_request.Request(
|
|
66
|
+
f"{self.api_url}/ingest",
|
|
67
|
+
data=data,
|
|
68
|
+
headers={
|
|
69
|
+
"Content-Type": "application/json",
|
|
70
|
+
"Authorization": f"Bearer {self.api_key}",
|
|
71
|
+
},
|
|
72
|
+
method="POST",
|
|
73
|
+
)
|
|
74
|
+
_urllib_request.urlopen(req, timeout=10)
|
|
75
|
+
except Exception:
|
|
76
|
+
pass # never block the application
|
|
77
|
+
|
|
78
|
+
# ── Run context ───────────────────────────────────────────────────────────
|
|
79
|
+
|
|
80
|
+
@contextmanager
|
|
81
|
+
def run(self, run_id: str | None = None) -> Generator["RunContext", None, None]:
|
|
82
|
+
"""Context manager that sets a run ID for the duration of a block.
|
|
83
|
+
|
|
84
|
+
Use this when you're not using LangChain and want to group manual
|
|
85
|
+
ingest() calls into a single run::
|
|
86
|
+
|
|
87
|
+
with tracer.run() as run:
|
|
88
|
+
tracer.ingest(run_id=run.run_id, step_name="step1", ...)
|
|
89
|
+
tracer.ingest(run_id=run.run_id, step_name="step2", ...)
|
|
90
|
+
"""
|
|
91
|
+
rid = run_id or _new_uuid()
|
|
92
|
+
token_id = _active_run_id.set(rid)
|
|
93
|
+
token_idx = _active_step_index.set(0)
|
|
94
|
+
ctx = RunContext(run_id=rid)
|
|
95
|
+
try:
|
|
96
|
+
yield ctx
|
|
97
|
+
finally:
|
|
98
|
+
_active_run_id.reset(token_id)
|
|
99
|
+
_active_step_index.reset(token_idx)
|
|
100
|
+
|
|
101
|
+
# ── Helpers for handlers ──────────────────────────────────────────────────
|
|
102
|
+
|
|
103
|
+
def get_active_run_id(self) -> str | None:
|
|
104
|
+
return _active_run_id.get()
|
|
105
|
+
|
|
106
|
+
def next_step_index(self) -> int:
|
|
107
|
+
idx = _active_step_index.get()
|
|
108
|
+
_active_step_index.set(idx + 1)
|
|
109
|
+
return idx
|
|
110
|
+
|
|
111
|
+
|
|
112
|
+
class RunContext:
|
|
113
|
+
"""Returned by Tracer.run() — holds the run_id for the current block."""
|
|
114
|
+
|
|
115
|
+
def __init__(self, run_id: str) -> None:
|
|
116
|
+
self.run_id = run_id
|