adding canonical tracing support via bright-staff

This commit is contained in:
Salman Paracha 2025-12-09 17:45:22 -08:00
parent 09c0b999b2
commit c748ea0a71
27 changed files with 1803 additions and 273 deletions

View file

@ -200,6 +200,17 @@ impl ProviderRequest for ConverseRequest {
})
}
fn get_tool_names(&self) -> Option<Vec<String>> {
self.tool_config.as_ref()?.tools.as_ref().map(|tools| {
tools
.iter()
.filter_map(|tool| match tool {
Tool::ToolSpec { tool_spec } => Some(tool_spec.name.clone()),
})
.collect()
})
}
fn to_bytes(&self) -> Result<Vec<u8>, ProviderRequestError> {
serde_json::to_vec(self).map_err(|e| ProviderRequestError {
message: format!("Failed to serialize Bedrock request: {}", e),

View file

@ -513,6 +513,12 @@ impl ProviderRequest for MessagesRequest {
None
}
fn get_tool_names(&self) -> Option<Vec<String>> {
self.tools.as_ref().map(|tools| {
tools.iter().map(|tool| tool.name.clone()).collect()
})
}
fn to_bytes(&self) -> Result<Vec<u8>, ProviderRequestError> {
serde_json::to_vec(self).map_err(|e| ProviderRequestError {
message: format!("Failed to serialize MessagesRequest: {}", e),

View file

@ -687,6 +687,32 @@ impl ProviderRequest for ChatCompletionsRequest {
})
}
fn get_tool_names(&self) -> Option<Vec<String>> {
// First check the 'tools' field (current API)
if let Some(tools) = &self.tools {
let names: Vec<String> = tools
.iter()
.map(|tool| tool.function.name.clone())
.collect();
if !names.is_empty() {
return Some(names);
}
}
// Fallback to 'functions' field (deprecated but still supported)
if let Some(functions) = &self.functions {
let names: Vec<String> = functions
.iter()
.map(|func| func.function.name.clone())
.collect();
if !names.is_empty() {
return Some(names);
}
}
None
}
fn to_bytes(&self) -> Result<Vec<u8>, ProviderRequestError> {
serde_json::to_vec(&self).map_err(|e| ProviderRequestError {
message: format!("Failed to serialize OpenAI request: {}", e),

View file

@ -1063,6 +1063,19 @@ impl ProviderRequest for ResponsesAPIRequest {
}
}
fn get_tool_names(&self) -> Option<Vec<String>> {
self.tools.as_ref().map(|tools| {
tools
.iter()
.filter_map(|tool| match tool {
Tool::Function { name, .. } => Some(name.clone()),
// Other tool types don't have user-defined names
_ => None,
})
.collect()
})
}
fn to_bytes(&self) -> Result<Vec<u8>, ProviderRequestError> {
serde_json::to_vec(&self).map_err(|e| ProviderRequestError {
message: format!("Failed to serialize Responses API request: {}", e),

View file

@ -35,6 +35,9 @@ pub trait ProviderRequest: Send + Sync {
/// Extract the user message for tracing/logging purposes
fn get_recent_user_message(&self) -> Option<String>;
/// Get tool names if tools are defined in the request
fn get_tool_names(&self) -> Option<Vec<String>>;
/// Convert the request to bytes for transmission
fn to_bytes(&self) -> Result<Vec<u8>, ProviderRequestError>;
@ -95,6 +98,16 @@ impl ProviderRequest for ProviderRequestType {
}
}
fn get_tool_names(&self) -> Option<Vec<String>> {
match self {
Self::ChatCompletionsRequest(r) => r.get_tool_names(),
Self::MessagesRequest(r) => r.get_tool_names(),
Self::BedrockConverse(r) => r.get_tool_names(),
Self::BedrockConverseStream(r) => r.get_tool_names(),
Self::ResponsesAPIRequest(r) => r.get_tool_names(),
}
}
fn to_bytes(&self) -> Result<Vec<u8>, ProviderRequestError> {
match self {
Self::ChatCompletionsRequest(r) => r.to_bytes(),