Skip to content

feat(observability): add messages genai span attributes #68

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
bzp2010 merged 5 commits into from
Apr 29, 2026
Merged

feat(observability): add messages genai span attributes #68

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
370d53d
feat(observability): add messages genai span attributes
bzp2010 Apr 29, 2026
c12cd96
fix
bzp2010 Apr 29, 2026
8b5ee93
fix comment
bzp2010 Apr 29, 2026
7049c31
fix comment
bzp2010 Apr 29, 2026
2fe06de
fix: ensure only one fta when multiple content block
bzp2010 Apr 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
263 changes: 10 additions & 253 deletions src/proxy/handlers/chat_completions/span_attributes/message_attributes.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,43 +1,16 @@
use serde::Serialize;
use serde_json::{Map, Value};

use crate::gateway::types::openai::{
ChatCompletionResponse, ChatMessage, ContentPart, MessageContent, Tool,
pub(super) use crate::proxy::utils::trace::span_message_attributes::{
ContentPartView, MessageContentView, MessageView, OutputMessageView, ToolCallView,
append_openinference_message_properties, append_openinference_output_message_properties,
gen_ai_input_messages_json, gen_ai_output_messages_json,
};
use crate::{
gateway::types::openai::{
ChatCompletionResponse, ChatMessage, ContentPart, MessageContent, Tool,
},
proxy::utils::trace::span_message_attributes::serialize_to_json_string,
};

#[derive(Clone)]
pub(super) enum MessageContentView {
Text(String),
Parts(Vec<ContentPartView>),
}

#[derive(Clone)]
pub(super) enum ContentPartView {
Text(String),
ImageUrl { url: String },
}

#[derive(Clone)]
pub(super) struct ToolCallView {
pub(super) id: Option<String>,
pub(super) name: String,
pub(super) arguments: String,
}

#[derive(Clone)]
pub(super) struct MessageView {
pub(super) role: String,
pub(super) content: Option<MessageContentView>,
pub(super) name: Option<String>,
pub(super) tool_calls: Vec<ToolCallView>,
pub(super) tool_call_id: Option<String>,
}

#[derive(Clone)]
pub(super) struct OutputMessageView {
pub(super) message: MessageView,
pub(super) finish_reason: Option<String>,
}

pub(super) fn message_view_from_chat_message(message: &ChatMessage) -> MessageView {
MessageView {
Expand Down Expand Up @@ -75,64 +48,6 @@ pub(super) fn response_output_message_views(
.collect()
}

pub(super) fn gen_ai_input_messages_json(messages: &[MessageView]) -> Option<String> {
if messages.is_empty() {
return None;
}

let values: Vec<Value> = messages
.iter()
.map(|message| {
let mut value = Map::new();
value.insert("role".into(), Value::String(message.role.clone()));
value.insert("parts".into(), Value::Array(gen_ai_message_parts(message)));

if let Some(name) = &message.name {
value.insert("name".into(), Value::String(name.clone()));
}

Value::Object(value)
})
.collect();

serialize_to_json_string(&values)
}

pub(super) fn gen_ai_output_messages_json(messages: &[OutputMessageView]) -> Option<String> {
if messages.is_empty() {
return None;
}

let values: Vec<Value> = messages
.iter()
.map(|message| {
let mut value = Map::new();
value.insert("role".into(), Value::String(message.message.role.clone()));
value.insert(
"parts".into(),
Value::Array(gen_ai_message_parts(&message.message)),
);
value.insert(
"finish_reason".into(),
Value::String(
message
.finish_reason
.clone()
.unwrap_or_else(|| "unknown".into()),
),
);

if let Some(name) = &message.message.name {
value.insert("name".into(), Value::String(name.clone()));
}

Value::Object(value)
})
.collect();

serialize_to_json_string(&values)
}

pub(super) fn gen_ai_tool_definitions_json(tools: &[Tool]) -> Option<String> {
if tools.is_empty() {
return None;
Expand Down Expand Up @@ -160,80 +75,6 @@ pub(super) fn gen_ai_tool_definitions_json(tools: &[Tool]) -> Option<String> {
serialize_to_json_string(&values)
}

pub(super) fn append_openinference_message_properties(
properties: &mut Vec<(String, String)>,
prefix: &str,
messages: &[MessageView],
) {
for (message_index, message) in messages.iter().enumerate() {
let prefix = format!("{prefix}.{message_index}.message");
properties.push((format!("{prefix}.role"), message.role.clone()));

if let Some(name) = &message.name {
properties.push((format!("{prefix}.name"), name.clone()));
}

if let Some(tool_call_id) = &message.tool_call_id {
properties.push((format!("{prefix}.tool_call_id"), tool_call_id.clone()));
}

match &message.content {
Some(MessageContentView::Text(text)) if !text.is_empty() => {
properties.push((format!("{prefix}.content"), text.clone()));
}
Some(MessageContentView::Parts(parts)) => {
for (content_index, part) in parts.iter().enumerate() {
let part_prefix = format!("{prefix}.contents.{content_index}.message_content");
match part {
ContentPartView::Text(text) => {
properties.push((format!("{part_prefix}.type"), "text".into()));
properties.push((format!("{part_prefix}.text"), text.clone()));
}
ContentPartView::ImageUrl { url } => {
properties.push((format!("{part_prefix}.type"), "image".into()));
properties
.push((format!("{part_prefix}.image.image.url"), url.clone()));
}
}
}
}
_ => {}
}

for (tool_call_index, tool_call) in message.tool_calls.iter().enumerate() {
let tool_call_prefix = format!("{prefix}.tool_calls.{tool_call_index}.tool_call");

if let Some(id) = &tool_call.id {
properties.push((format!("{tool_call_prefix}.id"), id.clone()));
}

properties.push((
format!("{tool_call_prefix}.function.name"),
tool_call.name.clone(),
));

if !tool_call.arguments.is_empty() {
properties.push((
format!("{tool_call_prefix}.function.arguments"),
tool_call.arguments.clone(),
));
}
}
}
}

pub(super) fn append_openinference_output_message_properties(
properties: &mut Vec<(String, String)>,
prefix: &str,
messages: &[OutputMessageView],
) {
let message_views: Vec<_> = messages
.iter()
.map(|message| message.message.clone())
.collect();
append_openinference_message_properties(properties, prefix, &message_views);
}

pub(super) fn append_openinference_tool_properties(
properties: &mut Vec<(String, String)>,
tools: &[Tool],
Expand All @@ -258,14 +99,6 @@ pub(super) fn append_openinference_tool_properties(
}
}

fn serialize_to_json_string(value: &impl Serialize) -> Option<String> {
serde_json::to_string(value).ok()
}

fn parse_json_or_string(raw: &str) -> Value {
serde_json::from_str(raw).unwrap_or_else(|_| Value::String(raw.to_string()))
}

fn content_part_view_from_content_part(part: &ContentPart) -> ContentPartView {
match part {
ContentPart::Text { text } => ContentPartView::Text(text.clone()),
Expand All @@ -286,79 +119,3 @@ fn message_content_view_from_message_content(content: &MessageContent) -> Messag
),
}
}

fn gen_ai_message_parts(message: &MessageView) -> Vec<Value> {
if message.role == "tool" {
let mut part = Map::new();
part.insert("type".into(), Value::String("tool_call_response".into()));

if let Some(tool_call_id) = &message.tool_call_id {
part.insert("id".into(), Value::String(tool_call_id.clone()));
}

let response = match &message.content {
Some(MessageContentView::Text(text)) => parse_json_or_string(text),
Some(MessageContentView::Parts(parts)) => {
Value::Array(parts.iter().map(gen_ai_content_part_value).collect())
}
None => Value::Null,
};
part.insert("response".into(), response);

return vec![Value::Object(part)];
}

let mut parts = Vec::new();

match &message.content {
Some(MessageContentView::Text(text)) if !text.is_empty() => {
parts.push(gen_ai_text_part(text));
}
Some(MessageContentView::Parts(content_parts)) => {
parts.extend(content_parts.iter().map(gen_ai_content_part_value));
}
_ => {}
}

for tool_call in &message.tool_calls {
let mut part = Map::new();
part.insert("type".into(), Value::String("tool_call".into()));

if let Some(id) = &tool_call.id {
part.insert("id".into(), Value::String(id.clone()));
}

part.insert("name".into(), Value::String(tool_call.name.clone()));

if !tool_call.arguments.is_empty() {
part.insert(
"arguments".into(),
parse_json_or_string(&tool_call.arguments),
);
}

parts.push(Value::Object(part));
}

parts
}

fn gen_ai_text_part(text: &str) -> Value {
let mut part = Map::new();
part.insert("type".into(), Value::String("text".into()));
part.insert("content".into(), Value::String(text.to_string()));
Value::Object(part)
}

fn gen_ai_content_part_value(part: &ContentPartView) -> Value {
match part {
ContentPartView::Text(text) => gen_ai_text_part(text),
ContentPartView::ImageUrl { url } => {
let mut part = Map::new();
part.insert("type".into(), Value::String("uri".into()));
part.insert("modality".into(), Value::String("image".into()));
part.insert("uri".into(), Value::String(url.clone()));
Value::Object(part)
}
}
}
7 changes: 5 additions & 2 deletions src/proxy/handlers/chat_completions/span_attributes/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,11 @@ mod telemetry;

pub(super) use stream_output::StreamOutputCollector;
pub(super) use telemetry::{
apply_span_properties, chunk_span_properties, request_span_properties,
response_span_properties, usage_span_properties,
chunk_span_properties, request_span_properties, response_span_properties,
};

pub(super) use crate::proxy::utils::trace::span_attributes::{
apply_span_properties, usage_span_properties,
};

#[cfg(test)]
Expand Down
23 changes: 5 additions & 18 deletions src/proxy/handlers/chat_completions/span_attributes/stream_output.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,10 +1,10 @@
use std::collections::BTreeMap;

use super::message_attributes::{
MessageContentView, MessageView, OutputMessageView, ToolCallView,
append_openinference_output_message_properties, gen_ai_output_messages_json,
use super::message_attributes::{MessageContentView, MessageView, OutputMessageView, ToolCallView};
use crate::{
gateway::types::openai::ChatCompletionChunk,
proxy::utils::trace::span_message_attributes::output_message_span_properties,
};
use crate::gateway::types::openai::ChatCompletionChunk;

#[derive(Default)]
struct StreamOutputToolCall {
Expand Down Expand Up @@ -72,20 +72,7 @@ impl StreamOutputCollector {
pub(in crate::proxy::handlers::chat_completions) fn output_message_span_properties(
&self,
) -> Vec<(String, String)> {
let output_messages = self.output_message_views();
let mut properties = Vec::new();

append_openinference_output_message_properties(
&mut properties,
"llm.output_messages",
&output_messages,
);

if let Some(value) = gen_ai_output_messages_json(&output_messages) {
properties.push(("gen_ai.output.messages".into(), value));
}

properties
output_message_span_properties(&self.output_message_views())
}

fn output_message_views(&self) -> Vec<OutputMessageView> {
Expand Down
Loading