Anthropic instrumentation (#181)

Co-authored-by: Alex Hall <alex.mojaki@gmail.com>

Author: willbakst

logfire/__init__.py

@@ -21,6 +21,7 @@
 install_auto_tracing = DEFAULT_LOGFIRE_INSTANCE.install_auto_tracing
 instrument_fastapi = DEFAULT_LOGFIRE_INSTANCE.instrument_fastapi
 instrument_openai = DEFAULT_LOGFIRE_INSTANCE.instrument_openai
+instrument_anthropic = DEFAULT_LOGFIRE_INSTANCE.instrument_anthropic
 instrument_asyncpg = DEFAULT_LOGFIRE_INSTANCE.instrument_asyncpg
 instrument_psycopg = DEFAULT_LOGFIRE_INSTANCE.instrument_psycopg
 shutdown = DEFAULT_LOGFIRE_INSTANCE.shutdown

logfire/_internal/integrations/llm_providers/anthropic.py

@@ -0,0 +1,76 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import anthropic
+from anthropic.types import ContentBlockDeltaEvent, ContentBlockStartEvent, Message
+from anthropic.types.beta.tools import ToolsBetaMessage
+
+from .types import EndpointConfig
+
+if TYPE_CHECKING:
+    from anthropic._models import FinalRequestOptions
+    from anthropic._types import ResponseT
+
+    from ...main import LogfireSpan
+
+
+__all__ = (
+    'get_endpoint_config',
+    'on_response',
+    'is_async_client',
+)
+
+
+def get_endpoint_config(options: FinalRequestOptions) -> EndpointConfig:
+    """Returns the endpoint config for Anthropic depending on the url."""
+    url = options.url
+    json_data = options.json_data
+    if not isinstance(json_data, dict):
+        raise ValueError('Expected `options.json_data` to be a dictionary')
+
+    if url == '/v1/messages' or url == '/v1/messages?beta=tools':
+        return EndpointConfig(
+            message_template='Message with {request_data[model]!r}',
+            span_data={'request_data': json_data},
+            content_from_stream=content_from_messages,
+        )
+    else:
+        raise ValueError(f'Unknown Anthropic API endpoint: `{url}`')
+
+
+def content_from_messages(chunk: anthropic.types.MessageStreamEvent) -> str | None:
+    if isinstance(chunk, ContentBlockStartEvent):
+        return chunk.content_block.text
+    if isinstance(chunk, ContentBlockDeltaEvent):
+        return chunk.delta.text
+    return None
+
+
+def on_response(response: ResponseT, span: LogfireSpan) -> ResponseT:
+    """Updates the span based on the type of response."""
+    if isinstance(response, (Message, ToolsBetaMessage)):  # pragma: no branch
+        block = response.content[0]
+        message: dict[str, Any] = {'role': 'assistant'}
+        if block.type == 'text':
+            message['content'] = block.text
+        else:
+            message['tool_calls'] = [
+                {
+                    'function': {
+                        'arguments': block.model_dump_json(include={'input'}),
+                        'name': block.name,  # type: ignore
+                    }
+                }
+                for block in response.content
+            ]
+        span.set_attribute('response_data', {'message': message, 'usage': response.usage})
+    return response
+
+
+def is_async_client(client: anthropic.Anthropic | anthropic.AsyncAnthropic):
+    """Returns whether or not `client` is async."""
+    if isinstance(client, anthropic.Anthropic):
+        return False
+    assert isinstance(client, anthropic.AsyncAnthropic), f'Unexpected Anthropic or AsyncAnthropic type, got: {client}'
+    return True

logfire/_internal/integrations/llm_providers/llm_provider.py

@@ -0,0 +1,165 @@
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, AsyncIterator, Callable, ContextManager, Iterator
+
+from opentelemetry import context
+
+from ...constants import ONE_SECOND_IN_NANOSECONDS
+
+if TYPE_CHECKING:
+    from ...main import Logfire, LogfireSpan
+    from .types import EndpointConfig
+
+
+__all__ = ('instrument_llm_provider',)
+
+
+def instrument_llm_provider(
+    logfire: Logfire,
+    client: Any,
+    suppress_otel: bool,
+    scope_suffix: str,
+    get_endpoint_config_fn: Callable[[Any], EndpointConfig],
+    on_response_fn: Callable[[Any, LogfireSpan], Any],
+    is_async_client_fn: Callable[[Any], bool],
+) -> ContextManager[None]:
+    """Instruments the provided `client` with `logfire`."""
+    logfire_llm = logfire.with_settings(custom_scope_suffix=scope_suffix.lower(), tags=['LLM'])
+
+    client._is_instrumented_by_logfire = True
+    client._original_request_method = original_request_method = client._request
+
+    is_async = is_async_client_fn(client)
+
+    def _instrumentation_setup(**kwargs: Any) -> Any:
+        if context.get_value('suppress_instrumentation'):
+            return None, None, kwargs
+
+        options = kwargs['options']
+        try:
+            message_template, span_data, content_from_stream = get_endpoint_config_fn(options)
+        except ValueError as exc:
+            logfire_llm.warn(
+                'Unable to instrument {suffix} API call: {error}', suffix=scope_suffix, error=str(exc), kwargs=kwargs
+            )
+            return None, None, kwargs
+
+        span_data['async'] = is_async
+
+        stream = kwargs['stream']
+
+        if stream and content_from_stream:
+            stream_cls = kwargs['stream_cls']
+            assert stream_cls is not None, 'Expected `stream_cls` when streaming'
+
+            if is_async:
+
+                class LogfireInstrumentedAsyncStream(stream_cls):
+                    async def __stream__(self) -> AsyncIterator[Any]:
+                        with record_streaming(logfire_llm, span_data, content_from_stream) as record_chunk:
+                            async for chunk in super().__stream__():  # type: ignore
+                                record_chunk(chunk)
+                                yield chunk
+
+                kwargs['stream_cls'] = LogfireInstrumentedAsyncStream
+            else:
+
+                class LogfireInstrumentedStream(stream_cls):
+                    def __stream__(self) -> Iterator[Any]:
+                        with record_streaming(logfire_llm, span_data, content_from_stream) as record_chunk:
+                            for chunk in super().__stream__():  # type: ignore
+                                record_chunk(chunk)
+                                yield chunk
+
+                kwargs['stream_cls'] = LogfireInstrumentedStream
+
+        return message_template, span_data, kwargs
+
+    def instrumented_llm_request_sync(**kwargs: Any) -> Any:
+        message_template, span_data, kwargs = _instrumentation_setup(**kwargs)
+        if message_template is None:
+            return original_request_method(**kwargs)
+        stream = kwargs['stream']
+        with logfire_llm.span(message_template, **span_data) as span:
+            with maybe_suppress_instrumentation(suppress_otel):
+                if stream:
+                    return original_request_method(**kwargs)
+                else:
+                    response = on_response_fn(original_request_method(**kwargs), span)
+                    return response
+
+    async def instrumented_llm_request_async(**kwargs: Any) -> Any:
+        message_template, span_data, kwargs = _instrumentation_setup(**kwargs)
+        if message_template is None:
+            return await original_request_method(**kwargs)
+        stream = kwargs['stream']
+        with logfire_llm.span(message_template, **span_data) as span:
+            with maybe_suppress_instrumentation(suppress_otel):
+                if stream:
+                    return await original_request_method(**kwargs)
+                else:
+                    response = on_response_fn(await original_request_method(**kwargs), span)
+                    return response
+
+    if is_async:
+        client._request = instrumented_llm_request_async
+    else:
+        client._request = instrumented_llm_request_sync
+
+    @contextmanager
+    def uninstrument_context():
+        """Context manager to remove instrumentation from the LLM client.
+
+        The user isn't required (or even expected) to use this context manager,
+        which is why the instrumenting has already happened before.
+        It exists mostly for tests and just in case users want it.
+        """
+        try:
+            yield
+        finally:
+            client._request = client._original_request_method
+            del client._original_request_method
+            client._is_instrumented_by_logfire = False
+
+    return uninstrument_context()
+
+
+@contextmanager
+def maybe_suppress_instrumentation(suppress: bool) -> Iterator[None]:
+    if suppress:
+        new_context = context.set_value('suppress_instrumentation', True)
+        token = context.attach(new_context)
+        try:
+            yield
+        finally:
+            context.detach(token)
+    else:
+        yield
+
+
+@contextmanager
+def record_streaming(
+    logire_llm: Logfire,
+    span_data: dict[str, Any],
+    content_from_stream: Callable[[Any], str | None],
+):
+    content: list[str] = []
+
+    def record_chunk(chunk: Any) -> Any:
+        chunk_content = content_from_stream(chunk)
+        if chunk_content is not None:
+            content.append(chunk_content)
+
+    timer = logire_llm._config.ns_timestamp_generator  # type: ignore
+    start = timer()
+    try:
+        yield record_chunk
+    finally:
+        duration = (timer() - start) / ONE_SECOND_IN_NANOSECONDS
+        logire_llm.info(
+            'streaming response from {request_data[model]!r} took {duration:.2f}s',
+            **span_data,
+            duration=duration,
+            response_data={'combined_chunk_content': ''.join(content), 'chunk_count': len(content)},
+        )

logfire/_internal/integrations/llm_providers/openai.py

@@ -0,0 +1,107 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, cast
+
+import openai
+from openai._legacy_response import LegacyAPIResponse
+from openai.types.chat.chat_completion import ChatCompletion
+from openai.types.chat.chat_completion_chunk import ChatCompletionChunk
+from openai.types.completion import Completion
+from openai.types.create_embedding_response import CreateEmbeddingResponse
+from openai.types.images_response import ImagesResponse
+
+from .types import EndpointConfig
+
+if TYPE_CHECKING:
+    from openai._models import FinalRequestOptions
+    from openai._types import ResponseT
+
+    from ...main import LogfireSpan
+
+__all__ = (
+    'get_endpoint_config',
+    'on_response',
+    'is_async_client',
+)
+
+
+def get_endpoint_config(options: FinalRequestOptions) -> EndpointConfig:
+    """Returns the endpoint config for OpenAI depending on the url."""
+    url = options.url
+    json_data = options.json_data
+    if not isinstance(json_data, dict):
+        raise ValueError('Expected `options.json_data` to be a dictionary')
+    if 'model' not in json_data:
+        # all OpenAI API calls have a model AFAIK
+        raise ValueError('`model` not found in request data')
+
+    if url == '/chat/completions':
+        return EndpointConfig(
+            message_template='Chat Completion with {request_data[model]!r}',
+            span_data={'request_data': json_data},
+            content_from_stream=content_from_chat_completions,
+        )
+    elif url == '/completions':
+        return EndpointConfig(
+            message_template='Completion with {request_data[model]!r}',
+            span_data={'request_data': json_data},
+            content_from_stream=content_from_completions,
+        )
+    elif url == '/embeddings':
+        return EndpointConfig(
+            message_template='Embedding Creation with {request_data[model]!r}',
+            span_data={'request_data': json_data},
+            content_from_stream=None,
+        )
+    elif url == '/images/generations':
+        return EndpointConfig(
+            message_template='Image Generation with {request_data[model]!r}',
+            span_data={'request_data': json_data},
+            content_from_stream=None,
+        )
+    else:
+        raise ValueError(f'Unknown OpenAI API endpoint: `{url}`')
+
+
+def content_from_completions(chunk: Completion | None) -> str | None:
+    if chunk and chunk.choices:
+        return chunk.choices[0].text
+    return None  # pragma: no cover
+
+
+def content_from_chat_completions(chunk: ChatCompletionChunk | None) -> str | None:
+    if chunk and chunk.choices:
+        return chunk.choices[0].delta.content
+    return None
+
+
+def on_response(response: ResponseT, span: LogfireSpan) -> ResponseT:
+    """Updates the span based on the type of response."""
+    if isinstance(response, LegacyAPIResponse):  # pragma: no cover
+        on_response(response.parse(), span)  # type: ignore
+        return cast('ResponseT', response)
+
+    if isinstance(response, ChatCompletion):
+        span.set_attribute(
+            'response_data',
+            {'message': response.choices[0].message, 'usage': response.usage},
+        )
+    elif isinstance(response, Completion):
+        first_choice = response.choices[0]
+        span.set_attribute(
+            'response_data',
+            {'finish_reason': first_choice.finish_reason, 'text': first_choice.text, 'usage': response.usage},
+        )
+    elif isinstance(response, CreateEmbeddingResponse):
+        span.set_attribute('response_data', {'usage': response.usage})
+    elif isinstance(response, ImagesResponse):  # pragma: no branch
+        span.set_attribute('response_data', {'images': response.data})
+    return response
+
+
+def is_async_client(client: openai.OpenAI | openai.AsyncOpenAI):
+    """Returns whether or not `client` is async."""
+    if isinstance(client, openai.OpenAI):
+        return False
+    assert isinstance(client, openai.AsyncOpenAI), f'Unexpected OpenAI or AsyncOpenAI type, got: {client}'
+    return True

logfire/_internal/integrations/llm_providers/types.py

@@ -0,0 +1,13 @@
+from __future__ import annotations
+
+from typing import Any, Callable, NamedTuple
+
+from typing_extensions import LiteralString
+
+
+class EndpointConfig(NamedTuple):
+    """The configuration for the endpoint of a provider based on request url."""
+
+    message_template: LiteralString
+    span_data: dict[str, Any]
+    content_from_stream: Callable[[Any], str | None] | None