traceroai 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.
- traceroai-0.1.0/.gitignore +45 -0
- traceroai-0.1.0/PKG-INFO +78 -0
- traceroai-0.1.0/README.md +61 -0
- traceroai-0.1.0/pyproject.toml +32 -0
- traceroai-0.1.0/test_sdk_ergonomics.py +70 -0
- traceroai-0.1.0/test_sdk_send.py +52 -0
- traceroai-0.1.0/traceroai/__init__.py +4 -0
- traceroai-0.1.0/traceroai/client.py +89 -0
- traceroai-0.1.0/traceroai/trace.py +107 -0
|
@@ -0,0 +1,45 @@
|
|
|
1
|
+
# Environment files
|
|
2
|
+
.env
|
|
3
|
+
.env.local
|
|
4
|
+
.env.*.local
|
|
5
|
+
|
|
6
|
+
# Node / Next.js
|
|
7
|
+
node_modules/
|
|
8
|
+
.next/
|
|
9
|
+
out/
|
|
10
|
+
build/
|
|
11
|
+
coverage/
|
|
12
|
+
.vercel/
|
|
13
|
+
*.tsbuildinfo
|
|
14
|
+
next-env.d.ts
|
|
15
|
+
|
|
16
|
+
# Python
|
|
17
|
+
__pycache__/
|
|
18
|
+
*.py[cod]
|
|
19
|
+
.venv/
|
|
20
|
+
venv/
|
|
21
|
+
.pytest_cache/
|
|
22
|
+
.mypy_cache/
|
|
23
|
+
.ruff_cache/
|
|
24
|
+
|
|
25
|
+
# Logs
|
|
26
|
+
*.log
|
|
27
|
+
npm-debug.log*
|
|
28
|
+
yarn-debug.log*
|
|
29
|
+
pnpm-debug.log*
|
|
30
|
+
|
|
31
|
+
# OS / editor
|
|
32
|
+
.DS_Store
|
|
33
|
+
Thumbs.db
|
|
34
|
+
.vscode/
|
|
35
|
+
.idea/
|
|
36
|
+
|
|
37
|
+
# Secrets / certificates
|
|
38
|
+
*.pem
|
|
39
|
+
*.key
|
|
40
|
+
|
|
41
|
+
.venv/
|
|
42
|
+
venv/
|
|
43
|
+
docs/reference-traces/
|
|
44
|
+
# Python packaging
|
|
45
|
+
*.egg-info/
|
traceroai-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,78 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: traceroai
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Python SDK for sending RAG traces to TraceroAI
|
|
5
|
+
Project-URL: Homepage, https://github.com/chinmai-sd-123/TraceroAI
|
|
6
|
+
Project-URL: Repository, https://github.com/chinmai-sd-123/TraceroAI
|
|
7
|
+
Author: chinmai-sd-123
|
|
8
|
+
License: MIT
|
|
9
|
+
Keywords: evaluation,llm,observability,rag,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.11
|
|
15
|
+
Requires-Dist: httpx>=0.27.0
|
|
16
|
+
Description-Content-Type: text/markdown
|
|
17
|
+
|
|
18
|
+
# TraceroAI Python SDK
|
|
19
|
+
|
|
20
|
+
Send RAG traces to [TraceroAI](https://github.com/chinmai-sd-123/TraceroAI) — a
|
|
21
|
+
RAG observability and evaluation platform. Instrument any RAG pipeline
|
|
22
|
+
(LangChain, LlamaIndex, or your own) and every answer becomes a debuggable trace.
|
|
23
|
+
|
|
24
|
+
## Install
|
|
25
|
+
|
|
26
|
+
```bash
|
|
27
|
+
pip install traceroai
|
|
28
|
+
```
|
|
29
|
+
|
|
30
|
+
## Usage
|
|
31
|
+
|
|
32
|
+
### Context manager (recommended)
|
|
33
|
+
|
|
34
|
+
Times the block and sends the trace automatically:
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
from traceroai import TraceroClient
|
|
38
|
+
|
|
39
|
+
client = TraceroClient(base_url="http://localhost:8000")
|
|
40
|
+
|
|
41
|
+
with client.trace("How long does a refund take?") as t:
|
|
42
|
+
t.log_retrieval(chunks, strategy="hybrid", config={"final_top_k": 3})
|
|
43
|
+
t.log_prompt(prompt_text, version="grounded_v1")
|
|
44
|
+
t.log_generation(answer, model="gpt-4o-mini")
|
|
45
|
+
|
|
46
|
+
print(t.trace_id)
|
|
47
|
+
```
|
|
48
|
+
|
|
49
|
+
### Decorator
|
|
50
|
+
|
|
51
|
+
For a function that returns `(answer, chunks)`:
|
|
52
|
+
|
|
53
|
+
```python
|
|
54
|
+
@client.traced(model="gpt-4o-mini", strategy="hybrid")
|
|
55
|
+
def answer(query: str):
|
|
56
|
+
chunks = retrieve(query)
|
|
57
|
+
return generate(query, chunks), chunks
|
|
58
|
+
|
|
59
|
+
answer("What is the maximum file upload size?") # traced automatically
|
|
60
|
+
```
|
|
61
|
+
|
|
62
|
+
### Low-level
|
|
63
|
+
|
|
64
|
+
```python
|
|
65
|
+
client.log_trace(
|
|
66
|
+
query={"original": question},
|
|
67
|
+
retrieval={"strategy": "hybrid", "chunks": chunks},
|
|
68
|
+
generation={"model": "gpt-4o-mini", "answer": answer},
|
|
69
|
+
)
|
|
70
|
+
```
|
|
71
|
+
|
|
72
|
+
## Authentication (multi-tenant)
|
|
73
|
+
|
|
74
|
+
Pass your project API key; the server attributes traces to your project:
|
|
75
|
+
|
|
76
|
+
```python
|
|
77
|
+
client = TraceroClient(base_url="https://api.traceroai.example", api_key="key_acme")
|
|
78
|
+
```
|
|
@@ -0,0 +1,61 @@
|
|
|
1
|
+
# TraceroAI Python SDK
|
|
2
|
+
|
|
3
|
+
Send RAG traces to [TraceroAI](https://github.com/chinmai-sd-123/TraceroAI) — a
|
|
4
|
+
RAG observability and evaluation platform. Instrument any RAG pipeline
|
|
5
|
+
(LangChain, LlamaIndex, or your own) and every answer becomes a debuggable trace.
|
|
6
|
+
|
|
7
|
+
## Install
|
|
8
|
+
|
|
9
|
+
```bash
|
|
10
|
+
pip install traceroai
|
|
11
|
+
```
|
|
12
|
+
|
|
13
|
+
## Usage
|
|
14
|
+
|
|
15
|
+
### Context manager (recommended)
|
|
16
|
+
|
|
17
|
+
Times the block and sends the trace automatically:
|
|
18
|
+
|
|
19
|
+
```python
|
|
20
|
+
from traceroai import TraceroClient
|
|
21
|
+
|
|
22
|
+
client = TraceroClient(base_url="http://localhost:8000")
|
|
23
|
+
|
|
24
|
+
with client.trace("How long does a refund take?") as t:
|
|
25
|
+
t.log_retrieval(chunks, strategy="hybrid", config={"final_top_k": 3})
|
|
26
|
+
t.log_prompt(prompt_text, version="grounded_v1")
|
|
27
|
+
t.log_generation(answer, model="gpt-4o-mini")
|
|
28
|
+
|
|
29
|
+
print(t.trace_id)
|
|
30
|
+
```
|
|
31
|
+
|
|
32
|
+
### Decorator
|
|
33
|
+
|
|
34
|
+
For a function that returns `(answer, chunks)`:
|
|
35
|
+
|
|
36
|
+
```python
|
|
37
|
+
@client.traced(model="gpt-4o-mini", strategy="hybrid")
|
|
38
|
+
def answer(query: str):
|
|
39
|
+
chunks = retrieve(query)
|
|
40
|
+
return generate(query, chunks), chunks
|
|
41
|
+
|
|
42
|
+
answer("What is the maximum file upload size?") # traced automatically
|
|
43
|
+
```
|
|
44
|
+
|
|
45
|
+
### Low-level
|
|
46
|
+
|
|
47
|
+
```python
|
|
48
|
+
client.log_trace(
|
|
49
|
+
query={"original": question},
|
|
50
|
+
retrieval={"strategy": "hybrid", "chunks": chunks},
|
|
51
|
+
generation={"model": "gpt-4o-mini", "answer": answer},
|
|
52
|
+
)
|
|
53
|
+
```
|
|
54
|
+
|
|
55
|
+
## Authentication (multi-tenant)
|
|
56
|
+
|
|
57
|
+
Pass your project API key; the server attributes traces to your project:
|
|
58
|
+
|
|
59
|
+
```python
|
|
60
|
+
client = TraceroClient(base_url="https://api.traceroai.example", api_key="key_acme")
|
|
61
|
+
```
|
|
@@ -0,0 +1,32 @@
|
|
|
1
|
+
[build-system]
|
|
2
|
+
requires = ["hatchling"]
|
|
3
|
+
build-backend = "hatchling.build"
|
|
4
|
+
|
|
5
|
+
[project]
|
|
6
|
+
name = "traceroai"
|
|
7
|
+
version = "0.1.0"
|
|
8
|
+
description = "Python SDK for sending RAG traces to TraceroAI"
|
|
9
|
+
readme = "README.md"
|
|
10
|
+
requires-python = ">=3.11"
|
|
11
|
+
license = { text = "MIT" }
|
|
12
|
+
authors = [{ name = "chinmai-sd-123" }]
|
|
13
|
+
keywords = ["rag", "observability", "evaluation", "llm", "tracing"]
|
|
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
|
+
"httpx>=0.27.0",
|
|
22
|
+
]
|
|
23
|
+
|
|
24
|
+
[project.urls]
|
|
25
|
+
Homepage = "https://github.com/chinmai-sd-123/TraceroAI"
|
|
26
|
+
Repository = "https://github.com/chinmai-sd-123/TraceroAI"
|
|
27
|
+
|
|
28
|
+
[tool.hatch.build.targets.wheel]
|
|
29
|
+
packages = ["traceroai"]
|
|
30
|
+
|
|
31
|
+
[tool.pytest.ini_options]
|
|
32
|
+
pythonpath = ["."]
|
|
@@ -0,0 +1,70 @@
|
|
|
1
|
+
"""Lightweight tests for the SDK's ergonomic tracing (context manager + decorator).
|
|
2
|
+
|
|
3
|
+
Run from sdks/python: python -m pytest test_sdk_ergonomics.py
|
|
4
|
+
No live API needed — TraceroClient.log_trace is stubbed to capture the payload.
|
|
5
|
+
"""
|
|
6
|
+
|
|
7
|
+
from uuid import uuid4
|
|
8
|
+
|
|
9
|
+
from traceroai import TraceroClient
|
|
10
|
+
|
|
11
|
+
|
|
12
|
+
def _client_capturing(captured: list[dict]) -> TraceroClient:
|
|
13
|
+
client = TraceroClient(base_url="http://test")
|
|
14
|
+
|
|
15
|
+
def fake_log_trace(**payload):
|
|
16
|
+
captured.append(payload)
|
|
17
|
+
return uuid4()
|
|
18
|
+
|
|
19
|
+
client.log_trace = fake_log_trace # type: ignore[method-assign]
|
|
20
|
+
return client
|
|
21
|
+
|
|
22
|
+
|
|
23
|
+
def test_context_manager_builds_and_sends_trace() -> None:
|
|
24
|
+
captured: list[dict] = []
|
|
25
|
+
client = _client_capturing(captured)
|
|
26
|
+
|
|
27
|
+
with client.trace("How long is a refund?") as t:
|
|
28
|
+
t.log_retrieval([{"rank": 1, "chunk_id": "c1", "text": "5 to 7 days"}], strategy="lexical")
|
|
29
|
+
t.log_generation("5 to 7 business days", model="gpt-4o-mini")
|
|
30
|
+
|
|
31
|
+
assert t.trace_id is not None
|
|
32
|
+
assert len(captured) == 1
|
|
33
|
+
payload = captured[0]
|
|
34
|
+
assert payload["query"]["original"] == "How long is a refund?"
|
|
35
|
+
assert payload["retrieval"]["strategy"] == "lexical"
|
|
36
|
+
assert payload["generation"]["answered"] is True
|
|
37
|
+
assert payload["generation"]["model"] == "gpt-4o-mini"
|
|
38
|
+
assert "total_ms" in payload["latency"]
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def test_exception_marks_trace_unanswered_and_still_sends() -> None:
|
|
42
|
+
captured: list[dict] = []
|
|
43
|
+
client = _client_capturing(captured)
|
|
44
|
+
|
|
45
|
+
try:
|
|
46
|
+
with client.trace("boom") as t:
|
|
47
|
+
t.log_generation("partial", model="gpt-4o-mini")
|
|
48
|
+
raise ValueError("pipeline failed")
|
|
49
|
+
except ValueError:
|
|
50
|
+
pass
|
|
51
|
+
|
|
52
|
+
# Trace still sent (for debugging the failure), marked unanswered.
|
|
53
|
+
assert len(captured) == 1
|
|
54
|
+
assert captured[0]["generation"]["answered"] is False
|
|
55
|
+
|
|
56
|
+
|
|
57
|
+
def test_decorator_traces_and_returns_answer() -> None:
|
|
58
|
+
captured: list[dict] = []
|
|
59
|
+
client = _client_capturing(captured)
|
|
60
|
+
|
|
61
|
+
@client.traced(model="gpt-4o-mini", strategy="lexical_top_k")
|
|
62
|
+
def answer(query: str):
|
|
63
|
+
return "an answer", [{"rank": 1, "chunk_id": "c1", "text": "ctx"}]
|
|
64
|
+
|
|
65
|
+
result = answer("a question")
|
|
66
|
+
|
|
67
|
+
assert result == "an answer" # caller gets the answer transparently
|
|
68
|
+
assert len(captured) == 1
|
|
69
|
+
assert captured[0]["query"]["original"] == "a question"
|
|
70
|
+
assert captured[0]["retrieval"]["strategy"] == "lexical_top_k"
|
|
@@ -0,0 +1,52 @@
|
|
|
1
|
+
from traceroai import TraceroClient
|
|
2
|
+
|
|
3
|
+
client = TraceroClient(base_url="http://127.0.0.1:8000")
|
|
4
|
+
|
|
5
|
+
trace_id = client.log_trace(
|
|
6
|
+
query={
|
|
7
|
+
"original": "Can admins change the workspace region themselves?",
|
|
8
|
+
"rewritten": "Can admins change the workspace region themselves?",
|
|
9
|
+
"rewrite_changed": False,
|
|
10
|
+
"rewrite_method": "rule_based_v1",
|
|
11
|
+
},
|
|
12
|
+
retrieval={
|
|
13
|
+
"strategy": "hybrid_rrf_rerank",
|
|
14
|
+
"config": {
|
|
15
|
+
"lexical_top_k": 5,
|
|
16
|
+
"dense_top_k": 5,
|
|
17
|
+
"final_top_k": 3,
|
|
18
|
+
"fusion": "rrf",
|
|
19
|
+
"reranker": "rule_based_v1",
|
|
20
|
+
},
|
|
21
|
+
"chunks": [
|
|
22
|
+
{
|
|
23
|
+
"rank": 1,
|
|
24
|
+
"chunk_id": "product_faq_2",
|
|
25
|
+
"document_id": "product_faq",
|
|
26
|
+
"document_title": "Product FAQ",
|
|
27
|
+
"section": "Can I change my workspace region?",
|
|
28
|
+
"source": "product_faq.md",
|
|
29
|
+
"final_score": 1.08,
|
|
30
|
+
"text": "Customers cannot directly change a workspace region after the workspace is created. To request a region change, customers must contact support.",
|
|
31
|
+
}
|
|
32
|
+
],
|
|
33
|
+
},
|
|
34
|
+
generation={
|
|
35
|
+
"provider": "openai",
|
|
36
|
+
"model": "gpt-4o-mini",
|
|
37
|
+
"temperature": 0,
|
|
38
|
+
"answer": "No, admins cannot change the workspace region themselves. They must contact support [1].",
|
|
39
|
+
"answered": True,
|
|
40
|
+
},
|
|
41
|
+
latency={
|
|
42
|
+
"retrieval_ms": 17,
|
|
43
|
+
"generation_ms": 1154,
|
|
44
|
+
"total_ms": 1171,
|
|
45
|
+
},
|
|
46
|
+
diagnosis={
|
|
47
|
+
"label": "healthy_answer",
|
|
48
|
+
"reason": "The retriever found useful context and the model answered the question.",
|
|
49
|
+
},
|
|
50
|
+
)
|
|
51
|
+
|
|
52
|
+
print(f"Sent trace: {trace_id}")
|
|
@@ -0,0 +1,89 @@
|
|
|
1
|
+
from typing import Any
|
|
2
|
+
from uuid import UUID
|
|
3
|
+
|
|
4
|
+
import httpx
|
|
5
|
+
|
|
6
|
+
from traceroai.trace import TraceContext
|
|
7
|
+
|
|
8
|
+
|
|
9
|
+
class TraceroClient:
|
|
10
|
+
def __init__(self, base_url: str, api_key: str| None = None, timeout_seconds: float = 10.0,)-> None:
|
|
11
|
+
self.base_url = base_url
|
|
12
|
+
self.api_key = api_key
|
|
13
|
+
self.timeout_seconds = timeout_seconds
|
|
14
|
+
|
|
15
|
+
def trace(
|
|
16
|
+
self,
|
|
17
|
+
query: str,
|
|
18
|
+
*,
|
|
19
|
+
rewritten: str | None = None,
|
|
20
|
+
project: dict[str, Any] | None = None,
|
|
21
|
+
metadata: dict[str, Any] | None = None,
|
|
22
|
+
) -> TraceContext:
|
|
23
|
+
"""Open a trace context manager that auto-times and auto-sends.
|
|
24
|
+
|
|
25
|
+
with client.trace("How long is a refund?") as t:
|
|
26
|
+
t.log_retrieval(chunks)
|
|
27
|
+
t.log_generation(answer, model="gpt-4o-mini")
|
|
28
|
+
"""
|
|
29
|
+
return TraceContext(
|
|
30
|
+
self, query, rewritten=rewritten, project=project, metadata=metadata
|
|
31
|
+
)
|
|
32
|
+
|
|
33
|
+
def traced(self, *, model: str, strategy: str = "vector"):
|
|
34
|
+
"""Decorator for a RAG function that returns (answer, chunks).
|
|
35
|
+
|
|
36
|
+
The first positional arg of the wrapped function is treated as the query.
|
|
37
|
+
|
|
38
|
+
@client.traced(model="gpt-4o-mini")
|
|
39
|
+
def answer(query: str) -> tuple[str, list[dict]]:
|
|
40
|
+
chunks = retrieve(query)
|
|
41
|
+
return generate(query, chunks), chunks
|
|
42
|
+
"""
|
|
43
|
+
|
|
44
|
+
def decorator(func):
|
|
45
|
+
def wrapper(query: str, *args: Any, **kwargs: Any):
|
|
46
|
+
with self.trace(query) as t:
|
|
47
|
+
answer, chunks = func(query, *args, **kwargs)
|
|
48
|
+
t.log_retrieval(chunks, strategy=strategy)
|
|
49
|
+
t.log_generation(answer, model=model)
|
|
50
|
+
return answer
|
|
51
|
+
|
|
52
|
+
return wrapper
|
|
53
|
+
|
|
54
|
+
return decorator
|
|
55
|
+
|
|
56
|
+
def log_trace(self , *,query:dict[str, Any],
|
|
57
|
+
retrieval: dict[str, Any],
|
|
58
|
+
generation: dict[str, Any],
|
|
59
|
+
prompt: dict[str, Any] | None = None,
|
|
60
|
+
latency: dict[str, Any] | None = None,
|
|
61
|
+
evaluation: dict[str, Any] | None = None,
|
|
62
|
+
diagnosis: dict[str, Any] | None = None,
|
|
63
|
+
project: dict[str, Any]| None = None,
|
|
64
|
+
metadata: dict[str, Any] | None = None,
|
|
65
|
+
)-> UUID:
|
|
66
|
+
payload = {
|
|
67
|
+
"query": query,
|
|
68
|
+
"retrieval": retrieval,
|
|
69
|
+
"generation": generation,
|
|
70
|
+
"prompt": prompt or {},
|
|
71
|
+
"latency": latency or {},
|
|
72
|
+
"evaluation": evaluation or {},
|
|
73
|
+
"diagnosis": diagnosis or {},
|
|
74
|
+
"project": project or {},
|
|
75
|
+
"metadata": metadata or {},
|
|
76
|
+
}
|
|
77
|
+
|
|
78
|
+
headers= {}
|
|
79
|
+
if self.api_key:
|
|
80
|
+
headers["Authorization"] = f"Bearer {self.api_key}"
|
|
81
|
+
|
|
82
|
+
response = httpx.post(f"{self.base_url}/v1/traces", json=payload, headers=headers, timeout=self.timeout_seconds,)
|
|
83
|
+
|
|
84
|
+
response.raise_for_status()
|
|
85
|
+
|
|
86
|
+
data = response.json()
|
|
87
|
+
return UUID(data["trace_id"])
|
|
88
|
+
|
|
89
|
+
|
|
@@ -0,0 +1,107 @@
|
|
|
1
|
+
"""Ergonomic tracing: a context manager that builds and sends a trace for you.
|
|
2
|
+
|
|
3
|
+
Instead of hand-assembling the log_trace(...) payload, wrap the RAG work:
|
|
4
|
+
|
|
5
|
+
with client.trace(query="How long is a refund?") as t:
|
|
6
|
+
t.log_retrieval(chunks, strategy="hybrid", config={"final_top_k": 3})
|
|
7
|
+
t.log_prompt(prompt_text, version="grounded_v1")
|
|
8
|
+
t.log_generation(answer, model="gpt-4o-mini")
|
|
9
|
+
|
|
10
|
+
The context manager times the block automatically, fills latency.total_ms, marks
|
|
11
|
+
the trace unanswered if an exception escapes, and sends it on exit.
|
|
12
|
+
"""
|
|
13
|
+
|
|
14
|
+
from __future__ import annotations
|
|
15
|
+
|
|
16
|
+
import time
|
|
17
|
+
from typing import TYPE_CHECKING, Any
|
|
18
|
+
from uuid import UUID
|
|
19
|
+
|
|
20
|
+
if TYPE_CHECKING:
|
|
21
|
+
from traceroai.client import TraceroClient
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
class TraceContext:
|
|
25
|
+
def __init__(
|
|
26
|
+
self,
|
|
27
|
+
client: "TraceroClient",
|
|
28
|
+
query: str,
|
|
29
|
+
*,
|
|
30
|
+
rewritten: str | None = None,
|
|
31
|
+
project: dict[str, Any] | None = None,
|
|
32
|
+
metadata: dict[str, Any] | None = None,
|
|
33
|
+
) -> None:
|
|
34
|
+
self._client = client
|
|
35
|
+
self._query: dict[str, Any] = {"original": query}
|
|
36
|
+
if rewritten is not None and rewritten != query:
|
|
37
|
+
self._query.update({"rewritten": rewritten, "rewrite_changed": True})
|
|
38
|
+
|
|
39
|
+
self._retrieval: dict[str, Any] = {"chunks": []}
|
|
40
|
+
self._prompt: dict[str, Any] | None = None
|
|
41
|
+
self._generation: dict[str, Any] = {"answer": "", "answered": False}
|
|
42
|
+
self._project = project
|
|
43
|
+
self._metadata = metadata
|
|
44
|
+
|
|
45
|
+
self._start: float = 0.0
|
|
46
|
+
self.trace_id: UUID | None = None
|
|
47
|
+
|
|
48
|
+
# --- piece loggers (call these inside the `with` block) ---
|
|
49
|
+
|
|
50
|
+
def log_retrieval(
|
|
51
|
+
self,
|
|
52
|
+
chunks: list[dict[str, Any]],
|
|
53
|
+
*,
|
|
54
|
+
strategy: str = "vector",
|
|
55
|
+
config: dict[str, Any] | None = None,
|
|
56
|
+
) -> None:
|
|
57
|
+
self._retrieval = {"strategy": strategy, "chunks": chunks}
|
|
58
|
+
if config is not None:
|
|
59
|
+
self._retrieval["config"] = config
|
|
60
|
+
|
|
61
|
+
def log_prompt(
|
|
62
|
+
self, content: str, *, version: str | None = None, template_name: str | None = None
|
|
63
|
+
) -> None:
|
|
64
|
+
self._prompt = {"content": content, "version": version, "template_name": template_name}
|
|
65
|
+
|
|
66
|
+
def log_generation(
|
|
67
|
+
self,
|
|
68
|
+
answer: str,
|
|
69
|
+
*,
|
|
70
|
+
model: str,
|
|
71
|
+
provider: str | None = None,
|
|
72
|
+
temperature: float | None = None,
|
|
73
|
+
) -> None:
|
|
74
|
+
self._generation = {
|
|
75
|
+
"answer": answer,
|
|
76
|
+
"answered": True,
|
|
77
|
+
"model": model,
|
|
78
|
+
"provider": provider,
|
|
79
|
+
"temperature": temperature,
|
|
80
|
+
}
|
|
81
|
+
|
|
82
|
+
# --- context manager protocol ---
|
|
83
|
+
|
|
84
|
+
def __enter__(self) -> "TraceContext":
|
|
85
|
+
self._start = time.perf_counter()
|
|
86
|
+
return self
|
|
87
|
+
|
|
88
|
+
def __exit__(self, exc_type, exc, tb) -> bool:
|
|
89
|
+
total_ms = int((time.perf_counter() - self._start) * 1000)
|
|
90
|
+
|
|
91
|
+
# A generation is only "answered" if one was logged and nothing blew up.
|
|
92
|
+
if exc_type is not None:
|
|
93
|
+
self._generation["answered"] = False
|
|
94
|
+
|
|
95
|
+
# The schema requires a model; default it so a half-finished trace still sends.
|
|
96
|
+
self._generation.setdefault("model", "unknown")
|
|
97
|
+
|
|
98
|
+
self.trace_id = self._client.log_trace(
|
|
99
|
+
query=self._query,
|
|
100
|
+
retrieval=self._retrieval,
|
|
101
|
+
generation=self._generation,
|
|
102
|
+
prompt=self._prompt,
|
|
103
|
+
latency={"total_ms": total_ms},
|
|
104
|
+
project=self._project,
|
|
105
|
+
metadata=self._metadata,
|
|
106
|
+
)
|
|
107
|
+
return False # never suppress exceptions
|