152 lines
4.4 KiB
Python
152 lines
4.4 KiB
Python
from llama_index.core.base.llms.types import MessageRole
|
|
from opentelemetry.instrumentation.llamaindex.utils import (
|
|
dont_throw,
|
|
should_send_prompts,
|
|
)
|
|
from opentelemetry.semconv._incubating.attributes import (
|
|
gen_ai_attributes as GenAIAttributes,
|
|
)
|
|
from opentelemetry.semconv_ai import (
|
|
LLMRequestTypeValues,
|
|
SpanAttributes,
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_llm_chat_request(event, span) -> None:
|
|
if not span.is_recording():
|
|
return
|
|
|
|
if should_send_prompts():
|
|
for idx, message in enumerate(event.messages):
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_PROMPT}.{idx}.role", message.role.value
|
|
)
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_PROMPT}.{idx}.content", message.content
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_llm_chat_request_model_attributes(event, span):
|
|
if span and not span.is_recording():
|
|
return
|
|
|
|
model_dict = event.model_dict
|
|
span.set_attribute(SpanAttributes.LLM_REQUEST_TYPE, LLMRequestTypeValues.CHAT.value)
|
|
|
|
# For StructuredLLM, the model and temperature are nested under model_dict.llm
|
|
if "llm" in model_dict:
|
|
model_dict = model_dict.get("llm", {})
|
|
|
|
span.set_attribute(GenAIAttributes.GEN_AI_REQUEST_MODEL, model_dict.get("model"))
|
|
span.set_attribute(
|
|
GenAIAttributes.GEN_AI_REQUEST_TEMPERATURE, model_dict.get("temperature")
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_llm_chat_response(event, span) -> None:
|
|
if not span.is_recording():
|
|
return
|
|
|
|
response = event.response
|
|
if should_send_prompts():
|
|
for idx, message in enumerate(event.messages):
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_PROMPT}.{idx}.role", message.role.value
|
|
)
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_PROMPT}.{idx}.content", message.content
|
|
)
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_COMPLETION}.0.role",
|
|
response.message.role.value,
|
|
)
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_COMPLETION}.0.content",
|
|
response.message.content,
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_llm_chat_response_model_attributes(event, span):
|
|
if not span.is_recording():
|
|
return
|
|
|
|
response = event.response
|
|
|
|
if not (raw := response.raw):
|
|
return
|
|
|
|
span.set_attribute(
|
|
GenAIAttributes.GEN_AI_RESPONSE_MODEL,
|
|
(
|
|
raw.get("model") if "model" in raw else raw.model
|
|
), # raw can be Any, not just ChatCompletion
|
|
)
|
|
if usage := raw.get("usage") if "usage" in raw else raw.usage:
|
|
span.set_attribute(
|
|
GenAIAttributes.GEN_AI_USAGE_OUTPUT_TOKENS, usage.completion_tokens
|
|
)
|
|
span.set_attribute(GenAIAttributes.GEN_AI_USAGE_INPUT_TOKENS, usage.prompt_tokens)
|
|
span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, usage.total_tokens)
|
|
if choices := raw.choices:
|
|
span.set_attribute(
|
|
SpanAttributes.LLM_RESPONSE_FINISH_REASON, choices[0].finish_reason
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_llm_predict_response(event, span) -> None:
|
|
if should_send_prompts():
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_COMPLETION}.role",
|
|
MessageRole.ASSISTANT.value,
|
|
)
|
|
span.set_attribute(
|
|
f"{GenAIAttributes.GEN_AI_COMPLETION}.content",
|
|
event.output,
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_embedding(event, span) -> None:
|
|
model_dict = event.model_dict
|
|
span.set_attribute(
|
|
f"{LLMRequestTypeValues.EMBEDDING.value}.model_name",
|
|
model_dict.get("model_name"),
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_rerank(event, span) -> None:
|
|
if not span.is_recording():
|
|
return
|
|
if should_send_prompts():
|
|
span.set_attribute(
|
|
f"{LLMRequestTypeValues.RERANK.value}.query",
|
|
event.query.query_str,
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_rerank_model_attributes(event, span):
|
|
if not span.is_recording():
|
|
return
|
|
span.set_attribute(
|
|
f"{LLMRequestTypeValues.RERANK.value}.model_name",
|
|
event.model_name,
|
|
)
|
|
span.set_attribute(
|
|
f"{LLMRequestTypeValues.RERANK.value}.top_n",
|
|
event.top_n,
|
|
)
|
|
|
|
|
|
@dont_throw
|
|
def set_tool(event, span) -> None:
|
|
span.set_attribute("tool.name", event.tool.name)
|
|
span.set_attribute("tool.description", event.tool.description)
|
|
span.set_attribute("tool.arguments", event.arguments)
|