ai-station/.venv/lib/python3.12/site-packages/opentelemetry/instrumentation/anthropic/span_utils.py

407 lines
16 KiB
Python

import json
import logging
from typing import Any, Dict
from opentelemetry.instrumentation.anthropic.config import Config
from opentelemetry.instrumentation.anthropic.utils import (
JSONEncoder,
dont_throw,
model_as_dict,
should_send_prompts,
_extract_response_data,
)
from opentelemetry.semconv._incubating.attributes import (
gen_ai_attributes as GenAIAttributes,
)
from opentelemetry.semconv_ai import SpanAttributes
logger = logging.getLogger(__name__)
def _is_base64_image(item: Dict[str, Any]) -> bool:
if not isinstance(item, dict):
return False
if not isinstance(item.get("source"), dict):
return False
if item.get("type") != "image" or item["source"].get("type") != "base64":
return False
return True
async def _process_image_item(item, trace_id, span_id, message_index, content_index):
if not Config.upload_base64_image:
return item
image_format = item.get("source").get("media_type").split("/")[1]
image_name = f"message_{message_index}_content_{content_index}.{image_format}"
base64_string = item.get("source").get("data")
url = await Config.upload_base64_image(trace_id, span_id, image_name, base64_string)
return {"type": "image_url", "image_url": {"url": url}}
async def _dump_content(message_index, content, span):
if isinstance(content, str):
return content
elif isinstance(content, list):
# If the content is a list of text blocks, concatenate them.
# This is more commonly used in prompt caching.
if all([model_as_dict(item).get("type") == "text" for item in content]):
return "".join([model_as_dict(item).get("text") for item in content])
content = [
(
await _process_image_item(
model_as_dict(item),
span.context.trace_id,
span.context.span_id,
message_index,
j,
)
if _is_base64_image(model_as_dict(item))
else model_as_dict(item)
)
for j, item in enumerate(content)
]
return json.dumps(content, cls=JSONEncoder)
@dont_throw
async def aset_input_attributes(span, kwargs):
from opentelemetry.instrumentation.anthropic import set_span_attribute
set_span_attribute(span, GenAIAttributes.GEN_AI_REQUEST_MODEL, kwargs.get("model"))
set_span_attribute(
span, GenAIAttributes.GEN_AI_REQUEST_MAX_TOKENS, kwargs.get("max_tokens_to_sample")
)
set_span_attribute(
span, GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, kwargs.get("temperature")
)
set_span_attribute(span, GenAIAttributes.GEN_AI_REQUEST_TOP_P, kwargs.get("top_p"))
set_span_attribute(
span, SpanAttributes.LLM_FREQUENCY_PENALTY, kwargs.get("frequency_penalty")
)
set_span_attribute(
span, SpanAttributes.LLM_PRESENCE_PENALTY, kwargs.get("presence_penalty")
)
set_span_attribute(span, SpanAttributes.LLM_IS_STREAMING, kwargs.get("stream"))
if should_send_prompts():
if kwargs.get("prompt") is not None:
set_span_attribute(
span, f"{GenAIAttributes.GEN_AI_PROMPT}.0.user", kwargs.get("prompt")
)
elif kwargs.get("messages") is not None:
has_system_message = False
if kwargs.get("system"):
has_system_message = True
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.0.content",
await _dump_content(
message_index=0, span=span, content=kwargs.get("system")
),
)
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.0.role",
"system",
)
for i, message in enumerate(kwargs.get("messages")):
prompt_index = i + (1 if has_system_message else 0)
content = message.get("content")
tool_use_blocks = []
other_blocks = []
if isinstance(content, list):
for block in content:
if dict(block).get("type") == "tool_use":
tool_use_blocks.append(dict(block))
else:
other_blocks.append(block)
content = other_blocks
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.content",
await _dump_content(
message_index=i, span=span, content=message.get("content")
),
)
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.role",
message.get("role"),
)
if tool_use_blocks:
for tool_num, tool_use_block in enumerate(tool_use_blocks):
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.tool_calls.{tool_num}.id",
tool_use_block.get("id"),
)
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.tool_calls.{tool_num}.name",
tool_use_block.get("name"),
)
set_span_attribute(
span,
f"{GenAIAttributes.GEN_AI_PROMPT}.{prompt_index}.tool_calls.{tool_num}.arguments",
json.dumps(tool_use_block.get("input")),
)
if kwargs.get("tools") is not None:
for i, tool in enumerate(kwargs.get("tools")):
prefix = f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{i}"
set_span_attribute(span, f"{prefix}.name", tool.get("name"))
set_span_attribute(
span, f"{prefix}.description", tool.get("description")
)
input_schema = tool.get("input_schema")
if input_schema is not None:
set_span_attribute(
span, f"{prefix}.input_schema", json.dumps(input_schema)
)
output_format = kwargs.get("output_format")
if output_format and isinstance(output_format, dict):
if output_format.get("type") == "json_schema":
schema = output_format.get("schema")
if schema:
set_span_attribute(
span,
SpanAttributes.LLM_REQUEST_STRUCTURED_OUTPUT_SCHEMA,
json.dumps(schema),
)
async def _aset_span_completions(span, response):
if not should_send_prompts():
return
from opentelemetry.instrumentation.anthropic import set_span_attribute
from opentelemetry.instrumentation.anthropic.utils import _aextract_response_data
response = await _aextract_response_data(response)
index = 0
prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
set_span_attribute(span, f"{prefix}.finish_reason", response.get("stop_reason"))
if response.get("role"):
set_span_attribute(span, f"{prefix}.role", response.get("role"))
if response.get("completion"):
set_span_attribute(span, f"{prefix}.content", response.get("completion"))
elif response.get("content"):
tool_call_index = 0
text = ""
for content in response.get("content"):
content_block_type = content.type
# usually, Antrhopic responds with just one text block,
# but the API allows for multiple text blocks, so concatenate them
if content_block_type == "text" and hasattr(content, "text"):
text += content.text
elif content_block_type == "thinking":
content = dict(content)
# override the role to thinking
set_span_attribute(
span,
f"{prefix}.role",
"thinking",
)
set_span_attribute(
span,
f"{prefix}.content",
content.get("thinking"),
)
# increment the index for subsequent content blocks
index += 1
prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
# set the role to the original role on the next completions
set_span_attribute(
span,
f"{prefix}.role",
response.get("role"),
)
elif content_block_type == "tool_use":
content = dict(content)
set_span_attribute(
span,
f"{prefix}.tool_calls.{tool_call_index}.id",
content.get("id"),
)
set_span_attribute(
span,
f"{prefix}.tool_calls.{tool_call_index}.name",
content.get("name"),
)
tool_arguments = content.get("input")
if tool_arguments is not None:
set_span_attribute(
span,
f"{prefix}.tool_calls.{tool_call_index}.arguments",
json.dumps(tool_arguments),
)
tool_call_index += 1
set_span_attribute(span, f"{prefix}.content", text)
def _set_span_completions(span, response):
if not should_send_prompts():
return
from opentelemetry.instrumentation.anthropic import set_span_attribute
response = _extract_response_data(response)
index = 0
prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
set_span_attribute(span, f"{prefix}.finish_reason", response.get("stop_reason"))
if response.get("role"):
set_span_attribute(span, f"{prefix}.role", response.get("role"))
if response.get("completion"):
set_span_attribute(span, f"{prefix}.content", response.get("completion"))
elif response.get("content"):
tool_call_index = 0
text = ""
for content in response.get("content"):
content_block_type = content.type
# usually, Antrhopic responds with just one text block,
# but the API allows for multiple text blocks, so concatenate them
if content_block_type == "text" and hasattr(content, "text"):
text += content.text
elif content_block_type == "thinking":
content = dict(content)
# override the role to thinking
set_span_attribute(
span,
f"{prefix}.role",
"thinking",
)
set_span_attribute(
span,
f"{prefix}.content",
content.get("thinking"),
)
# increment the index for subsequent content blocks
index += 1
prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
# set the role to the original role on the next completions
set_span_attribute(
span,
f"{prefix}.role",
response.get("role"),
)
elif content_block_type == "tool_use":
content = dict(content)
set_span_attribute(
span,
f"{prefix}.tool_calls.{tool_call_index}.id",
content.get("id"),
)
set_span_attribute(
span,
f"{prefix}.tool_calls.{tool_call_index}.name",
content.get("name"),
)
tool_arguments = content.get("input")
if tool_arguments is not None:
set_span_attribute(
span,
f"{prefix}.tool_calls.{tool_call_index}.arguments",
json.dumps(tool_arguments),
)
tool_call_index += 1
set_span_attribute(span, f"{prefix}.content", text)
@dont_throw
async def aset_response_attributes(span, response):
from opentelemetry.instrumentation.anthropic import set_span_attribute
from opentelemetry.instrumentation.anthropic.utils import _aextract_response_data
response = await _aextract_response_data(response)
set_span_attribute(span, GenAIAttributes.GEN_AI_RESPONSE_MODEL, response.get("model"))
set_span_attribute(span, GenAIAttributes.GEN_AI_RESPONSE_ID, response.get("id"))
if response.get("usage"):
prompt_tokens = response.get("usage").input_tokens
completion_tokens = response.get("usage").output_tokens
set_span_attribute(span, GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens)
set_span_attribute(
span, GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
)
set_span_attribute(
span,
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
prompt_tokens + completion_tokens,
)
await _aset_span_completions(span, response)
@dont_throw
def set_response_attributes(span, response):
from opentelemetry.instrumentation.anthropic import set_span_attribute
response = _extract_response_data(response)
set_span_attribute(span, GenAIAttributes.GEN_AI_RESPONSE_MODEL, response.get("model"))
set_span_attribute(span, GenAIAttributes.GEN_AI_RESPONSE_ID, response.get("id"))
if response.get("usage"):
prompt_tokens = response.get("usage").input_tokens
completion_tokens = response.get("usage").output_tokens
set_span_attribute(span, GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS, prompt_tokens)
set_span_attribute(
span, GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, completion_tokens
)
set_span_attribute(
span,
SpanAttributes.LLM_USAGE_TOTAL_TOKENS,
prompt_tokens + completion_tokens,
)
_set_span_completions(span, response)
@dont_throw
def set_streaming_response_attributes(span, complete_response_events):
if not should_send_prompts():
return
from opentelemetry.instrumentation.anthropic import set_span_attribute
if not span.is_recording() or not complete_response_events:
return
index = 0
for event in complete_response_events:
prefix = f"{GenAIAttributes.GEN_AI_COMPLETION}.{index}"
set_span_attribute(span, f"{prefix}.finish_reason", event.get("finish_reason"))
role = "thinking" if event.get("type") == "thinking" else "assistant"
# Thinking is added as a separate completion, so we need to increment the index
if event.get("type") == "thinking":
index += 1
set_span_attribute(span, f"{prefix}.role", role)
if event.get("type") == "tool_use":
set_span_attribute(
span,
f"{prefix}.tool_calls.0.id",
event.get("id"),
)
set_span_attribute(
span,
f"{prefix}.tool_calls.0.name",
event.get("name"),
)
tool_arguments = event.get("input")
if tool_arguments is not None:
set_span_attribute(
span,
f"{prefix}.tool_calls.0.arguments",
tool_arguments,
)
else:
set_span_attribute(span, f"{prefix}.content", event.get("text"))