tkmind/tkmind_go
1
0
Fork 0
Code Issues Pull Requests Actions 7 Packages Projects Releases Wiki Activity

improve provider request logging a bit (#5236)

Browse Source
This commit is contained in:
Jack Amadeo
2025-10-18 11:35:42 -04:00
committed by GitHub
parent 3b04b2ebf4
commit 9b4c8726f1
19 changed files with 200 additions and 100 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/goose/src/providers/anthropic.rs
+12 -6
View File
@@ -14,11 +14,12 @@ use super::errors::ProviderError;
use super::formats::anthropic::{
create_request, get_usage, response_to_message, response_to_streaming_message,
};
use super::utils::{emit_debug_trace, get_model, map_http_error_to_provider_error};
use super::utils::{get_model, map_http_error_to_provider_error};
use crate::config::declarative_providers::DeclarativeProviderConfig;
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::retry::ProviderRetry;
use crate::providers::utils::RequestLog;
use rmcp::model::Tool;
pub const ANTHROPIC_DEFAULT_MODEL: &str = "claude-sonnet-4-0";
@@ -204,7 +205,8 @@ impl Provider for AnthropicProvider {
usage.input_tokens, usage.output_tokens, usage.total_tokens);
let response_model = get_model(&json_response);
emit_debug_trace(&self.model, &payload, &json_response, &usage);
let mut log = RequestLog::start(&self.model, &payload)?;
log.write(&json_response, Some(&usage))?;
let provider_usage = ProviderUsage::new(response_model, usage);
tracing::debug!(
"🔍 Anthropic non-streaming returning ProviderUsage: {:?}",
@@ -258,22 +260,26 @@ impl Provider for AnthropicProvider {
.insert("stream".to_string(), Value::Bool(true));
let mut request = self.api_client.request("v1/messages");
let mut log = RequestLog::start(&self.model, &payload)?;
for (key, value) in self.get_conditional_headers() {
request = request.header(key, value)?;
}
let response = request.response_post(&payload).await?;
let response = request.response_post(&payload).await.inspect_err(|e| {
let _ = log.error(e);
})?;
if !response.status().is_success() {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
let error_json = serde_json::from_str::<Value>(&error_text).ok();
return Err(map_http_error_to_provider_error(status, error_json));
let error = map_http_error_to_provider_error(status, error_json);
let _ = log.error(&error);
return Err(error);
}
let stream = response.bytes_stream().map_err(io::Error::other);
let model = self.model.clone();
Ok(Box::pin(try_stream! {
let stream_reader = StreamReader::new(stream);
let framed = tokio_util::codec::FramedRead::new(stream_reader, tokio_util::codec::LinesCodec::new()).map_err(anyhow::Error::from);
@@ -282,7 +288,7 @@ impl Provider for AnthropicProvider {
pin!(message_stream);
while let Some(message) = futures::StreamExt::next(&mut message_stream).await {
let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
emit_debug_trace(&model, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default());
log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
yield (message, usage);
}
}))
crates/goose/src/providers/azure.rs
+4 -2
View File
@@ -9,9 +9,10 @@ use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
use super::retry::ProviderRetry;
use super::utils::{emit_debug_trace, get_model, handle_response_openai_compat, ImageFormat};
use super::utils::{get_model, handle_response_openai_compat, ImageFormat};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::utils::RequestLog;
use rmcp::model::Tool;
pub const AZURE_DEFAULT_MODEL: &str = "gpt-4o";
@@ -156,7 +157,8 @@ impl Provider for AzureProvider {
Usage::default()
});
let response_model = get_model(&response);
emit_debug_trace(model_config, &payload, &response, &usage);
let mut log = RequestLog::start(model_config, &payload)?;
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
}
crates/goose/src/providers/bedrock.rs
+5 -6
View File
@@ -5,7 +5,7 @@ use super::errors::ProviderError;
use super::retry::{ProviderRetry, RetryConfig};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::utils::emit_debug_trace;
use crate::providers::utils::RequestLog;
use anyhow::Result;
use async_trait::async_trait;
use aws_sdk_bedrockruntime::config::ProvideCredentials;
@@ -222,12 +222,11 @@ impl Provider for BedrockProvider {
"messages": messages,
"tools": tools
});
emit_debug_trace(
&self.model,
&debug_payload,
let mut log = RequestLog::start(&self.model, &debug_payload)?;
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
&usage,
);
Some(&usage),
)?;
let provider_usage = ProviderUsage::new(model_name.to_string(), usage);
Ok((message, provider_usage))
crates/goose/src/providers/claude_code.rs
+3 -2
View File
@@ -9,7 +9,7 @@ use tokio::process::Command;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::utils::emit_debug_trace;
use super::utils::RequestLog;
use crate::config::Config;
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
@@ -495,13 +495,14 @@ impl Provider for ClaudeCodeProvider {
"system": system,
"messages": messages.len()
});
let mut log = RequestLog::start(model_config, &payload)?;
let response = json!({
"lines": json_lines.len(),
"usage": usage
});
emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((
message,
crates/goose/src/providers/cursor_agent.rs
+3 -2
View File
@@ -9,7 +9,7 @@ use tokio::process::Command;
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::utils::emit_debug_trace;
use super::utils::RequestLog;
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
use rmcp::model::Tool;
@@ -433,7 +433,8 @@ impl Provider for CursorAgentProvider {
"usage": usage
});
emit_debug_trace(model_config, &payload, &response, &usage);
let mut log = RequestLog::start(&self.model, &payload)?;
log.write(&response, Some(&usage))?;
Ok((
message,
crates/goose/src/providers/databricks.rs
+10 -4
View File
@@ -18,6 +18,7 @@ use super::oauth;
use super::retry::ProviderRetry;
use super::utils::{
get_model, handle_response_openai_compat, map_http_error_to_provider_error, ImageFormat,
RequestLog,
};
use crate::config::ConfigError;
use crate::conversation::message::Message;
@@ -286,6 +287,8 @@ impl Provider for DatabricksProvider {
.expect("payload should have model key")
.remove("model");
let mut log = RequestLog::start(&self.model, &payload)?;
let response = self
.with_retry(|| self.post(payload.clone(), Some(&model_config.model_name)))
.await?;
@@ -296,7 +299,7 @@ impl Provider for DatabricksProvider {
Usage::default()
});
let response_model = get_model(&response);
super::utils::emit_debug_trace(&self.model, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
@@ -322,6 +325,7 @@ impl Provider for DatabricksProvider {
.insert("stream".to_string(), Value::Bool(true));
let path = self.get_endpoint_path(&model_config.model_name, false);
let mut log = RequestLog::start(&self.model, &payload)?;
let response = self
.with_retry(|| async {
let resp = self.api_client.response_post(&path, &payload).await?;
@@ -335,11 +339,13 @@ impl Provider for DatabricksProvider {
}
Ok(resp)
})
.await?;
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
let stream = response.bytes_stream().map_err(io::Error::other);
let model = self.model.clone();
Ok(Box::pin(try_stream! {
let stream_reader = StreamReader::new(stream);
let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from);
@@ -348,7 +354,7 @@ impl Provider for DatabricksProvider {
pin!(message_stream);
while let Some(message) = message_stream.next().await {
let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
super::utils::emit_debug_trace(&model, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default());
log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
yield (message, usage);
}
}))
crates/goose/src/providers/gcpvertexai.rs
+3 -2
View File
@@ -21,7 +21,7 @@ use crate::providers::formats::gcpvertexai::{
use crate::providers::formats::gcpvertexai::GcpLocation::Iowa;
use crate::providers::gcpauth::GcpAuth;
use crate::providers::retry::RetryConfig;
use crate::providers::utils::emit_debug_trace;
use crate::providers::utils::RequestLog;
use rmcp::model::Tool;
/// Base URL for GCP Vertex AI documentation
@@ -518,7 +518,8 @@ impl Provider for GcpVertexAIProvider {
let response = self.post(&request, &context).await?;
let usage = get_usage(&response, &context)?;
emit_debug_trace(model_config, &request, &response, &usage);
let mut log = RequestLog::start(model_config, &request)?;
log.write(&response, Some(&usage))?;
// Convert response to message
let message = response_to_message(response, context)?;
crates/goose/src/providers/gemini_cli.rs
+14 -8
View File
@@ -8,7 +8,7 @@ use tokio::process::Command;
use super::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::utils::emit_debug_trace;
use super::utils::RequestLog;
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
@@ -317,12 +317,12 @@ impl Provider for GeminiCliProvider {
}
#[tracing::instrument(
skip(self, model_config, system, messages, tools),
skip(self, _model_config, system, messages, tools),
fields(model_config, input, output, input_tokens, output_tokens, total_tokens)
)]
async fn complete_with_model(
&self,
model_config: &ModelConfig,
_model_config: &ModelConfig,
system: &str,
messages: &[Message],
tools: &[Tool],
@@ -332,10 +332,6 @@ impl Provider for GeminiCliProvider {
return self.generate_simple_session_description(messages);
}
let lines = self.execute_command(system, messages, tools).await?;
let (message, usage) = self.parse_response(&lines)?;
// Create a dummy payload for debug tracing
let payload = json!({
"command": self.command,
@@ -344,12 +340,22 @@ impl Provider for GeminiCliProvider {
"messages": messages.len()
});
let mut log = RequestLog::start(&self.model, &payload).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to start request log: {}", e))
})?;
let lines = self.execute_command(system, messages, tools).await?;
let (message, usage) = self.parse_response(&lines)?;
let response = json!({
"lines": lines.len(),
"usage": usage
});
emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage)).map_err(|e| {
ProviderError::RequestFailed(format!("Failed to write request log: {}", e))
})?;
Ok((
message,
crates/goose/src/providers/githubcopilot.rs
+3 -2
View File
@@ -15,7 +15,7 @@ use super::base::{Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
use super::retry::ProviderRetry;
use super::utils::{emit_debug_trace, get_model, handle_response_openai_compat, ImageFormat};
use super::utils::{get_model, handle_response_openai_compat, ImageFormat, RequestLog};
use crate::config::{Config, ConfigError};
use crate::conversation::message::Message;
@@ -408,6 +408,7 @@ impl Provider for GithubCopilotProvider {
tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
let payload = create_request(model_config, system, messages, tools, &ImageFormat::OpenAi)?;
let mut log = RequestLog::start(model_config, &payload)?;
// Make request with retry
let response = self
@@ -424,7 +425,7 @@ impl Provider for GithubCopilotProvider {
Usage::default()
});
let response_model = get_model(&response);
emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
crates/goose/src/providers/google.rs
+3 -2
View File
@@ -1,7 +1,7 @@
use super::api_client::{ApiClient, AuthMethod};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{emit_debug_trace, handle_response_google_compat, unescape_json_values};
use super::utils::{handle_response_google_compat, unescape_json_values, RequestLog};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
@@ -113,6 +113,7 @@ impl Provider for GoogleProvider {
tools: &[Tool],
) -> Result<(Message, ProviderUsage), ProviderError> {
let payload = create_request(model_config, system, messages, tools)?;
let mut log = RequestLog::start(model_config, &payload)?;
// Make request
let response = self
@@ -129,7 +130,7 @@ impl Provider for GoogleProvider {
Some(model_version) => model_version.as_str().unwrap_or_default().to_string(),
None => model_config.model_name.clone(),
};
emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
let provider_usage = ProviderUsage::new(response_model, usage);
Ok((message, provider_usage))
}
crates/goose/src/providers/litellm.rs
+3 -2
View File
@@ -8,7 +8,7 @@ use super::base::{ConfigKey, ModelInfo, Provider, ProviderMetadata, ProviderUsag
use super::embedding::EmbeddingCapable;
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{emit_debug_trace, get_model, handle_response_openai_compat, ImageFormat};
use super::utils::{get_model, handle_response_openai_compat, ImageFormat, RequestLog};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
@@ -188,7 +188,8 @@ impl Provider for LiteLLMProvider {
let message = super::formats::openai::response_to_message(&response)?;
let usage = super::formats::openai::get_usage(&response);
let response_model = get_model(&response);
emit_debug_trace(model_config, &payload, &response, &usage);
let mut log = RequestLog::start(model_config, &payload)?;
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
crates/goose/src/providers/ollama.rs
+7 -3
View File
@@ -2,7 +2,9 @@ use super::api_client::{ApiClient, AuthMethod};
use super::base::{ConfigKey, MessageStream, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{get_model, handle_response_openai_compat, handle_status_openai_compat};
use super::utils::{
get_model, handle_response_openai_compat, handle_status_openai_compat, RequestLog,
};
use crate::config::declarative_providers::DeclarativeProviderConfig;
use crate::conversation::message::Message;
use crate::conversation::Conversation;
@@ -213,7 +215,8 @@ impl Provider for OllamaProvider {
Usage::default()
});
let response_model = get_model(&response);
super::utils::emit_debug_trace(model_config, &payload, &response, &usage);
let mut log = RequestLog::start(model_config, &payload)?;
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
@@ -270,13 +273,14 @@ impl Provider for OllamaProvider {
let model_config = self.model.clone();
Ok(Box::pin(try_stream! {
let mut log = RequestLog::start(&model_config, &payload)?;
let stream_reader = StreamReader::new(stream);
let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from);
let message_stream = response_to_streaming_message(framed);
pin!(message_stream);
while let Some(message) = message_stream.next().await {
let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
super::utils::emit_debug_trace(&model_config, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default());
log.write(&message, usage.as_ref().map(|f| &f.usage))?;
yield (message, usage);
}
}))
crates/goose/src/providers/openai.rs
+18 -9
View File
@@ -17,8 +17,7 @@ use super::embedding::{EmbeddingCapable, EmbeddingRequest, EmbeddingResponse};
use super::errors::ProviderError;
use super::formats::openai::{create_request, get_usage, response_to_message};
use super::utils::{
emit_debug_trace, get_model, handle_response_openai_compat, handle_status_openai_compat,
ImageFormat,
get_model, handle_response_openai_compat, handle_status_openai_compat, ImageFormat,
};
use crate::config::declarative_providers::DeclarativeProviderConfig;
use crate::conversation::message::Message;
@@ -26,6 +25,7 @@ use crate::conversation::message::Message;
use crate::model::ModelConfig;
use crate::providers::base::MessageStream;
use crate::providers::formats::openai::response_to_streaming_message;
use crate::providers::utils::RequestLog;
use rmcp::model::Tool;
pub const OPEN_AI_DEFAULT_MODEL: &str = "gpt-4o";
@@ -218,7 +218,10 @@ impl Provider for OpenAiProvider {
) -> Result<(Message, ProviderUsage), ProviderError> {
let payload = create_request(model_config, system, messages, tools, &ImageFormat::OpenAi)?;
let json_response = self.post(&payload).await?;
let mut log = RequestLog::start(&self.model, &payload)?;
let json_response = self.post(&payload).await.inspect_err(|e| {
let _ = log.error(e);
})?;
let message = response_to_message(&json_response)?;
let usage = json_response
@@ -229,7 +232,7 @@ impl Provider for OpenAiProvider {
Usage::default()
});
let model = get_model(&json_response);
emit_debug_trace(&self.model, &payload, &json_response, &usage);
log.write(&json_response, Some(&usage))?;
Ok((message, ProviderUsage::new(model, usage)))
}
@@ -282,17 +285,23 @@ impl Provider for OpenAiProvider {
payload["stream_options"] = json!({
"include_usage": true,
});
let mut log = RequestLog::start(&self.model, &payload)?;
let response = self
.api_client
.response_post(&self.base_path, &payload)
.await?;
let response = handle_status_openai_compat(response).await?;
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
let response = handle_status_openai_compat(response)
.await
.inspect_err(|e| {
let _ = log.error(e);
})?;
let stream = response.bytes_stream().map_err(io::Error::other);
let model_config = self.model.clone();
Ok(Box::pin(try_stream! {
let stream_reader = StreamReader::new(stream);
let framed = FramedRead::new(stream_reader, LinesCodec::new()).map_err(anyhow::Error::from);
@@ -301,7 +310,7 @@ impl Provider for OpenAiProvider {
pin!(message_stream);
while let Some(message) = message_stream.next().await {
let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
emit_debug_trace(&model_config, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default());
log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
yield (message, usage);
}
}))
crates/goose/src/providers/openrouter.rs
+4 -3
View File
@@ -7,8 +7,8 @@ use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{
emit_debug_trace, get_model, handle_response_google_compat, handle_response_openai_compat,
is_google_model,
get_model, handle_response_google_compat, handle_response_openai_compat, is_google_model,
RequestLog,
};
use crate::conversation::message::Message;
@@ -259,6 +259,7 @@ impl Provider for OpenRouterProvider {
) -> Result<(Message, ProviderUsage), ProviderError> {
// Create the base payload
let payload = create_request_based_on_model(self, system, messages, tools)?;
let mut log = RequestLog::start(model_config, &payload)?;
// Make request
let response = self
@@ -275,7 +276,7 @@ impl Provider for OpenRouterProvider {
Usage::default()
});
let response_model = get_model(&response);
emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
crates/goose/src/providers/sagemaker_tgi.rs
+5 -6
View File
@@ -12,7 +12,7 @@ use serde_json::{json, Value};
use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage, Usage};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::emit_debug_trace;
use super::utils::RequestLog;
use crate::conversation::message::{Message, MessageContent};
use crate::model::ModelConfig;
@@ -312,12 +312,11 @@ impl Provider for SageMakerTgiProvider {
"messages": messages,
"tools": tools
});
emit_debug_trace(
&self.model,
&debug_payload,
let mut log = RequestLog::start(&self.model, &debug_payload)?;
log.write(
&serde_json::to_value(&message).unwrap_or_default(),
&usage,
);
Some(&usage),
)?;
let provider_usage = ProviderUsage::new(model_name.to_string(), usage);
Ok((message, provider_usage))
crates/goose/src/providers/snowflake.rs
+5 -3
View File
@@ -8,7 +8,7 @@ use super::base::{ConfigKey, Provider, ProviderMetadata, ProviderUsage};
use super::errors::ProviderError;
use super::formats::snowflake::{create_request, get_usage, response_to_message};
use super::retry::ProviderRetry;
use super::utils::{get_model, map_http_error_to_provider_error, ImageFormat};
use super::utils::{get_model, map_http_error_to_provider_error, ImageFormat, RequestLog};
use crate::config::ConfigError;
use crate::conversation::message::Message;
@@ -309,6 +309,8 @@ impl Provider for SnowflakeProvider {
) -> Result<(Message, ProviderUsage), ProviderError> {
let payload = create_request(model_config, system, messages, tools)?;
let mut log = RequestLog::start(&self.model, &payload)?;
let response = self
.with_retry(|| async {
let payload_clone = payload.clone();
@@ -316,11 +318,11 @@ impl Provider for SnowflakeProvider {
})
.await?;
// Parse response
let message = response_to_message(&response)?;
let usage = get_usage(&response)?;
let response_model = get_model(&response);
super::utils::emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
crates/goose/src/providers/tetrate.rs
+6 -5
View File
@@ -15,8 +15,8 @@ use super::errors::ProviderError;
use super::formats::openai::response_to_streaming_message;
use super::retry::ProviderRetry;
use super::utils::{
emit_debug_trace, get_model, handle_response_google_compat, handle_response_openai_compat,
handle_status_openai_compat, is_google_model,
get_model, handle_response_google_compat, handle_response_openai_compat,
handle_status_openai_compat, is_google_model, RequestLog,
};
use crate::config::signup_tetrate::TETRATE_DEFAULT_MODEL;
use crate::conversation::message::Message;
@@ -173,6 +173,7 @@ impl Provider for TetrateProvider {
tools,
&super::utils::ImageFormat::OpenAi,
)?;
let mut log = RequestLog::start(model_config, &payload)?;
// Make request
let response = self
@@ -189,7 +190,7 @@ impl Provider for TetrateProvider {
Usage::default()
});
let model = get_model(&response);
emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(model, usage)))
}
@@ -220,7 +221,7 @@ impl Provider for TetrateProvider {
let response = handle_status_openai_compat(response).await?;
let stream = response.bytes_stream().map_err(io::Error::other);
let model_config = self.model.clone();
let mut log = RequestLog::start(&self.model, &payload)?;
Ok(Box::pin(try_stream! {
let stream_reader = StreamReader::new(stream);
@@ -230,7 +231,7 @@ impl Provider for TetrateProvider {
pin!(message_stream);
while let Some(message) = message_stream.next().await {
let (message, usage) = message.map_err(|e| ProviderError::RequestFailed(format!("Stream decode error: {}", e)))?;
emit_debug_trace(&model_config, &payload, &message, &usage.as_ref().map(|f| f.usage).unwrap_or_default());
log.write(&message, usage.as_ref().map(|f| f.usage).as_ref())?;
yield (message, usage);
}
}))
crates/goose/src/providers/utils.rs
+89 -31
View File
@@ -3,16 +3,19 @@ use super::errors::GoogleErrorCode;
use crate::config::paths::Paths;
use crate::model::ModelConfig;
use crate::providers::errors::{OpenAIError, ProviderError};
use anyhow::Result;
use anyhow::{anyhow, Result};
use base64::Engine;
use regex::Regex;
use reqwest::{Response, StatusCode};
use rmcp::model::{AnnotateAble, ImageContent, RawImageContent};
use serde::{Deserialize, Serialize};
use serde_json::{json, Map, Value};
use std::io::Read;
use std::path::Path;
use std::fmt::Display;
use std::fs::File;
use std::io::{BufWriter, Read, Write};
use std::path::{Path, PathBuf};
use std::time::Duration;
use uuid::Uuid;
#[derive(serde::Deserialize)]
struct OpenAIErrorResponse {
@@ -452,40 +455,95 @@ pub fn unescape_json_values(value: &Value) -> Value {
}
}
pub fn emit_debug_trace<T1, T2>(
model_config: &ModelConfig,
payload: &T1,
response: &T2,
usage: &Usage,
) where
T1: ?Sized + Serialize,
T2: ?Sized + Serialize,
{
let logs_dir = Paths::in_state_dir("logs");
pub struct RequestLog {
writer: Option<BufWriter<File>>,
temp_path: PathBuf,
}
if let Err(e) = std::fs::create_dir_all(&logs_dir) {
tracing::warn!("Failed to create logs directory: {}", e);
return;
const LOGS_TO_KEEP: usize = 10;
impl RequestLog {
pub fn start<Payload>(model_config: &ModelConfig, payload: &Payload) -> Result<Self>
where
Payload: Serialize,
{
let logs_dir = Paths::in_state_dir("logs");
std::fs::create_dir_all(&logs_dir)?;
let request_id = Uuid::new_v4();
let temp_name = format!("llm_request.{request_id}.jsonl");
let temp_path = logs_dir.join(PathBuf::from(temp_name));
let mut writer = BufWriter::new(
File::options()
.write(true)
.create(true)
.truncate(true)
.open(&temp_path)?,
);
let data = serde_json::json!({
"model_config": model_config,
"input": payload,
});
writeln!(writer, "{}", serde_json::to_string(&data)?)?;
Ok(Self {
writer: Some(writer),
temp_path,
})
}
let log_path = |i| logs_dir.join(format!("llm_request.{}.json", i));
for i in (0..4).rev() {
let _ = std::fs::rename(log_path(i), log_path(i + 1));
fn write_json(&mut self, line: &serde_json::Value) -> Result<()> {
let writer = self
.writer
.as_mut()
.ok_or_else(|| anyhow!("logger is finished"))?;
writeln!(writer, "{}", serde_json::to_string(line)?)?;
Ok(())
}
let data = serde_json::json!({
"model_config": model_config,
"input": payload,
"output": response,
"usage": usage,
});
pub fn error<E>(&mut self, error: E) -> Result<()>
where
E: Display,
{
self.write_json(&serde_json::json!({
"error": format!("{}", error),
}))
}
if let Err(e) = std::fs::write(
log_path(0),
serde_json::to_string_pretty(&data).unwrap_or_default(),
) {
tracing::warn!("Failed to write log file: {}", e);
pub fn write<Payload>(&mut self, data: &Payload, usage: Option<&Usage>) -> Result<()>
where
Payload: Serialize,
{
self.write_json(&serde_json::json!({
"data": data,
"usage": usage,
}))
}
fn finish(&mut self) -> Result<()> {
if let Some(mut writer) = self.writer.take() {
writer.flush()?;
let logs_dir = Paths::in_state_dir("logs");
let log_path = |i| logs_dir.join(format!("llm_request.{}.jsonl", i));
for i in (0..LOGS_TO_KEEP - 1).rev() {
let _ = std::fs::rename(log_path(i), log_path(i + 1));
}
std::fs::rename(&self.temp_path, log_path(0))?;
}
Ok(())
}
}
impl Drop for RequestLog {
fn drop(&mut self) {
if std::thread::panicking() {
return;
}
let _ = self.finish();
}
}
crates/goose/src/providers/xai.rs
+3 -2
View File
@@ -1,7 +1,7 @@
use super::api_client::{ApiClient, AuthMethod};
use super::errors::ProviderError;
use super::retry::ProviderRetry;
use super::utils::{get_model, handle_response_openai_compat};
use super::utils::{get_model, handle_response_openai_compat, RequestLog};
use crate::conversation::message::Message;
use crate::model::ModelConfig;
@@ -110,6 +110,7 @@ impl Provider for XaiProvider {
&super::utils::ImageFormat::OpenAi,
)?;
let mut log = RequestLog::start(&self.model, &payload)?;
let response = self.with_retry(|| self.post(payload.clone())).await?;
let message = response_to_message(&response)?;
@@ -118,7 +119,7 @@ impl Provider for XaiProvider {
Usage::default()
});
let response_model = get_model(&response);
super::utils::emit_debug_trace(model_config, &payload, &response, &usage);
log.write(&response, Some(&usage))?;
Ok((message, ProviderUsage::new(response_model, usage)))
}
}