telemetry-dev-anthropic 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.
|
@@ -0,0 +1,101 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: telemetry-dev-anthropic
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: Anthropic integration for telemetry.dev Python SDK
|
|
5
|
+
Keywords: telemetry,opentelemetry,anthropic,claude,llm,genai,tracing
|
|
6
|
+
Author: telemetry.dev
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
Classifier: Development Status :: 4 - Beta
|
|
9
|
+
Classifier: Intended Audience :: Developers
|
|
10
|
+
Classifier: Programming Language :: Python :: 3
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.10
|
|
12
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
13
|
+
Classifier: Programming Language :: Python :: 3.12
|
|
14
|
+
Classifier: Programming Language :: Python :: 3.13
|
|
15
|
+
Classifier: Typing :: Typed
|
|
16
|
+
Requires-Dist: telemetry-dev>=0.1.0
|
|
17
|
+
Requires-Dist: anthropic>=0.116,<1
|
|
18
|
+
Requires-Python: >=3.10
|
|
19
|
+
Project-URL: Homepage, https://telemetry.dev
|
|
20
|
+
Project-URL: Repository, https://github.com/telemetry-dev/telemetry.dev
|
|
21
|
+
Description-Content-Type: text/markdown
|
|
22
|
+
|
|
23
|
+
# telemetry-dev-anthropic
|
|
24
|
+
|
|
25
|
+
Anthropic Claude SDK instrumentation for telemetry.dev. It wraps the official `anthropic` Python SDK and emits telemetry.dev generation spans through `telemetry-dev`.
|
|
26
|
+
|
|
27
|
+
## Install
|
|
28
|
+
|
|
29
|
+
```sh
|
|
30
|
+
pip install telemetry-dev-anthropic
|
|
31
|
+
```
|
|
32
|
+
|
|
33
|
+
Initialize the core SDK first:
|
|
34
|
+
|
|
35
|
+
```python
|
|
36
|
+
import telemetry_dev
|
|
37
|
+
|
|
38
|
+
telemetry_dev.init(
|
|
39
|
+
api_key="td_live_...",
|
|
40
|
+
service_name="my-service",
|
|
41
|
+
)
|
|
42
|
+
```
|
|
43
|
+
|
|
44
|
+
## Per-client wrapping
|
|
45
|
+
|
|
46
|
+
```python
|
|
47
|
+
from anthropic import Anthropic
|
|
48
|
+
from telemetry_dev_anthropic import wrap_anthropic
|
|
49
|
+
|
|
50
|
+
client = wrap_anthropic(Anthropic())
|
|
51
|
+
message = client.messages.create(
|
|
52
|
+
model="claude-sonnet-4-6",
|
|
53
|
+
max_tokens=1024,
|
|
54
|
+
messages=[{"role": "user", "content": "Tell me a joke about OpenTelemetry"}],
|
|
55
|
+
)
|
|
56
|
+
```
|
|
57
|
+
|
|
58
|
+
`wrap_anthropic` also supports `AsyncAnthropic`, `AnthropicBedrock`, `AsyncAnthropicBedrock`, `AnthropicVertex`, and `AsyncAnthropicVertex` clients.
|
|
59
|
+
|
|
60
|
+
## Global instrumentation
|
|
61
|
+
|
|
62
|
+
```python
|
|
63
|
+
from anthropic import Anthropic
|
|
64
|
+
from telemetry_dev_anthropic import instrument_anthropic, uninstrument_anthropic
|
|
65
|
+
|
|
66
|
+
instrument_anthropic()
|
|
67
|
+
try:
|
|
68
|
+
client = Anthropic()
|
|
69
|
+
client.messages.create(
|
|
70
|
+
model="claude-sonnet-4-6",
|
|
71
|
+
max_tokens=1024,
|
|
72
|
+
messages=[{"role": "user", "content": "Hello"}],
|
|
73
|
+
)
|
|
74
|
+
finally:
|
|
75
|
+
uninstrument_anthropic()
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
## Instrumented surfaces
|
|
79
|
+
|
|
80
|
+
- `client.messages.create(...)`
|
|
81
|
+
- `client.messages.create(..., stream=True)`
|
|
82
|
+
- `client.messages.stream(...)` context managers, sync and async
|
|
83
|
+
|
|
84
|
+
The integration maps native Anthropic request and response shapes directly into telemetry.dev fields. It does not normalize messages into another schema.
|
|
85
|
+
|
|
86
|
+
## Streaming
|
|
87
|
+
|
|
88
|
+
Native Anthropic stream events pass through unmodified. The span records time to first chunk on the first received event, merges usage from `message_start` and `message_delta`, aggregates text, tool-use JSON, and thinking blocks, and ends on stream exhaustion, close, context-manager exit, or error.
|
|
89
|
+
|
|
90
|
+
`messages.stream()` starts the span when the context manager is entered, because that is when the Anthropic SDK opens the HTTP stream.
|
|
91
|
+
|
|
92
|
+
## Bedrock and Vertex
|
|
93
|
+
|
|
94
|
+
Class instrumentation covers Bedrock and Vertex clients because the Anthropic SDK reuses the same `Messages` and `AsyncMessages` resource classes. Provider attribution is recorded as `aws.bedrock` or `gcp.vertex_ai` when the client class identifies those runtimes.
|
|
95
|
+
|
|
96
|
+
## Limitations
|
|
97
|
+
|
|
98
|
+
- `client.beta.messages` is not instrumented.
|
|
99
|
+
- `messages.parse()` and `messages.count_tokens()` are not instrumented.
|
|
100
|
+
- `with_raw_response` snapshots bound methods; wrap or instrument clients before creating raw-response wrappers.
|
|
101
|
+
- Unconsumed streams end spans only on exhaustion, close, context-manager exit, or error.
|
|
@@ -0,0 +1,79 @@
|
|
|
1
|
+
# telemetry-dev-anthropic
|
|
2
|
+
|
|
3
|
+
Anthropic Claude SDK instrumentation for telemetry.dev. It wraps the official `anthropic` Python SDK and emits telemetry.dev generation spans through `telemetry-dev`.
|
|
4
|
+
|
|
5
|
+
## Install
|
|
6
|
+
|
|
7
|
+
```sh
|
|
8
|
+
pip install telemetry-dev-anthropic
|
|
9
|
+
```
|
|
10
|
+
|
|
11
|
+
Initialize the core SDK first:
|
|
12
|
+
|
|
13
|
+
```python
|
|
14
|
+
import telemetry_dev
|
|
15
|
+
|
|
16
|
+
telemetry_dev.init(
|
|
17
|
+
api_key="td_live_...",
|
|
18
|
+
service_name="my-service",
|
|
19
|
+
)
|
|
20
|
+
```
|
|
21
|
+
|
|
22
|
+
## Per-client wrapping
|
|
23
|
+
|
|
24
|
+
```python
|
|
25
|
+
from anthropic import Anthropic
|
|
26
|
+
from telemetry_dev_anthropic import wrap_anthropic
|
|
27
|
+
|
|
28
|
+
client = wrap_anthropic(Anthropic())
|
|
29
|
+
message = client.messages.create(
|
|
30
|
+
model="claude-sonnet-4-6",
|
|
31
|
+
max_tokens=1024,
|
|
32
|
+
messages=[{"role": "user", "content": "Tell me a joke about OpenTelemetry"}],
|
|
33
|
+
)
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
`wrap_anthropic` also supports `AsyncAnthropic`, `AnthropicBedrock`, `AsyncAnthropicBedrock`, `AnthropicVertex`, and `AsyncAnthropicVertex` clients.
|
|
37
|
+
|
|
38
|
+
## Global instrumentation
|
|
39
|
+
|
|
40
|
+
```python
|
|
41
|
+
from anthropic import Anthropic
|
|
42
|
+
from telemetry_dev_anthropic import instrument_anthropic, uninstrument_anthropic
|
|
43
|
+
|
|
44
|
+
instrument_anthropic()
|
|
45
|
+
try:
|
|
46
|
+
client = Anthropic()
|
|
47
|
+
client.messages.create(
|
|
48
|
+
model="claude-sonnet-4-6",
|
|
49
|
+
max_tokens=1024,
|
|
50
|
+
messages=[{"role": "user", "content": "Hello"}],
|
|
51
|
+
)
|
|
52
|
+
finally:
|
|
53
|
+
uninstrument_anthropic()
|
|
54
|
+
```
|
|
55
|
+
|
|
56
|
+
## Instrumented surfaces
|
|
57
|
+
|
|
58
|
+
- `client.messages.create(...)`
|
|
59
|
+
- `client.messages.create(..., stream=True)`
|
|
60
|
+
- `client.messages.stream(...)` context managers, sync and async
|
|
61
|
+
|
|
62
|
+
The integration maps native Anthropic request and response shapes directly into telemetry.dev fields. It does not normalize messages into another schema.
|
|
63
|
+
|
|
64
|
+
## Streaming
|
|
65
|
+
|
|
66
|
+
Native Anthropic stream events pass through unmodified. The span records time to first chunk on the first received event, merges usage from `message_start` and `message_delta`, aggregates text, tool-use JSON, and thinking blocks, and ends on stream exhaustion, close, context-manager exit, or error.
|
|
67
|
+
|
|
68
|
+
`messages.stream()` starts the span when the context manager is entered, because that is when the Anthropic SDK opens the HTTP stream.
|
|
69
|
+
|
|
70
|
+
## Bedrock and Vertex
|
|
71
|
+
|
|
72
|
+
Class instrumentation covers Bedrock and Vertex clients because the Anthropic SDK reuses the same `Messages` and `AsyncMessages` resource classes. Provider attribution is recorded as `aws.bedrock` or `gcp.vertex_ai` when the client class identifies those runtimes.
|
|
73
|
+
|
|
74
|
+
## Limitations
|
|
75
|
+
|
|
76
|
+
- `client.beta.messages` is not instrumented.
|
|
77
|
+
- `messages.parse()` and `messages.count_tokens()` are not instrumented.
|
|
78
|
+
- `with_raw_response` snapshots bound methods; wrap or instrument clients before creating raw-response wrappers.
|
|
79
|
+
- Unconsumed streams end spans only on exhaustion, close, context-manager exit, or error.
|
|
@@ -0,0 +1,58 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "telemetry-dev-anthropic"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "Anthropic integration for telemetry.dev Python SDK"
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
requires-python = ">=3.10"
|
|
8
|
+
authors = [{ name = "telemetry.dev" }]
|
|
9
|
+
keywords = ["telemetry", "opentelemetry", "anthropic", "claude", "llm", "genai", "tracing"]
|
|
10
|
+
classifiers = [
|
|
11
|
+
"Development Status :: 4 - Beta",
|
|
12
|
+
"Intended Audience :: Developers",
|
|
13
|
+
"Programming Language :: Python :: 3",
|
|
14
|
+
"Programming Language :: Python :: 3.10",
|
|
15
|
+
"Programming Language :: Python :: 3.11",
|
|
16
|
+
"Programming Language :: Python :: 3.12",
|
|
17
|
+
"Programming Language :: Python :: 3.13",
|
|
18
|
+
"Typing :: Typed",
|
|
19
|
+
]
|
|
20
|
+
dependencies = [
|
|
21
|
+
"telemetry-dev>=0.1.0",
|
|
22
|
+
"anthropic>=0.116,<1",
|
|
23
|
+
]
|
|
24
|
+
|
|
25
|
+
[project.urls]
|
|
26
|
+
Homepage = "https://telemetry.dev"
|
|
27
|
+
Repository = "https://github.com/telemetry-dev/telemetry.dev"
|
|
28
|
+
|
|
29
|
+
[tool.uv.sources]
|
|
30
|
+
telemetry-dev = { path = "../python", editable = true }
|
|
31
|
+
|
|
32
|
+
[dependency-groups]
|
|
33
|
+
dev = [
|
|
34
|
+
"pytest>=8.3",
|
|
35
|
+
"pytest-asyncio>=0.25",
|
|
36
|
+
"ruff>=0.9",
|
|
37
|
+
"pyright>=1.1.390",
|
|
38
|
+
]
|
|
39
|
+
|
|
40
|
+
[build-system]
|
|
41
|
+
requires = ["uv_build>=0.9.0,<0.10.0"]
|
|
42
|
+
build-backend = "uv_build"
|
|
43
|
+
|
|
44
|
+
[tool.pytest.ini_options]
|
|
45
|
+
asyncio_mode = "auto"
|
|
46
|
+
testpaths = ["tests"]
|
|
47
|
+
|
|
48
|
+
[tool.ruff]
|
|
49
|
+
line-length = 100
|
|
50
|
+
target-version = "py310"
|
|
51
|
+
|
|
52
|
+
[tool.ruff.lint]
|
|
53
|
+
select = ["E", "F", "I", "UP", "B", "RUF"]
|
|
54
|
+
|
|
55
|
+
[tool.pyright]
|
|
56
|
+
include = ["src", "tests"]
|
|
57
|
+
typeCheckingMode = "strict"
|
|
58
|
+
pythonVersion = "3.10"
|
|
@@ -0,0 +1,780 @@
|
|
|
1
|
+
from __future__ import annotations
|
|
2
|
+
|
|
3
|
+
import json
|
|
4
|
+
import threading
|
|
5
|
+
import time
|
|
6
|
+
from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence
|
|
7
|
+
from functools import wraps
|
|
8
|
+
from typing import Any, TypeVar, cast
|
|
9
|
+
|
|
10
|
+
import anthropic
|
|
11
|
+
import telemetry_dev
|
|
12
|
+
from anthropic.resources.messages import AsyncMessages, Messages
|
|
13
|
+
|
|
14
|
+
__version__ = "0.1.0"
|
|
15
|
+
|
|
16
|
+
ProviderResolver = Callable[[object | None], str]
|
|
17
|
+
RequestMapper = Callable[[Mapping[str, Any]], tuple[str, dict[str, Any]]]
|
|
18
|
+
ResponseMapper = Callable[[Any], dict[str, Any]]
|
|
19
|
+
|
|
20
|
+
_WRAPPED_ATTR = "_telemetry_dev_anthropic_wrapped"
|
|
21
|
+
_ORIGINAL_ATTR = "_telemetry_dev_anthropic_original"
|
|
22
|
+
_ORIGINALS: list[tuple[type[Any], str, Any]] = []
|
|
23
|
+
_installed = False
|
|
24
|
+
_install_lock = threading.Lock()
|
|
25
|
+
_T = TypeVar("_T")
|
|
26
|
+
|
|
27
|
+
|
|
28
|
+
def _field(value: Any, name: str) -> Any:
|
|
29
|
+
if isinstance(value, Mapping):
|
|
30
|
+
mapping = cast(Mapping[str, Any], value)
|
|
31
|
+
return mapping.get(name)
|
|
32
|
+
return getattr(value, name, None)
|
|
33
|
+
|
|
34
|
+
|
|
35
|
+
def _sequence_items(value: Any) -> list[Any]:
|
|
36
|
+
if isinstance(value, Sequence) and not isinstance(value, str | bytes | bytearray):
|
|
37
|
+
return list(cast(Sequence[Any], value))
|
|
38
|
+
return []
|
|
39
|
+
|
|
40
|
+
|
|
41
|
+
def _request_iterable(value: Any) -> Any:
|
|
42
|
+
if isinstance(value, str | bytes | bytearray):
|
|
43
|
+
return value
|
|
44
|
+
if callable(getattr(value, "model_dump", None)):
|
|
45
|
+
return value
|
|
46
|
+
if isinstance(value, Mapping):
|
|
47
|
+
mapping = cast(Mapping[Any, Any], value)
|
|
48
|
+
return {str(key): _request_iterable(item) for key, item in mapping.items()}
|
|
49
|
+
if isinstance(value, Iterable):
|
|
50
|
+
return [_request_iterable(item) for item in cast(Iterable[Any], value)]
|
|
51
|
+
return value
|
|
52
|
+
|
|
53
|
+
|
|
54
|
+
def _normalize_request_params(params: Mapping[str, Any]) -> dict[str, Any]:
|
|
55
|
+
normalized = dict(params)
|
|
56
|
+
for key in ("messages", "system", "tools"):
|
|
57
|
+
if key in normalized:
|
|
58
|
+
normalized[key] = _request_iterable(normalized[key])
|
|
59
|
+
return normalized
|
|
60
|
+
|
|
61
|
+
|
|
62
|
+
def _raw_response_requested(params: Mapping[str, Any]) -> bool:
|
|
63
|
+
extra_headers = params.get("extra_headers")
|
|
64
|
+
if not isinstance(extra_headers, Mapping):
|
|
65
|
+
return False
|
|
66
|
+
headers = cast(Mapping[str, Any], extra_headers)
|
|
67
|
+
return headers.get("X-Stainless-Raw-Response") in {"true", "raw", "stream"}
|
|
68
|
+
|
|
69
|
+
|
|
70
|
+
def _native(value: Any) -> Any:
|
|
71
|
+
model_dump = getattr(value, "model_dump", None)
|
|
72
|
+
if callable(model_dump):
|
|
73
|
+
return model_dump(mode="json", exclude_none=True)
|
|
74
|
+
if isinstance(value, Mapping):
|
|
75
|
+
mapping = cast(Mapping[Any, Any], value)
|
|
76
|
+
return {str(key): _native(item) for key, item in mapping.items() if item is not None}
|
|
77
|
+
sequence = _sequence_items(value)
|
|
78
|
+
if sequence:
|
|
79
|
+
return [_native(item) for item in sequence]
|
|
80
|
+
return value
|
|
81
|
+
|
|
82
|
+
|
|
83
|
+
def _number(value: Any) -> int | float | None:
|
|
84
|
+
if isinstance(value, bool):
|
|
85
|
+
return None
|
|
86
|
+
if isinstance(value, int | float):
|
|
87
|
+
return value
|
|
88
|
+
return None
|
|
89
|
+
|
|
90
|
+
|
|
91
|
+
def _string(value: Any) -> str | None:
|
|
92
|
+
return value if isinstance(value, str) else None
|
|
93
|
+
|
|
94
|
+
|
|
95
|
+
def _usage(fields: dict[str, int | float | None]) -> dict[str, int | float] | None:
|
|
96
|
+
usage = {key: value for key, value in fields.items() if value is not None}
|
|
97
|
+
return usage or None
|
|
98
|
+
|
|
99
|
+
|
|
100
|
+
def _merge_usage(
|
|
101
|
+
current: dict[str, int | float] | None, incoming: dict[str, int | float] | None
|
|
102
|
+
) -> dict[str, int | float] | None:
|
|
103
|
+
if incoming is None:
|
|
104
|
+
return current
|
|
105
|
+
merged = dict(current or {})
|
|
106
|
+
merged.update(incoming)
|
|
107
|
+
return merged or None
|
|
108
|
+
|
|
109
|
+
|
|
110
|
+
def _stop_sequences(value: Any) -> list[str] | None:
|
|
111
|
+
if isinstance(value, str):
|
|
112
|
+
return [value]
|
|
113
|
+
strings = [item for item in _sequence_items(value) if isinstance(item, str)]
|
|
114
|
+
return strings or None
|
|
115
|
+
|
|
116
|
+
|
|
117
|
+
def _messages_request(params: Mapping[str, Any]) -> tuple[str, dict[str, Any]]:
|
|
118
|
+
model = _string(params.get("model"))
|
|
119
|
+
input_value: Any
|
|
120
|
+
if "tools" in params or "tool_choice" in params:
|
|
121
|
+
input_value = {"messages": _native(params.get("messages"))}
|
|
122
|
+
tools = _native(params.get("tools"))
|
|
123
|
+
tool_choice = _native(params.get("tool_choice"))
|
|
124
|
+
if tools is not None:
|
|
125
|
+
input_value["tools"] = tools
|
|
126
|
+
if tool_choice is not None:
|
|
127
|
+
input_value["tool_choice"] = tool_choice
|
|
128
|
+
else:
|
|
129
|
+
input_value = _native(params.get("messages"))
|
|
130
|
+
return (
|
|
131
|
+
f"chat {model or 'unknown'}",
|
|
132
|
+
{
|
|
133
|
+
"type": "generation",
|
|
134
|
+
"model": model,
|
|
135
|
+
"input": input_value,
|
|
136
|
+
"system_instructions": _native(params.get("system")),
|
|
137
|
+
"temperature": _number(params.get("temperature")),
|
|
138
|
+
"top_p": _number(params.get("top_p")),
|
|
139
|
+
"top_k": _number(params.get("top_k")),
|
|
140
|
+
"max_tokens": _number(params.get("max_tokens")),
|
|
141
|
+
"stop_sequences": _stop_sequences(params.get("stop_sequences")),
|
|
142
|
+
},
|
|
143
|
+
)
|
|
144
|
+
|
|
145
|
+
|
|
146
|
+
def _messages_usage(raw: Any) -> dict[str, int | float] | None:
|
|
147
|
+
output_details = _field(raw, "output_tokens_details")
|
|
148
|
+
return _usage(
|
|
149
|
+
{
|
|
150
|
+
"input_tokens": _number(_field(raw, "input_tokens")),
|
|
151
|
+
"output_tokens": _number(_field(raw, "output_tokens")),
|
|
152
|
+
"cache_read_input_tokens": _number(_field(raw, "cache_read_input_tokens")),
|
|
153
|
+
"cache_creation_input_tokens": _number(_field(raw, "cache_creation_input_tokens")),
|
|
154
|
+
"reasoning_output_tokens": _number(_field(output_details, "thinking_tokens")),
|
|
155
|
+
}
|
|
156
|
+
)
|
|
157
|
+
|
|
158
|
+
|
|
159
|
+
def _messages_response(response: Any) -> dict[str, Any]:
|
|
160
|
+
content = _field(response, "content")
|
|
161
|
+
role = _string(_field(response, "role")) or "assistant"
|
|
162
|
+
return {
|
|
163
|
+
"response_model": _string(_field(response, "model")),
|
|
164
|
+
"response_id": _string(_field(response, "id")),
|
|
165
|
+
"finish_reason": _string(_field(response, "stop_reason")),
|
|
166
|
+
"output": [{"role": role, "content": _native(content)}] if content is not None else None,
|
|
167
|
+
"usage": _messages_usage(_field(response, "usage")),
|
|
168
|
+
}
|
|
169
|
+
|
|
170
|
+
|
|
171
|
+
def _is_anthropic_class(client: object | None, name: str) -> bool:
|
|
172
|
+
cls = getattr(anthropic, name, None)
|
|
173
|
+
return isinstance(cls, type) and isinstance(client, cls)
|
|
174
|
+
|
|
175
|
+
|
|
176
|
+
def _provider_for_client(client: object | None) -> str:
|
|
177
|
+
if _is_anthropic_class(client, "AnthropicBedrock") or _is_anthropic_class(
|
|
178
|
+
client, "AsyncAnthropicBedrock"
|
|
179
|
+
):
|
|
180
|
+
return "aws.bedrock"
|
|
181
|
+
if _is_anthropic_class(client, "AnthropicVertex") or _is_anthropic_class(
|
|
182
|
+
client, "AsyncAnthropicVertex"
|
|
183
|
+
):
|
|
184
|
+
return "gcp.vertex_ai"
|
|
185
|
+
return "anthropic"
|
|
186
|
+
|
|
187
|
+
|
|
188
|
+
def _provider_for_resource(resource: object | None) -> str:
|
|
189
|
+
return _provider_for_client(getattr(resource, "_client", None))
|
|
190
|
+
|
|
191
|
+
|
|
192
|
+
def _clean_fields(fields: Mapping[str, Any]) -> dict[str, Any]:
|
|
193
|
+
return {key: value for key, value in fields.items() if value is not None}
|
|
194
|
+
|
|
195
|
+
|
|
196
|
+
def _end_once(handle: telemetry_dev.SpanHandle) -> Callable[..., None]:
|
|
197
|
+
ended = False
|
|
198
|
+
|
|
199
|
+
def end(**fields: Any) -> None:
|
|
200
|
+
nonlocal ended
|
|
201
|
+
if ended:
|
|
202
|
+
return
|
|
203
|
+
ended = True
|
|
204
|
+
handle.end(**_clean_fields(fields))
|
|
205
|
+
|
|
206
|
+
return end
|
|
207
|
+
|
|
208
|
+
|
|
209
|
+
class _StreamState:
|
|
210
|
+
def __init__(self) -> None:
|
|
211
|
+
self.blocks: dict[int, dict[str, Any]] = {}
|
|
212
|
+
self.tool_json: dict[int, str] = {}
|
|
213
|
+
self.usage: dict[str, int | float] | None = None
|
|
214
|
+
self.finish_reason: str | None = None
|
|
215
|
+
|
|
216
|
+
|
|
217
|
+
def _parse_tool_input(raw: str) -> Any:
|
|
218
|
+
if raw == "":
|
|
219
|
+
return {}
|
|
220
|
+
try:
|
|
221
|
+
return json.loads(raw)
|
|
222
|
+
except ValueError:
|
|
223
|
+
return raw
|
|
224
|
+
|
|
225
|
+
|
|
226
|
+
def _finalize_block(index: int, block: Mapping[str, Any], state: _StreamState) -> dict[str, Any]:
|
|
227
|
+
raw_input = state.tool_json.get(index)
|
|
228
|
+
if raw_input is None:
|
|
229
|
+
return dict(block)
|
|
230
|
+
finalized = dict(block)
|
|
231
|
+
if raw_input == "" and "input" in finalized:
|
|
232
|
+
return finalized
|
|
233
|
+
parsed_input = _parse_tool_input(raw_input)
|
|
234
|
+
existing_input = finalized.get("input")
|
|
235
|
+
if isinstance(existing_input, Mapping) and isinstance(parsed_input, Mapping):
|
|
236
|
+
finalized["input"] = {**existing_input, **parsed_input}
|
|
237
|
+
else:
|
|
238
|
+
finalized["input"] = parsed_input
|
|
239
|
+
return finalized
|
|
240
|
+
|
|
241
|
+
|
|
242
|
+
def _stream_output(state: _StreamState) -> list[dict[str, Any]] | None:
|
|
243
|
+
if not state.blocks:
|
|
244
|
+
return None
|
|
245
|
+
return [
|
|
246
|
+
{
|
|
247
|
+
"role": "assistant",
|
|
248
|
+
"content": [
|
|
249
|
+
_finalize_block(index, block, state)
|
|
250
|
+
for index, block in sorted(state.blocks.items())
|
|
251
|
+
],
|
|
252
|
+
}
|
|
253
|
+
]
|
|
254
|
+
|
|
255
|
+
|
|
256
|
+
def _stream_partial(state: _StreamState) -> dict[str, Any]:
|
|
257
|
+
return {
|
|
258
|
+
"output": _stream_output(state),
|
|
259
|
+
"usage": state.usage,
|
|
260
|
+
"finish_reason": state.finish_reason,
|
|
261
|
+
}
|
|
262
|
+
|
|
263
|
+
|
|
264
|
+
def _append_string(target: dict[str, Any], key: str, value: Any) -> None:
|
|
265
|
+
text = _string(value)
|
|
266
|
+
if text is None:
|
|
267
|
+
return
|
|
268
|
+
target[key] = f"{_string(target.get(key)) or ''}{text}"
|
|
269
|
+
|
|
270
|
+
|
|
271
|
+
def _append_item(target: dict[str, Any], key: str, value: Any) -> None:
|
|
272
|
+
if value is None:
|
|
273
|
+
return
|
|
274
|
+
existing = target.get(key)
|
|
275
|
+
items: list[Any] = (
|
|
276
|
+
list(cast(Sequence[Any], existing))
|
|
277
|
+
if isinstance(existing, Sequence) and not isinstance(existing, str)
|
|
278
|
+
else []
|
|
279
|
+
)
|
|
280
|
+
items.append(_native(value))
|
|
281
|
+
target[key] = items
|
|
282
|
+
|
|
283
|
+
|
|
284
|
+
def _record_content_block_start(event: Any, state: _StreamState) -> None:
|
|
285
|
+
index = _field(event, "index")
|
|
286
|
+
block_index = index if isinstance(index, int) else len(state.blocks)
|
|
287
|
+
content_block = _field(event, "content_block")
|
|
288
|
+
block_type = _string(_field(content_block, "type"))
|
|
289
|
+
if block_type == "text":
|
|
290
|
+
state.blocks[block_index] = {
|
|
291
|
+
"type": "text",
|
|
292
|
+
"text": _string(_field(content_block, "text")) or "",
|
|
293
|
+
}
|
|
294
|
+
return
|
|
295
|
+
if block_type in {"tool_use", "server_tool_use"}:
|
|
296
|
+
native = _native(content_block)
|
|
297
|
+
state.blocks[block_index] = native if isinstance(native, dict) else {"type": block_type}
|
|
298
|
+
state.tool_json[block_index] = ""
|
|
299
|
+
return
|
|
300
|
+
if block_type == "thinking":
|
|
301
|
+
state.blocks[block_index] = {
|
|
302
|
+
"type": "thinking",
|
|
303
|
+
"thinking": _string(_field(content_block, "thinking")) or "",
|
|
304
|
+
}
|
|
305
|
+
return
|
|
306
|
+
native = _native(content_block)
|
|
307
|
+
state.blocks[block_index] = native if isinstance(native, dict) else {"type": block_type}
|
|
308
|
+
|
|
309
|
+
|
|
310
|
+
def _block_for_delta(index: int, delta_type: str | None, state: _StreamState) -> dict[str, Any]:
|
|
311
|
+
if index in state.blocks:
|
|
312
|
+
return state.blocks[index]
|
|
313
|
+
if delta_type == "input_json_delta":
|
|
314
|
+
state.blocks[index] = {"type": "tool_use"}
|
|
315
|
+
state.tool_json[index] = ""
|
|
316
|
+
elif delta_type == "thinking_delta":
|
|
317
|
+
state.blocks[index] = {"type": "thinking", "thinking": ""}
|
|
318
|
+
else:
|
|
319
|
+
state.blocks[index] = {"type": "text", "text": ""}
|
|
320
|
+
return state.blocks[index]
|
|
321
|
+
|
|
322
|
+
|
|
323
|
+
def _record_content_block_delta(event: Any, state: _StreamState) -> None:
|
|
324
|
+
index = _field(event, "index")
|
|
325
|
+
block_index = index if isinstance(index, int) else 0
|
|
326
|
+
delta = _field(event, "delta")
|
|
327
|
+
delta_type = _string(_field(delta, "type"))
|
|
328
|
+
block = _block_for_delta(block_index, delta_type, state)
|
|
329
|
+
if delta_type == "input_json_delta":
|
|
330
|
+
state.tool_json[block_index] = (
|
|
331
|
+
f"{state.tool_json.get(block_index, '')}{_string(_field(delta, 'partial_json')) or ''}"
|
|
332
|
+
)
|
|
333
|
+
elif delta_type == "text_delta":
|
|
334
|
+
_append_string(block, "text", _field(delta, "text"))
|
|
335
|
+
elif delta_type == "citations_delta":
|
|
336
|
+
_append_item(block, "citations", _field(delta, "citation"))
|
|
337
|
+
elif delta_type == "thinking_delta":
|
|
338
|
+
_append_string(block, "thinking", _field(delta, "thinking"))
|
|
339
|
+
elif delta_type == "signature_delta" and _field(delta, "signature") is not None:
|
|
340
|
+
block["signature"] = _field(delta, "signature")
|
|
341
|
+
|
|
342
|
+
|
|
343
|
+
def _record_stream_event(event: Any, state: _StreamState) -> dict[str, Any]:
|
|
344
|
+
event_type = _string(_field(event, "type"))
|
|
345
|
+
update: dict[str, Any] = {}
|
|
346
|
+
if event_type == "message_start":
|
|
347
|
+
message = _field(event, "message")
|
|
348
|
+
update["response_id"] = _string(_field(message, "id"))
|
|
349
|
+
update["response_model"] = _string(_field(message, "model"))
|
|
350
|
+
state.usage = _merge_usage(state.usage, _messages_usage(_field(message, "usage")))
|
|
351
|
+
elif event_type == "content_block_start":
|
|
352
|
+
_record_content_block_start(event, state)
|
|
353
|
+
elif event_type == "content_block_delta":
|
|
354
|
+
_record_content_block_delta(event, state)
|
|
355
|
+
elif event_type == "message_delta":
|
|
356
|
+
delta = _field(event, "delta")
|
|
357
|
+
state.finish_reason = _string(_field(delta, "stop_reason")) or state.finish_reason
|
|
358
|
+
state.usage = _merge_usage(state.usage, _messages_usage(_field(event, "usage")))
|
|
359
|
+
return update
|
|
360
|
+
|
|
361
|
+
|
|
362
|
+
class _InstrumentedStream:
|
|
363
|
+
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
364
|
+
self._inner = inner
|
|
365
|
+
self._handle = handle
|
|
366
|
+
self._end = _end_once(handle)
|
|
367
|
+
self._started_at = started_at
|
|
368
|
+
self._state = _StreamState()
|
|
369
|
+
self._saw_first = False
|
|
370
|
+
self._consume: Iterator[Any] | None = None
|
|
371
|
+
|
|
372
|
+
def _update_first(self, update: Mapping[str, Any]) -> None:
|
|
373
|
+
if self._saw_first:
|
|
374
|
+
return
|
|
375
|
+
self._saw_first = True
|
|
376
|
+
self._handle.update(
|
|
377
|
+
time_to_first_chunk_ms=(time.perf_counter() - self._started_at) * 1000,
|
|
378
|
+
response_id=update.get("response_id"),
|
|
379
|
+
response_model=update.get("response_model"),
|
|
380
|
+
)
|
|
381
|
+
|
|
382
|
+
def finish(self, error: BaseException | None = None) -> None:
|
|
383
|
+
fields = _stream_partial(self._state)
|
|
384
|
+
if error is not None:
|
|
385
|
+
fields["error"] = error
|
|
386
|
+
self._end(**fields)
|
|
387
|
+
|
|
388
|
+
def _iterate(self) -> Iterator[Any]:
|
|
389
|
+
try:
|
|
390
|
+
while True:
|
|
391
|
+
try:
|
|
392
|
+
event = next(self._inner)
|
|
393
|
+
except StopIteration:
|
|
394
|
+
break
|
|
395
|
+
except BaseException as exc:
|
|
396
|
+
self.finish(exc)
|
|
397
|
+
raise
|
|
398
|
+
update = _record_stream_event(event, self._state)
|
|
399
|
+
self._update_first(update)
|
|
400
|
+
yield event
|
|
401
|
+
finally:
|
|
402
|
+
self.close()
|
|
403
|
+
|
|
404
|
+
def __iter__(self) -> Iterator[Any]:
|
|
405
|
+
return self._iterate()
|
|
406
|
+
|
|
407
|
+
def __next__(self) -> Any:
|
|
408
|
+
if self._consume is None:
|
|
409
|
+
self._consume = self._iterate()
|
|
410
|
+
return next(self._consume)
|
|
411
|
+
|
|
412
|
+
def __enter__(self) -> _InstrumentedStream:
|
|
413
|
+
enter = getattr(self._inner, "__enter__", None)
|
|
414
|
+
if enter is not None:
|
|
415
|
+
enter()
|
|
416
|
+
return self
|
|
417
|
+
|
|
418
|
+
def __exit__(
|
|
419
|
+
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any
|
|
420
|
+
) -> None:
|
|
421
|
+
if exc is not None:
|
|
422
|
+
self.finish(exc)
|
|
423
|
+
self.close()
|
|
424
|
+
|
|
425
|
+
def close(self) -> None:
|
|
426
|
+
self.finish()
|
|
427
|
+
close = getattr(self._inner, "close", None)
|
|
428
|
+
if close is not None:
|
|
429
|
+
close()
|
|
430
|
+
|
|
431
|
+
def __getattr__(self, name: str) -> Any:
|
|
432
|
+
return getattr(self._inner, name)
|
|
433
|
+
|
|
434
|
+
|
|
435
|
+
class _InstrumentedAsyncStream:
|
|
436
|
+
def __init__(self, inner: Any, handle: telemetry_dev.SpanHandle, started_at: float) -> None:
|
|
437
|
+
self._inner = inner
|
|
438
|
+
self._handle = handle
|
|
439
|
+
self._end = _end_once(handle)
|
|
440
|
+
self._started_at = started_at
|
|
441
|
+
self._state = _StreamState()
|
|
442
|
+
self._saw_first = False
|
|
443
|
+
self._consume: AsyncIterator[Any] | None = None
|
|
444
|
+
|
|
445
|
+
def _update_first(self, update: Mapping[str, Any]) -> None:
|
|
446
|
+
if self._saw_first:
|
|
447
|
+
return
|
|
448
|
+
self._saw_first = True
|
|
449
|
+
self._handle.update(
|
|
450
|
+
time_to_first_chunk_ms=(time.perf_counter() - self._started_at) * 1000,
|
|
451
|
+
response_id=update.get("response_id"),
|
|
452
|
+
response_model=update.get("response_model"),
|
|
453
|
+
)
|
|
454
|
+
|
|
455
|
+
def finish(self, error: BaseException | None = None) -> None:
|
|
456
|
+
fields = _stream_partial(self._state)
|
|
457
|
+
if error is not None:
|
|
458
|
+
fields["error"] = error
|
|
459
|
+
self._end(**fields)
|
|
460
|
+
|
|
461
|
+
async def _aiterate(self) -> AsyncIterator[Any]:
|
|
462
|
+
try:
|
|
463
|
+
while True:
|
|
464
|
+
try:
|
|
465
|
+
event = await self._inner.__anext__()
|
|
466
|
+
except StopAsyncIteration:
|
|
467
|
+
break
|
|
468
|
+
except BaseException as exc:
|
|
469
|
+
self.finish(exc)
|
|
470
|
+
raise
|
|
471
|
+
update = _record_stream_event(event, self._state)
|
|
472
|
+
self._update_first(update)
|
|
473
|
+
yield event
|
|
474
|
+
finally:
|
|
475
|
+
await self.close()
|
|
476
|
+
|
|
477
|
+
def __aiter__(self) -> AsyncIterator[Any]:
|
|
478
|
+
return self._aiterate()
|
|
479
|
+
|
|
480
|
+
async def __anext__(self) -> Any:
|
|
481
|
+
if self._consume is None:
|
|
482
|
+
self._consume = self._aiterate()
|
|
483
|
+
return await self._consume.__anext__()
|
|
484
|
+
|
|
485
|
+
async def __aenter__(self) -> _InstrumentedAsyncStream:
|
|
486
|
+
enter = getattr(self._inner, "__aenter__", None)
|
|
487
|
+
if enter is not None:
|
|
488
|
+
await enter()
|
|
489
|
+
return self
|
|
490
|
+
|
|
491
|
+
async def __aexit__(
|
|
492
|
+
self,
|
|
493
|
+
exc_type: type[BaseException] | None,
|
|
494
|
+
exc: BaseException | None,
|
|
495
|
+
tb: Any,
|
|
496
|
+
) -> None:
|
|
497
|
+
if exc is not None:
|
|
498
|
+
self.finish(exc)
|
|
499
|
+
await self.close()
|
|
500
|
+
|
|
501
|
+
async def close(self) -> None:
|
|
502
|
+
self.finish()
|
|
503
|
+
close = getattr(self._inner, "close", None)
|
|
504
|
+
if close is not None:
|
|
505
|
+
result = close()
|
|
506
|
+
if hasattr(result, "__await__"):
|
|
507
|
+
await result
|
|
508
|
+
|
|
509
|
+
def __getattr__(self, name: str) -> Any:
|
|
510
|
+
return getattr(self._inner, name)
|
|
511
|
+
|
|
512
|
+
|
|
513
|
+
class _InstrumentedMessageStreamManager:
|
|
514
|
+
def __init__(self, inner: Any, name: str, fields: Mapping[str, Any], provider: str) -> None:
|
|
515
|
+
self._inner = inner
|
|
516
|
+
self._name = name
|
|
517
|
+
self._fields = fields
|
|
518
|
+
self._provider = provider
|
|
519
|
+
self._proxy: _InstrumentedStream | None = None
|
|
520
|
+
|
|
521
|
+
def __enter__(self) -> Any:
|
|
522
|
+
handle = telemetry_dev.start_span(
|
|
523
|
+
self._name, provider=self._provider, **_clean_fields(self._fields)
|
|
524
|
+
)
|
|
525
|
+
end = _end_once(handle)
|
|
526
|
+
started_at = time.perf_counter()
|
|
527
|
+
try:
|
|
528
|
+
message_stream = self._inner.__enter__()
|
|
529
|
+
except BaseException as exc:
|
|
530
|
+
end(error=exc)
|
|
531
|
+
raise
|
|
532
|
+
raw_stream = getattr(message_stream, "_raw_stream", None)
|
|
533
|
+
proxy = _InstrumentedStream(raw_stream, handle, started_at)
|
|
534
|
+
message_stream._raw_stream = proxy
|
|
535
|
+
self._proxy = proxy
|
|
536
|
+
return message_stream
|
|
537
|
+
|
|
538
|
+
def __exit__(
|
|
539
|
+
self, exc_type: type[BaseException] | None, exc: BaseException | None, tb: Any
|
|
540
|
+
) -> Any:
|
|
541
|
+
if exc is not None and self._proxy is not None:
|
|
542
|
+
self._proxy.finish(exc)
|
|
543
|
+
result = self._inner.__exit__(exc_type, exc, tb)
|
|
544
|
+
if exc is None and self._proxy is not None:
|
|
545
|
+
self._proxy.finish()
|
|
546
|
+
return result
|
|
547
|
+
|
|
548
|
+
def __getattr__(self, name: str) -> Any:
|
|
549
|
+
return getattr(self._inner, name)
|
|
550
|
+
|
|
551
|
+
|
|
552
|
+
class _InstrumentedAsyncMessageStreamManager:
|
|
553
|
+
def __init__(self, inner: Any, name: str, fields: Mapping[str, Any], provider: str) -> None:
|
|
554
|
+
self._inner = inner
|
|
555
|
+
self._name = name
|
|
556
|
+
self._fields = fields
|
|
557
|
+
self._provider = provider
|
|
558
|
+
self._proxy: _InstrumentedAsyncStream | None = None
|
|
559
|
+
|
|
560
|
+
async def __aenter__(self) -> Any:
|
|
561
|
+
handle = telemetry_dev.start_span(
|
|
562
|
+
self._name, provider=self._provider, **_clean_fields(self._fields)
|
|
563
|
+
)
|
|
564
|
+
end = _end_once(handle)
|
|
565
|
+
started_at = time.perf_counter()
|
|
566
|
+
try:
|
|
567
|
+
message_stream = await self._inner.__aenter__()
|
|
568
|
+
except BaseException as exc:
|
|
569
|
+
end(error=exc)
|
|
570
|
+
raise
|
|
571
|
+
raw_stream = getattr(message_stream, "_raw_stream", None)
|
|
572
|
+
proxy = _InstrumentedAsyncStream(raw_stream, handle, started_at)
|
|
573
|
+
message_stream._raw_stream = proxy
|
|
574
|
+
self._proxy = proxy
|
|
575
|
+
return message_stream
|
|
576
|
+
|
|
577
|
+
async def __aexit__(
|
|
578
|
+
self,
|
|
579
|
+
exc_type: type[BaseException] | None,
|
|
580
|
+
exc: BaseException | None,
|
|
581
|
+
tb: Any,
|
|
582
|
+
) -> Any:
|
|
583
|
+
if exc is not None and self._proxy is not None:
|
|
584
|
+
self._proxy.finish(exc)
|
|
585
|
+
result = await self._inner.__aexit__(exc_type, exc, tb)
|
|
586
|
+
if exc is None and self._proxy is not None:
|
|
587
|
+
self._proxy.finish()
|
|
588
|
+
return result
|
|
589
|
+
|
|
590
|
+
def __getattr__(self, name: str) -> Any:
|
|
591
|
+
return getattr(self._inner, name)
|
|
592
|
+
|
|
593
|
+
|
|
594
|
+
def _start_span(
|
|
595
|
+
params: Mapping[str, Any], mapper: RequestMapper, provider: str
|
|
596
|
+
) -> tuple[telemetry_dev.SpanHandle, Callable[..., None], float]:
|
|
597
|
+
name, fields = mapper(params)
|
|
598
|
+
handle = telemetry_dev.start_span(name, provider=provider, **_clean_fields(fields))
|
|
599
|
+
return handle, _end_once(handle), time.perf_counter()
|
|
600
|
+
|
|
601
|
+
|
|
602
|
+
def _wrap_sync(
|
|
603
|
+
original: Callable[..., Any],
|
|
604
|
+
request_mapper: RequestMapper,
|
|
605
|
+
response_mapper: ResponseMapper,
|
|
606
|
+
provider: ProviderResolver,
|
|
607
|
+
) -> Callable[..., Any]:
|
|
608
|
+
@wraps(original)
|
|
609
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
610
|
+
resource = args[0] if args else None
|
|
611
|
+
call_kwargs = _normalize_request_params(kwargs)
|
|
612
|
+
if _raw_response_requested(call_kwargs):
|
|
613
|
+
return original(*args, **call_kwargs)
|
|
614
|
+
streaming = call_kwargs.get("stream") is True
|
|
615
|
+
handle, end, started_at = _start_span(call_kwargs, request_mapper, provider(resource))
|
|
616
|
+
try:
|
|
617
|
+
result = original(*args, **call_kwargs)
|
|
618
|
+
except BaseException as exc:
|
|
619
|
+
end(error=exc)
|
|
620
|
+
raise
|
|
621
|
+
if streaming:
|
|
622
|
+
return _InstrumentedStream(result, handle, started_at)
|
|
623
|
+
end(**response_mapper(result))
|
|
624
|
+
return result
|
|
625
|
+
|
|
626
|
+
setattr(wrapper, _WRAPPED_ATTR, True)
|
|
627
|
+
setattr(wrapper, _ORIGINAL_ATTR, original)
|
|
628
|
+
return wrapper
|
|
629
|
+
|
|
630
|
+
|
|
631
|
+
def _wrap_async(
|
|
632
|
+
original: Callable[..., Any],
|
|
633
|
+
request_mapper: RequestMapper,
|
|
634
|
+
response_mapper: ResponseMapper,
|
|
635
|
+
provider: ProviderResolver,
|
|
636
|
+
) -> Callable[..., Any]:
|
|
637
|
+
@wraps(original)
|
|
638
|
+
async def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
639
|
+
resource = args[0] if args else None
|
|
640
|
+
call_kwargs = _normalize_request_params(kwargs)
|
|
641
|
+
if _raw_response_requested(call_kwargs):
|
|
642
|
+
return await original(*args, **call_kwargs)
|
|
643
|
+
streaming = call_kwargs.get("stream") is True
|
|
644
|
+
handle, end, started_at = _start_span(call_kwargs, request_mapper, provider(resource))
|
|
645
|
+
try:
|
|
646
|
+
result = await original(*args, **call_kwargs)
|
|
647
|
+
except BaseException as exc:
|
|
648
|
+
end(error=exc)
|
|
649
|
+
raise
|
|
650
|
+
if streaming:
|
|
651
|
+
return _InstrumentedAsyncStream(result, handle, started_at)
|
|
652
|
+
end(**response_mapper(result))
|
|
653
|
+
return result
|
|
654
|
+
|
|
655
|
+
setattr(wrapper, _WRAPPED_ATTR, True)
|
|
656
|
+
setattr(wrapper, _ORIGINAL_ATTR, original)
|
|
657
|
+
return wrapper
|
|
658
|
+
|
|
659
|
+
|
|
660
|
+
def _wrap_stream_manager_sync(
|
|
661
|
+
original: Callable[..., Any], request_mapper: RequestMapper, provider: ProviderResolver
|
|
662
|
+
) -> Callable[..., Any]:
|
|
663
|
+
@wraps(original)
|
|
664
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
665
|
+
resource = args[0] if args else None
|
|
666
|
+
call_kwargs = _normalize_request_params(kwargs)
|
|
667
|
+
name, fields = request_mapper(call_kwargs)
|
|
668
|
+
manager = original(*args, **call_kwargs)
|
|
669
|
+
return _InstrumentedMessageStreamManager(manager, name, fields, provider(resource))
|
|
670
|
+
|
|
671
|
+
setattr(wrapper, _WRAPPED_ATTR, True)
|
|
672
|
+
setattr(wrapper, _ORIGINAL_ATTR, original)
|
|
673
|
+
return wrapper
|
|
674
|
+
|
|
675
|
+
|
|
676
|
+
def _wrap_stream_manager_async(
|
|
677
|
+
original: Callable[..., Any], request_mapper: RequestMapper, provider: ProviderResolver
|
|
678
|
+
) -> Callable[..., Any]:
|
|
679
|
+
@wraps(original)
|
|
680
|
+
def wrapper(*args: Any, **kwargs: Any) -> Any:
|
|
681
|
+
resource = args[0] if args else None
|
|
682
|
+
call_kwargs = _normalize_request_params(kwargs)
|
|
683
|
+
name, fields = request_mapper(call_kwargs)
|
|
684
|
+
manager = original(*args, **call_kwargs)
|
|
685
|
+
return _InstrumentedAsyncMessageStreamManager(manager, name, fields, provider(resource))
|
|
686
|
+
|
|
687
|
+
setattr(wrapper, _WRAPPED_ATTR, True)
|
|
688
|
+
setattr(wrapper, _ORIGINAL_ATTR, original)
|
|
689
|
+
return wrapper
|
|
690
|
+
|
|
691
|
+
|
|
692
|
+
def _own_method(resource: object, name: str) -> bool:
|
|
693
|
+
namespace = getattr(resource, "__dict__", {})
|
|
694
|
+
return isinstance(namespace, Mapping) and name in namespace
|
|
695
|
+
|
|
696
|
+
|
|
697
|
+
def _original_method(current: Any, resource: object) -> Any:
|
|
698
|
+
function = getattr(current, "__func__", current)
|
|
699
|
+
original = getattr(function, _ORIGINAL_ATTR, None)
|
|
700
|
+
if original is None:
|
|
701
|
+
return current
|
|
702
|
+
bind = getattr(original, "__get__", None)
|
|
703
|
+
return bind(resource, type(resource)) if bind is not None else original
|
|
704
|
+
|
|
705
|
+
|
|
706
|
+
def _patch_instance_create(resource: object, async_resource: bool, provider_name: str) -> None:
|
|
707
|
+
resource_any: Any = resource
|
|
708
|
+
current = resource_any.create
|
|
709
|
+
if getattr(current, _WRAPPED_ATTR, False) and _own_method(resource, "create"):
|
|
710
|
+
return
|
|
711
|
+
original = _original_method(current, resource)
|
|
712
|
+
factory = _wrap_async if async_resource else _wrap_sync
|
|
713
|
+
wrapped = factory(original, _messages_request, _messages_response, lambda _: provider_name)
|
|
714
|
+
resource_any.create = wrapped
|
|
715
|
+
|
|
716
|
+
|
|
717
|
+
def _patch_instance_stream(resource: object, async_resource: bool, provider_name: str) -> None:
|
|
718
|
+
resource_any: Any = resource
|
|
719
|
+
current = resource_any.stream
|
|
720
|
+
if getattr(current, _WRAPPED_ATTR, False) and _own_method(resource, "stream"):
|
|
721
|
+
return
|
|
722
|
+
original = _original_method(current, resource)
|
|
723
|
+
factory = _wrap_stream_manager_async if async_resource else _wrap_stream_manager_sync
|
|
724
|
+
wrapped = factory(original, _messages_request, lambda _: provider_name)
|
|
725
|
+
resource_any.stream = wrapped
|
|
726
|
+
|
|
727
|
+
|
|
728
|
+
def _patch_class_create(cls: type[Any], async_resource: bool) -> None:
|
|
729
|
+
original = cls.create
|
|
730
|
+
if getattr(original, _WRAPPED_ATTR, False):
|
|
731
|
+
return
|
|
732
|
+
_ORIGINALS.append((cls, "create", original))
|
|
733
|
+
factory = _wrap_async if async_resource else _wrap_sync
|
|
734
|
+
cls.create = factory(original, _messages_request, _messages_response, _provider_for_resource)
|
|
735
|
+
|
|
736
|
+
|
|
737
|
+
def _patch_class_stream(cls: type[Any], async_resource: bool) -> None:
|
|
738
|
+
original = cls.stream
|
|
739
|
+
if getattr(original, _WRAPPED_ATTR, False):
|
|
740
|
+
return
|
|
741
|
+
_ORIGINALS.append((cls, "stream", original))
|
|
742
|
+
factory = _wrap_stream_manager_async if async_resource else _wrap_stream_manager_sync
|
|
743
|
+
cls.stream = factory(original, _messages_request, _provider_for_resource)
|
|
744
|
+
|
|
745
|
+
|
|
746
|
+
def wrap_anthropic(client: _T) -> _T:
|
|
747
|
+
if getattr(client, _WRAPPED_ATTR, False):
|
|
748
|
+
return client
|
|
749
|
+
client_any: Any = client
|
|
750
|
+
messages = client_any.messages
|
|
751
|
+
async_resource = isinstance(messages, AsyncMessages)
|
|
752
|
+
provider_name = _provider_for_client(cast(object, client))
|
|
753
|
+
_patch_instance_create(messages, async_resource, provider_name)
|
|
754
|
+
_patch_instance_stream(messages, async_resource, provider_name)
|
|
755
|
+
setattr(client, _WRAPPED_ATTR, True)
|
|
756
|
+
return client
|
|
757
|
+
|
|
758
|
+
|
|
759
|
+
def instrument_anthropic() -> None:
|
|
760
|
+
global _installed
|
|
761
|
+
with _install_lock:
|
|
762
|
+
if _installed:
|
|
763
|
+
return
|
|
764
|
+
_patch_class_create(Messages, async_resource=False)
|
|
765
|
+
_patch_class_stream(Messages, async_resource=False)
|
|
766
|
+
_patch_class_create(AsyncMessages, async_resource=True)
|
|
767
|
+
_patch_class_stream(AsyncMessages, async_resource=True)
|
|
768
|
+
_installed = True
|
|
769
|
+
|
|
770
|
+
|
|
771
|
+
def uninstrument_anthropic() -> None:
|
|
772
|
+
global _installed
|
|
773
|
+
with _install_lock:
|
|
774
|
+
while _ORIGINALS:
|
|
775
|
+
cls, method, original = _ORIGINALS.pop()
|
|
776
|
+
setattr(cls, method, original)
|
|
777
|
+
_installed = False
|
|
778
|
+
|
|
779
|
+
|
|
780
|
+
__all__ = ["__version__", "instrument_anthropic", "uninstrument_anthropic", "wrap_anthropic"]
|
|
File without changes
|