Skip to content

Commit db1f57a

Browse files
stephenhibbertStephen Hibbertalexmojaki
authored
Support for AnthropicBedrock client (#701)
Co-authored-by: Stephen Hibbert <[email protected]> Co-authored-by: Alex Hall <[email protected]>
1 parent ecdd723 commit db1f57a

File tree

6 files changed

+256
-18
lines changed

6 files changed

+256
-18
lines changed

‎docs/integrations/llms/anthropic.md

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,3 +103,21 @@ Shows up like this in Logfire:
103103
![Logfire Anthropic Streaming](../../images/logfire-screenshot-anthropic-stream.png){ width="500" }
104104
<figcaption>Anthropic streaming response</figcaption>
105105
</figure>
106+
107+
## Amazon Bedrock
108+
109+
You can also log Anthropic LLM calls to Amazon Bedrock using the `AmazonBedrock` and `AsyncAmazonBedrock` clients.
110+
111+
```python
112+
import anthropic
113+
import logfire
114+
115+
client = anthropic.AnthropicBedrock(
116+
aws_region='us-east-1',
117+
aws_access_key='access-key',
118+
aws_secret_key='secret-key',
119+
)
120+
121+
logfire.configure()
122+
logfire.instrument_anthropic(client)
123+
```

‎logfire/_internal/integrations/llm_providers/anthropic.py

Lines changed: 11 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@
2121

2222

2323
def get_endpoint_config(options: FinalRequestOptions) -> EndpointConfig:
24-
"""Returns the endpoint config for Anthropic depending on the url."""
24+
"""Returns the endpoint config for Anthropic or Bedrock depending on the url."""
2525
url = options.url
2626
json_data = options.json_data
2727
if not isinstance(json_data, dict): # pragma: no cover
@@ -83,9 +83,16 @@ def on_response(response: ResponseT, span: LogfireSpan) -> ResponseT:
8383
return response
8484

8585

86-
def is_async_client(client: type[anthropic.Anthropic] | type[anthropic.AsyncAnthropic]):
86+
def is_async_client(
87+
client: type[anthropic.Anthropic]
88+
| type[anthropic.AsyncAnthropic]
89+
| type[anthropic.AnthropicBedrock]
90+
| type[anthropic.AsyncAnthropicBedrock],
91+
):
8792
"""Returns whether or not the `client` class is async."""
88-
if issubclass(client, anthropic.Anthropic):
93+
if issubclass(client, (anthropic.Anthropic, anthropic.AnthropicBedrock)):
8994
return False
90-
assert issubclass(client, anthropic.AsyncAnthropic), f'Expected Anthropic or AsyncAnthropic type, got: {client}'
95+
assert issubclass(
96+
client, (anthropic.AsyncAnthropic, anthropic.AsyncAnthropicBedrock)
97+
), f'Expected Anthropic, AsyncAnthropic, AnthropicBedrock or AsyncAnthropicBedrock type, got: {client}'
9198
return True

‎logfire/_internal/main.py

Lines changed: 24 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -1072,17 +1072,23 @@ def instrument_openai(
10721072

10731073
def instrument_anthropic(
10741074
self,
1075-
anthropic_client: anthropic.Anthropic
1076-
| anthropic.AsyncAnthropic
1077-
| type[anthropic.Anthropic]
1078-
| type[anthropic.AsyncAnthropic]
1079-
| None = None,
1075+
anthropic_client: (
1076+
anthropic.Anthropic
1077+
| anthropic.AsyncAnthropic
1078+
| anthropic.AnthropicBedrock
1079+
| anthropic.AsyncAnthropicBedrock
1080+
| type[anthropic.Anthropic]
1081+
| type[anthropic.AsyncAnthropic]
1082+
| type[anthropic.AnthropicBedrock]
1083+
| type[anthropic.AsyncAnthropicBedrock]
1084+
| None
1085+
) = None,
10801086
*,
10811087
suppress_other_instrumentation: bool = True,
10821088
) -> ContextManager[None]:
10831089
"""Instrument an Anthropic client so that spans are automatically created for each request.
10841090
1085-
The following methods are instrumented for both the sync and the async clients:
1091+
The following methods are instrumented for both the sync and async clients:
10861092
10871093
- [`client.messages.create`](https://docs.anthropic.com/en/api/messages)
10881094
- [`client.messages.stream`](https://docs.anthropic.com/en/api/messages-streaming)
@@ -1097,6 +1103,7 @@ def instrument_anthropic(
10971103
import anthropic
10981104
10991105
client = anthropic.Anthropic()
1106+
11001107
logfire.configure()
11011108
logfire.instrument_anthropic(client)
11021109
@@ -1112,13 +1119,10 @@ def instrument_anthropic(
11121119
11131120
Args:
11141121
anthropic_client: The Anthropic client or class to instrument:
1115-
1116-
- `None` (the default) to instrument both the
1117-
`anthropic.Anthropic` and `anthropic.AsyncAnthropic` classes.
1118-
- The `anthropic.Anthropic` class or a subclass
1119-
- The `anthropic.AsyncAnthropic` class or a subclass
1120-
- An instance of `anthropic.Anthropic`
1121-
- An instance of `anthropic.AsyncAnthropic`
1122+
- `None` (the default) to instrument all Anthropic client types
1123+
- The `anthropic.Anthropic` or `anthropic.AnthropicBedrock` class or subclass
1124+
- The `anthropic.AsyncAnthropic` or `anthropic.AsyncAnthropicBedrock` class or subclass
1125+
- An instance of any of the above classes
11221126
11231127
suppress_other_instrumentation: If True, suppress any other OTEL instrumentation that may be otherwise
11241128
enabled. In reality, this means the HTTPX instrumentation, which could otherwise be called since
@@ -1136,7 +1140,13 @@ def instrument_anthropic(
11361140
self._warn_if_not_initialized_for_instrumentation()
11371141
return instrument_llm_provider(
11381142
self,
1139-
anthropic_client or (anthropic.Anthropic, anthropic.AsyncAnthropic),
1143+
anthropic_client
1144+
or (
1145+
anthropic.Anthropic,
1146+
anthropic.AsyncAnthropic,
1147+
anthropic.AnthropicBedrock,
1148+
anthropic.AsyncAnthropicBedrock,
1149+
),
11401150
suppress_other_instrumentation,
11411151
'Anthropic',
11421152
get_endpoint_config,

‎pyproject.toml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -160,6 +160,9 @@ dev = [
160160
"requests",
161161
"setuptools>=75.3.0",
162162
"aiosqlite>=0.20.0",
163+
"boto3 >= 1.28.57",
164+
"botocore >= 1.31.57",
165+
163166
]
164167
docs = [
165168
"mkdocs>=1.5.0",

‎tests/otel_integrations/test_anthropic_bedrock.py

Lines changed: 147 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,147 @@
1+
from typing import Iterator
2+
3+
import httpx
4+
import pytest
5+
from anthropic import Anthropic, AnthropicBedrock, AsyncAnthropic, AsyncAnthropicBedrock
6+
from anthropic.types import Message, TextBlock, Usage
7+
from dirty_equals import IsJson
8+
from httpx._transports.mock import MockTransport
9+
from inline_snapshot import snapshot
10+
11+
import logfire
12+
from logfire._internal.integrations.llm_providers.anthropic import is_async_client
13+
from logfire.testing import TestExporter
14+
15+
16+
def request_handler(request: httpx.Request) -> httpx.Response:
17+
"""Used to mock httpx requests"""
18+
model_id = 'anthropic.claude-3-haiku-20240307-v1:0'
19+
20+
assert request.method == 'POST'
21+
assert request.url == f'https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke'
22+
23+
return httpx.Response(
24+
200,
25+
json=Message(
26+
id='test_id',
27+
content=[
28+
TextBlock(
29+
text='Nine',
30+
type='text',
31+
)
32+
],
33+
model=model_id,
34+
role='assistant',
35+
type='message',
36+
usage=Usage(input_tokens=2, output_tokens=3), # Match the snapshot values
37+
).model_dump(mode='json'),
38+
)
39+
40+
41+
@pytest.fixture
42+
def mock_client() -> Iterator[AnthropicBedrock]:
43+
"""Fixture that provides a mocked Anthropic client with AWS credentials"""
44+
with httpx.Client(transport=MockTransport(request_handler)) as http_client:
45+
client = AnthropicBedrock(
46+
aws_region='us-east-1',
47+
aws_access_key='test-access-key',
48+
aws_secret_key='test-secret-key',
49+
aws_session_token='test-session-token',
50+
http_client=http_client,
51+
)
52+
with logfire.instrument_anthropic():
53+
yield client
54+
55+
56+
@pytest.mark.filterwarnings('ignore:datetime.datetime.utcnow:DeprecationWarning')
57+
def test_sync_messages(mock_client: AnthropicBedrock, exporter: TestExporter):
58+
"""Test basic synchronous message creation"""
59+
model_id = 'anthropic.claude-3-haiku-20240307-v1:0'
60+
response = mock_client.messages.create(
61+
max_tokens=1000,
62+
model=model_id,
63+
system='You are a helpful assistant.',
64+
messages=[{'role': 'user', 'content': 'What is four plus five?'}],
65+
)
66+
67+
# Verify response structure
68+
assert isinstance(response.content[0], TextBlock)
69+
assert response.content[0].text == 'Nine'
70+
71+
# Verify exported spans
72+
assert exporter.exported_spans_as_dict() == snapshot(
73+
[
74+
{
75+
'name': 'Message with {request_data[model]!r}',
76+
'context': {'trace_id': 1, 'span_id': 1, 'is_remote': False},
77+
'parent': None,
78+
'start_time': 1000000000,
79+
'end_time': 2000000000,
80+
'attributes': {
81+
'code.filepath': 'test_anthropic_bedrock.py',
82+
'code.function': 'test_sync_messages',
83+
'code.lineno': 123,
84+
'request_data': IsJson(
85+
{
86+
'max_tokens': 1000,
87+
'system': 'You are a helpful assistant.',
88+
'messages': [{'role': 'user', 'content': 'What is four plus five?'}],
89+
'model': model_id,
90+
}
91+
),
92+
'async': False,
93+
'logfire.msg_template': 'Message with {request_data[model]!r}',
94+
'logfire.msg': f"Message with '{model_id}'",
95+
'logfire.span_type': 'span',
96+
'logfire.tags': ('LLM',),
97+
'response_data': IsJson(
98+
{
99+
'message': {
100+
'content': 'Nine',
101+
'role': 'assistant',
102+
},
103+
'usage': {
104+
'input_tokens': 2,
105+
'output_tokens': 3,
106+
'cache_creation_input_tokens': None,
107+
'cache_read_input_tokens': None,
108+
},
109+
}
110+
),
111+
'logfire.json_schema': IsJson(
112+
{
113+
'type': 'object',
114+
'properties': {
115+
'request_data': {'type': 'object'},
116+
'async': {},
117+
'response_data': {
118+
'type': 'object',
119+
'properties': {
120+
'usage': {
121+
'type': 'object',
122+
'title': 'Usage',
123+
'x-python-datatype': 'PydanticModel',
124+
},
125+
},
126+
},
127+
},
128+
}
129+
),
130+
},
131+
}
132+
]
133+
)
134+
135+
136+
def test_is_async_client() -> None:
137+
# Test sync clients
138+
assert not is_async_client(Anthropic)
139+
assert not is_async_client(AnthropicBedrock)
140+
141+
# Test async clients
142+
assert is_async_client(AsyncAnthropic)
143+
assert is_async_client(AsyncAnthropicBedrock)
144+
145+
# Test invalid input
146+
with pytest.raises(AssertionError):
147+
is_async_client(str) # type: ignore

‎uv.lock

Lines changed: 53 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

0 commit comments

Comments
 (0)