Jack Amadeo
451 lines
15 KiB
Rust
451 lines
15 KiB
Rust
use crate::config::paths::Paths;
|
|
use anyhow::{anyhow, Result};
|
|
use fs_err::File;
|
|
use goose_providers::errors::{GoogleErrorCode, ProviderError};
|
|
use goose_providers::request_log::{install_logger, RequestLogHandle, RequestLogger};
|
|
use reqwest::{Response, StatusCode};
|
|
use serde_json::Value;
|
|
use std::error::Error;
|
|
use std::io::{BufWriter, Write};
|
|
use std::path::PathBuf;
|
|
use std::sync::OnceLock;
|
|
use std::time::Duration;
|
|
use uuid::Uuid;
|
|
|
|
pub fn filter_extensions_from_system_prompt(system: &str) -> String {
|
|
let Some(extensions_start) = system.find("# Extensions") else {
|
|
return system.to_string();
|
|
};
|
|
|
|
let Some(after_extensions) = system.get(extensions_start + 1..) else {
|
|
return system.to_string();
|
|
};
|
|
|
|
if let Some(next_section_pos) = after_extensions.find("\n# ") {
|
|
let Some(before) = system.get(..extensions_start) else {
|
|
return system.to_string();
|
|
};
|
|
let Some(after) = system.get(extensions_start + next_section_pos + 1..) else {
|
|
return system.to_string();
|
|
};
|
|
format!("{}{}", before.trim_end(), after)
|
|
} else {
|
|
system
|
|
.get(..extensions_start)
|
|
.map(|s| s.trim_end().to_string())
|
|
.unwrap_or_else(|| system.to_string())
|
|
}
|
|
}
|
|
|
|
fn format_server_error_message(status_code: StatusCode, payload: Option<&Value>) -> String {
|
|
match payload {
|
|
Some(Value::Null) | None => format!(
|
|
"HTTP {}: No response body received from server",
|
|
status_code.as_u16()
|
|
),
|
|
Some(p) => format!("HTTP {}: {}", status_code.as_u16(), p),
|
|
}
|
|
}
|
|
|
|
pub fn is_google_model(payload: &Value) -> bool {
|
|
payload
|
|
.get("model")
|
|
.and_then(|m| m.as_str())
|
|
.unwrap_or("")
|
|
.to_lowercase()
|
|
.contains("google")
|
|
}
|
|
|
|
/// Extracts `StatusCode` from response status or payload error code.
|
|
/// This function first checks the status code of the response. If the status is successful (2xx),
|
|
/// it then checks the payload for any error codes and maps them to appropriate `StatusCode`.
|
|
/// If the status is not successful (e.g., 4xx or 5xx), the original status code is returned.
|
|
fn get_google_final_status(status: StatusCode, payload: Option<&Value>) -> StatusCode {
|
|
// If the status is successful, check for an error in the payload
|
|
if status.is_success() {
|
|
if let Some(payload) = payload {
|
|
if let Some(error) = payload.get("error") {
|
|
if let Some(code) = error.get("code").and_then(|c| c.as_u64()) {
|
|
if let Some(google_error) = GoogleErrorCode::from_code(code) {
|
|
return google_error.to_status_code();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
status
|
|
}
|
|
|
|
fn parse_google_retry_delay(payload: &Value) -> Option<Duration> {
|
|
payload
|
|
.get("error")
|
|
.and_then(|error| error.get("details"))
|
|
.and_then(|details| details.as_array())
|
|
.and_then(|details_array| {
|
|
details_array.iter().find_map(|detail| {
|
|
if detail
|
|
.get("@type")
|
|
.and_then(|t| t.as_str())
|
|
.is_some_and(|s| s.ends_with("RetryInfo"))
|
|
{
|
|
detail
|
|
.get("retryDelay")
|
|
.and_then(|delay| delay.as_str())
|
|
.and_then(|s| s.strip_suffix('s'))
|
|
.and_then(|num| num.parse::<u64>().ok())
|
|
.map(Duration::from_secs)
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
})
|
|
}
|
|
|
|
/// Handle response from Google Gemini API-compatible endpoints.
|
|
///
|
|
/// Processes HTTP responses, handling specific statuses and parsing the payload
|
|
/// for error messages. Logs the response payload for debugging purposes.
|
|
///
|
|
/// ### References
|
|
/// - Error Codes: https://ai.google.dev/gemini-api/docs/troubleshooting?lang=python
|
|
///
|
|
/// ### Arguments
|
|
/// - `response`: The HTTP response to process.
|
|
///
|
|
/// ### Returns
|
|
/// - `Ok(Value)`: Parsed JSON on success.
|
|
/// - `Err(ProviderError)`: Describes the failure reason.
|
|
pub async fn handle_response_google_compat(response: Response) -> Result<Value, ProviderError> {
|
|
let status = response.status();
|
|
let url = super::http_status::sanitize_url(response.url().as_str());
|
|
let payload: Option<Value> = response.json().await.ok();
|
|
let final_status = get_google_final_status(status, payload.as_ref());
|
|
|
|
match final_status {
|
|
StatusCode::OK => payload.ok_or_else( || ProviderError::RequestFailed("Response body is not valid JSON".to_string()) ),
|
|
StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
|
|
Err(ProviderError::Authentication(format!("Authentication failed for {url}. Please ensure your API keys are valid and have the required permissions. \
|
|
Status: {}. Response: {:?}", final_status, payload )))
|
|
}
|
|
StatusCode::BAD_REQUEST | StatusCode::NOT_FOUND => {
|
|
let mut error_msg = "Unknown error".to_string();
|
|
if let Some(payload) = &payload {
|
|
if let Some(error) = payload.get("error") {
|
|
error_msg = error.get("message").and_then(|m| m.as_str()).unwrap_or("Unknown error").to_string();
|
|
let error_status = error.get("status").and_then(|s| s.as_str()).unwrap_or("Unknown status");
|
|
if error_status == "INVALID_ARGUMENT"
|
|
&& goose_providers::http_status::is_context_length_exceeded_message(&error_msg)
|
|
{
|
|
return Err(ProviderError::ContextLengthExceeded(error_msg.to_string()));
|
|
}
|
|
}
|
|
}
|
|
tracing::debug!(
|
|
"{}", format!("Provider request failed with status: {}. Payload: {:?}", final_status, payload)
|
|
);
|
|
Err(ProviderError::RequestFailed(format!("Request failed with status {} at {url}. Message: {}", final_status, error_msg)))
|
|
}
|
|
StatusCode::TOO_MANY_REQUESTS => {
|
|
let retry_delay = payload.as_ref().and_then(parse_google_retry_delay);
|
|
Err(ProviderError::RateLimitExceeded {
|
|
details: format!("{:?}", payload),
|
|
retry_delay,
|
|
})
|
|
}
|
|
_ if final_status.is_server_error() => Err(ProviderError::ServerError(
|
|
format!("Server error ({}) at {url}: {}", final_status, format_server_error_message(final_status, payload.as_ref())),
|
|
)),
|
|
_ => {
|
|
tracing::debug!(
|
|
"{}", format!("Provider request failed with status: {}. Payload: {:?}", final_status, payload)
|
|
);
|
|
Err(ProviderError::RequestFailed(format!("Request failed with status {} at {url}", final_status)))
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Extract the model name from a JSON object. Common with most providers to have this top level attribute.
|
|
pub fn get_model(data: &Value) -> String {
|
|
if let Some(model) = data.get("model") {
|
|
if let Some(model_str) = model.as_str() {
|
|
model_str.to_string()
|
|
} else {
|
|
"Unknown".to_string()
|
|
}
|
|
} else {
|
|
"Unknown".to_string()
|
|
}
|
|
}
|
|
|
|
pub fn unescape_json_values(value: &Value) -> Value {
|
|
let mut cloned = value.clone();
|
|
unescape_json_values_in_place(&mut cloned);
|
|
cloned
|
|
}
|
|
|
|
fn unescape_json_values_in_place(value: &mut Value) {
|
|
match value {
|
|
Value::Object(map) => {
|
|
for v in map.values_mut() {
|
|
unescape_json_values_in_place(v);
|
|
}
|
|
}
|
|
Value::Array(arr) => {
|
|
for v in arr.iter_mut() {
|
|
unescape_json_values_in_place(v);
|
|
}
|
|
}
|
|
Value::String(s) => {
|
|
if s.contains('\\') {
|
|
*s = s
|
|
.replace("\\\\n", "\n")
|
|
.replace("\\\\t", "\t")
|
|
.replace("\\\\r", "\r")
|
|
.replace("\\\\\"", "\"")
|
|
.replace("\\n", "\n")
|
|
.replace("\\t", "\t")
|
|
.replace("\\r", "\r")
|
|
.replace("\\\"", "\"");
|
|
}
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
|
|
pub const LOGS_TO_KEEP: usize = 10;
|
|
|
|
static INIT_LOGGER: OnceLock<Result<()>> = OnceLock::new();
|
|
|
|
pub fn init_goose_request_log() -> Result<()> {
|
|
INIT_LOGGER
|
|
.get_or_init(|| Ok(install_logger(RequestLog::new(LOGS_TO_KEEP)?)?))
|
|
.as_ref()
|
|
.map_err(|e| anyhow::anyhow!("failed to set up logger: {}", e))?;
|
|
Ok(())
|
|
}
|
|
|
|
pub struct RequestLog {
|
|
logs_to_keep: usize,
|
|
}
|
|
|
|
impl RequestLog {
|
|
pub fn new(logs_to_keep: usize) -> Result<Self> {
|
|
let logs_dir = Paths::in_state_dir("logs");
|
|
fs_err::create_dir_all(&logs_dir)?;
|
|
Ok(Self { logs_to_keep })
|
|
}
|
|
}
|
|
|
|
struct FileLogHandle {
|
|
writer: Option<BufWriter<File>>,
|
|
temp_path: PathBuf,
|
|
logs_to_keep: usize,
|
|
}
|
|
|
|
impl RequestLogger for RequestLog {
|
|
fn start(&self) -> Result<Box<dyn RequestLogHandle>, Box<dyn Error + Send + Sync>> {
|
|
let logs_dir = Paths::in_state_dir("logs");
|
|
fs_err::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 writer = BufWriter::new(
|
|
File::options()
|
|
.write(true)
|
|
.create(true)
|
|
.truncate(true)
|
|
.open(&temp_path)?,
|
|
);
|
|
|
|
Ok(Box::new(FileLogHandle {
|
|
writer: Some(writer),
|
|
temp_path,
|
|
logs_to_keep: self.logs_to_keep,
|
|
}))
|
|
}
|
|
}
|
|
|
|
impl RequestLogHandle for FileLogHandle {
|
|
fn write(&mut self, s: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
|
|
let writer = self
|
|
.writer
|
|
.as_mut()
|
|
.ok_or_else(|| anyhow!("logger is finished"))?;
|
|
writeln!(writer, "{}", s)?;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl FileLogHandle {
|
|
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));
|
|
|
|
if self.logs_to_keep == 0 {
|
|
fs_err::remove_file(&self.temp_path)?;
|
|
return Ok(());
|
|
}
|
|
|
|
for i in (0..self.logs_to_keep.saturating_sub(1)).rev() {
|
|
let _ = fs_err::rename(log_path(i), log_path(i + 1));
|
|
}
|
|
|
|
fs_err::rename(&self.temp_path, log_path(0))?;
|
|
}
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
impl Drop for FileLogHandle {
|
|
fn drop(&mut self) {
|
|
if std::thread::panicking() {
|
|
return;
|
|
}
|
|
let _ = self.finish();
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use serde_json::json;
|
|
|
|
#[test]
|
|
fn unescape_json_values_with_object() {
|
|
let value = json!({"text": "Hello\\nWorld"});
|
|
let unescaped_value = unescape_json_values(&value);
|
|
assert_eq!(unescaped_value, json!({"text": "Hello\nWorld"}));
|
|
}
|
|
|
|
#[test]
|
|
fn unescape_json_values_with_array() {
|
|
let value = json!(["Hello\\nWorld", "Goodbye\\tWorld"]);
|
|
let unescaped_value = unescape_json_values(&value);
|
|
assert_eq!(unescaped_value, json!(["Hello\nWorld", "Goodbye\tWorld"]));
|
|
}
|
|
|
|
#[test]
|
|
fn unescape_json_values_with_string() {
|
|
let value = json!("Hello\\nWorld");
|
|
let unescaped_value = unescape_json_values(&value);
|
|
assert_eq!(unescaped_value, json!("Hello\nWorld"));
|
|
}
|
|
|
|
#[test]
|
|
fn unescape_json_values_with_mixed_content() {
|
|
let value = json!({
|
|
"text": "Hello\\nWorld\\\\n!",
|
|
"array": ["Goodbye\\tWorld", "See you\\rlater"],
|
|
"nested": {
|
|
"inner_text": "Inner\\\"Quote\\\""
|
|
}
|
|
});
|
|
let unescaped_value = unescape_json_values(&value);
|
|
assert_eq!(
|
|
unescaped_value,
|
|
json!({
|
|
"text": "Hello\nWorld\n!",
|
|
"array": ["Goodbye\tWorld", "See you\rlater"],
|
|
"nested": {
|
|
"inner_text": "Inner\"Quote\""
|
|
}
|
|
})
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn unescape_json_values_with_no_escapes() {
|
|
let value = json!({"text": "Hello World"});
|
|
let unescaped_value = unescape_json_values(&value);
|
|
assert_eq!(unescaped_value, json!({"text": "Hello World"}));
|
|
}
|
|
|
|
#[test]
|
|
fn test_is_google_model() {
|
|
// Define the test cases as a vector of tuples
|
|
let test_cases = vec![
|
|
// (input, expected_result)
|
|
(json!({ "model": "google_gemini" }), true),
|
|
(json!({ "model": "microsoft_bing" }), false),
|
|
(json!({ "model": "" }), false),
|
|
(json!({}), false),
|
|
(json!({ "model": "Google_XYZ" }), true),
|
|
(json!({ "model": "google_abc" }), true),
|
|
];
|
|
|
|
// Iterate through each test case and assert the result
|
|
for (payload, expected_result) in test_cases {
|
|
assert_eq!(is_google_model(&payload), expected_result);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_google_final_status_success() {
|
|
let status = StatusCode::OK;
|
|
let payload = json!({});
|
|
let result = get_google_final_status(status, Some(&payload));
|
|
assert_eq!(result, StatusCode::OK);
|
|
}
|
|
|
|
#[test]
|
|
fn test_get_google_final_status_with_error_code() {
|
|
// Test error code mappings for different payload error codes
|
|
let test_cases = vec![
|
|
// (error code, status, expected status code)
|
|
(200, None, StatusCode::OK),
|
|
(429, Some(StatusCode::OK), StatusCode::TOO_MANY_REQUESTS),
|
|
(400, Some(StatusCode::OK), StatusCode::BAD_REQUEST),
|
|
(401, Some(StatusCode::OK), StatusCode::UNAUTHORIZED),
|
|
(403, Some(StatusCode::OK), StatusCode::FORBIDDEN),
|
|
(404, Some(StatusCode::OK), StatusCode::NOT_FOUND),
|
|
(500, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR),
|
|
(503, Some(StatusCode::OK), StatusCode::SERVICE_UNAVAILABLE),
|
|
(999, Some(StatusCode::OK), StatusCode::INTERNAL_SERVER_ERROR),
|
|
(500, Some(StatusCode::BAD_REQUEST), StatusCode::BAD_REQUEST),
|
|
(
|
|
404,
|
|
Some(StatusCode::INTERNAL_SERVER_ERROR),
|
|
StatusCode::INTERNAL_SERVER_ERROR,
|
|
),
|
|
];
|
|
|
|
for (error_code, status, expected_status) in test_cases {
|
|
let payload = if let Some(_status) = status {
|
|
json!({
|
|
"error": {
|
|
"code": error_code,
|
|
"message": "Error message"
|
|
}
|
|
})
|
|
} else {
|
|
json!({})
|
|
};
|
|
|
|
let result = get_google_final_status(status.unwrap_or(StatusCode::OK), Some(&payload));
|
|
assert_eq!(result, expected_status);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_parse_google_retry_delay() {
|
|
let payload = json!({
|
|
"error": {
|
|
"details": [
|
|
{
|
|
"@type": "type.googleapis.com/google.rpc.RetryInfo",
|
|
"retryDelay": "42s"
|
|
}
|
|
]
|
|
}
|
|
});
|
|
assert_eq!(
|
|
parse_google_retry_delay(&payload),
|
|
Some(Duration::from_secs(42))
|
|
);
|
|
}
|
|
}
|