diff --git a/python/instrumentation/openinference-instrumentation-anthropic/src/openinference/instrumentation/anthropic/_wrappers.py b/python/instrumentation/openinference-instrumentation-anthropic/src/openinference/instrumentation/anthropic/_wrappers.py index 050200479..fa3129a87 100644 --- a/python/instrumentation/openinference-instrumentation-anthropic/src/openinference/instrumentation/anthropic/_wrappers.py +++ b/python/instrumentation/openinference-instrumentation-anthropic/src/openinference/instrumentation/anthropic/_wrappers.py @@ -29,7 +29,7 @@ if TYPE_CHECKING: from pydantic import BaseModel - from anthropic.types import Usage + from anthropic.types import Message, Usage class _WithTracer(ABC): @@ -92,7 +92,7 @@ def __call__( attributes=dict( chain( get_attributes_from_context(), - _get_llm_model(arguments), + _get_llm_model_name_from_input(arguments), _get_llm_provider(), _get_llm_system(), _get_llm_span_kind(), @@ -145,7 +145,7 @@ async def __call__( attributes=dict( chain( get_attributes_from_context(), - _get_llm_model(arguments), + _get_llm_model_name_from_input(arguments), _get_llm_provider(), _get_llm_system(), _get_llm_span_kind(), @@ -198,7 +198,7 @@ def __call__( attributes=dict( chain( get_attributes_from_context(), - _get_llm_model(arguments), + _get_llm_model_name_from_input(arguments), _get_llm_provider(), _get_llm_system(), _get_llm_span_kind(), @@ -223,6 +223,7 @@ def __call__( span.set_attributes( dict( chain( + _get_llm_model_name_from_response(response), _get_output_messages(response), _get_llm_token_counts(response.usage), _get_outputs(response), @@ -260,7 +261,7 @@ async def __call__( get_attributes_from_context(), _get_llm_provider(), _get_llm_system(), - _get_llm_model(arguments), + _get_llm_model_name_from_input(arguments), _get_llm_span_kind(), _get_llm_input_messages(llm_input_messages), _get_llm_invocation_parameters(invocation_parameters), @@ -283,6 +284,7 @@ async def __call__( span.set_attributes( dict( chain( + _get_llm_model_name_from_response(response), _get_output_messages(response), _get_llm_token_counts(response.usage), _get_outputs(response), @@ -326,11 +328,16 @@ def _get_llm_token_counts(usage: "Usage") -> Iterator[Tuple[str, Any]]: yield LLM_TOKEN_COUNT_COMPLETION, usage.output_tokens -def _get_llm_model(arguments: Mapping[str, Any]) -> Iterator[Tuple[str, Any]]: +def _get_llm_model_name_from_input(arguments: Mapping[str, Any]) -> Iterator[Tuple[str, Any]]: if model_name := arguments.get("model"): yield LLM_MODEL_NAME, model_name +def _get_llm_model_name_from_response(message: "Message") -> Iterator[Tuple[str, Any]]: + if model_name := getattr(message, "model"): + yield LLM_MODEL_NAME, model_name + + def _get_llm_invocation_parameters( invocation_parameters: Mapping[str, Any], ) -> Iterator[Tuple[str, Any]]: diff --git a/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/cassettes/test_instrumentor/test_anthropic_instrumentation_messages.yaml b/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/cassettes/test_instrumentor/test_anthropic_instrumentation_messages.yaml index e8c640e51..d517e4882 100644 --- a/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/cassettes/test_instrumentor/test_anthropic_instrumentation_messages.yaml +++ b/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/cassettes/test_instrumentor/test_anthropic_instrumentation_messages.yaml @@ -1,7 +1,7 @@ interactions: - request: body: '{"max_tokens": 1024, "messages": [{"role": "user", "content": "What''s - the capital of France?"}], "model": "claude-3-opus-20240229"}' + the capital of France?"}], "model": "claude-3-opus-latest"}' headers: {} method: POST uri: https://api.anthropic.com/v1/messages diff --git a/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/test_instrumentor.py b/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/test_instrumentor.py index 541b65e87..540600e69 100644 --- a/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/test_instrumentor.py +++ b/python/instrumentation/openinference-instrumentation-anthropic/tests/openinference/anthropic/test_instrumentor.py @@ -279,7 +279,7 @@ def test_anthropic_instrumentation_messages( client = Anthropic(api_key="fake") input_message = "What's the capital of France?" - invocation_params = {"max_tokens": 1024, "model": "claude-3-opus-20240229"} + invocation_params = {"max_tokens": 1024, "model": "claude-3-opus-latest"} client.messages.create( max_tokens=1024, @@ -289,7 +289,7 @@ def test_anthropic_instrumentation_messages( "content": input_message, } ], - model="claude-3-opus-20240229", + model="claude-3-opus-latest", ) spans = in_memory_span_exporter.get_finished_spans()