mirror of
https://github.com/katanemo/plano.git
synced 2026-06-29 15:49:40 +02:00
merge main into musa/custom-trace-attributes
Co-authored-by: Cursor <cursoragent@cursor.com>
This commit is contained in:
commit
edea84b5c8
7 changed files with 280 additions and 161 deletions
3
crates/Cargo.lock
generated
3
crates/Cargo.lock
generated
|
|
@ -436,11 +436,14 @@ name = "common"
|
||||||
version = "0.1.0"
|
version = "0.1.0"
|
||||||
dependencies = [
|
dependencies = [
|
||||||
"axum",
|
"axum",
|
||||||
|
"bytes",
|
||||||
"derivative",
|
"derivative",
|
||||||
"duration-string",
|
"duration-string",
|
||||||
"governor",
|
"governor",
|
||||||
"hermesllm",
|
"hermesllm",
|
||||||
"hex",
|
"hex",
|
||||||
|
"http-body-util",
|
||||||
|
"hyper 1.6.0",
|
||||||
"log",
|
"log",
|
||||||
"pretty_assertions",
|
"pretty_assertions",
|
||||||
"proxy-wasm",
|
"proxy-wasm",
|
||||||
|
|
|
||||||
|
|
@ -3,6 +3,7 @@ use std::time::Instant;
|
||||||
|
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
use common::configuration::SpanAttributes;
|
use common::configuration::SpanAttributes;
|
||||||
|
use common::errors::BrightStaffError;
|
||||||
use common::llm_providers::LlmProviders;
|
use common::llm_providers::LlmProviders;
|
||||||
use hermesllm::apis::OpenAIMessage;
|
use hermesllm::apis::OpenAIMessage;
|
||||||
use hermesllm::clients::SupportedAPIsFromClient;
|
use hermesllm::clients::SupportedAPIsFromClient;
|
||||||
|
|
@ -25,12 +26,12 @@ use crate::tracing::{collect_custom_trace_attributes, operation_component, set_s
|
||||||
/// Main errors for agent chat completions
|
/// Main errors for agent chat completions
|
||||||
#[derive(Debug, thiserror::Error)]
|
#[derive(Debug, thiserror::Error)]
|
||||||
pub enum AgentFilterChainError {
|
pub enum AgentFilterChainError {
|
||||||
|
#[error("Forwarded error: {0}")]
|
||||||
|
Brightstaff(#[from] BrightStaffError),
|
||||||
#[error("Agent selection error: {0}")]
|
#[error("Agent selection error: {0}")]
|
||||||
Selection(#[from] AgentSelectionError),
|
Selection(#[from] AgentSelectionError),
|
||||||
#[error("Pipeline processing error: {0}")]
|
#[error("Pipeline processing error: {0}")]
|
||||||
Pipeline(#[from] PipelineError),
|
Pipeline(#[from] PipelineError),
|
||||||
#[error("Response handling error: {0}")]
|
|
||||||
Response(#[from] super::response_handler::ResponseError),
|
|
||||||
#[error("Request parsing error: {0}")]
|
#[error("Request parsing error: {0}")]
|
||||||
RequestParsing(#[from] serde_json::Error),
|
RequestParsing(#[from] serde_json::Error),
|
||||||
#[error("HTTP error: {0}")]
|
#[error("HTTP error: {0}")]
|
||||||
|
|
@ -108,16 +109,15 @@ pub async fn agent_chat(
|
||||||
"agent_response": body
|
"agent_response": body
|
||||||
});
|
});
|
||||||
|
|
||||||
|
let status_code = hyper::StatusCode::from_u16(*status)
|
||||||
|
.unwrap_or(hyper::StatusCode::INTERNAL_SERVER_ERROR);
|
||||||
|
|
||||||
let json_string = error_json.to_string();
|
let json_string = error_json.to_string();
|
||||||
let mut response =
|
return Ok(BrightStaffError::ForwardedError {
|
||||||
Response::new(ResponseHandler::create_full_body(json_string));
|
status_code,
|
||||||
*response.status_mut() = hyper::StatusCode::from_u16(*status)
|
message: json_string,
|
||||||
.unwrap_or(hyper::StatusCode::BAD_REQUEST);
|
}
|
||||||
response.headers_mut().insert(
|
.into_response());
|
||||||
hyper::header::CONTENT_TYPE,
|
|
||||||
"application/json".parse().unwrap(),
|
|
||||||
);
|
|
||||||
return Ok(response);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Print detailed error information with full error chain for other errors
|
// Print detailed error information with full error chain for other errors
|
||||||
|
|
@ -150,8 +150,11 @@ pub async fn agent_chat(
|
||||||
// Log the error for debugging
|
// Log the error for debugging
|
||||||
info!(error = %error_json, "structured error info");
|
info!(error = %error_json, "structured error info");
|
||||||
|
|
||||||
// Return JSON error response
|
Ok(BrightStaffError::ForwardedError {
|
||||||
Ok(ResponseHandler::create_json_error_response(&error_json))
|
status_code: StatusCode::BAD_REQUEST,
|
||||||
|
message: error_json.to_string(),
|
||||||
|
}
|
||||||
|
.into_response())
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
@ -258,10 +261,7 @@ async fn handle_agent_chat_inner(
|
||||||
None => {
|
None => {
|
||||||
let err_msg = "No model specified in request and no default provider configured";
|
let err_msg = "No model specified in request and no default provider configured";
|
||||||
warn!("{}", err_msg);
|
warn!("{}", err_msg);
|
||||||
let mut bad_request =
|
return Ok(BrightStaffError::NoModelSpecified.into_response());
|
||||||
Response::new(ResponseHandler::create_full_body(err_msg.to_string()));
|
|
||||||
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
|
||||||
return Ok(bad_request);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -5,9 +5,10 @@ use hyper::header::HeaderMap;
|
||||||
|
|
||||||
use crate::handlers::agent_selector::{AgentSelectionError, AgentSelector};
|
use crate::handlers::agent_selector::{AgentSelectionError, AgentSelector};
|
||||||
use crate::handlers::pipeline_processor::PipelineProcessor;
|
use crate::handlers::pipeline_processor::PipelineProcessor;
|
||||||
use crate::handlers::response_handler::ResponseHandler;
|
|
||||||
use crate::router::plano_orchestrator::OrchestratorService;
|
use crate::router::plano_orchestrator::OrchestratorService;
|
||||||
|
use common::errors::BrightStaffError;
|
||||||
|
use http_body_util::BodyExt;
|
||||||
|
use hyper::StatusCode;
|
||||||
/// Integration test that demonstrates the modular agent chat flow
|
/// Integration test that demonstrates the modular agent chat flow
|
||||||
/// This test shows how the three main components work together:
|
/// This test shows how the three main components work together:
|
||||||
/// 1. AgentSelector - selects the appropriate agents based on orchestration
|
/// 1. AgentSelector - selects the appropriate agents based on orchestration
|
||||||
|
|
@ -128,8 +129,24 @@ mod tests {
|
||||||
}
|
}
|
||||||
|
|
||||||
// Test 4: Error Response Creation
|
// Test 4: Error Response Creation
|
||||||
let error_response = ResponseHandler::create_bad_request("Test error");
|
let err = BrightStaffError::ModelNotFound("gpt-5-secret".to_string());
|
||||||
assert_eq!(error_response.status(), hyper::StatusCode::BAD_REQUEST);
|
let response = err.into_response();
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||||
|
|
||||||
|
// Helper to extract body as JSON
|
||||||
|
let body_bytes = response.into_body().collect().await.unwrap().to_bytes();
|
||||||
|
let body: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(body["error"]["code"], "ModelNotFound");
|
||||||
|
assert_eq!(
|
||||||
|
body["error"]["details"]["rejected_model_id"],
|
||||||
|
"gpt-5-secret"
|
||||||
|
);
|
||||||
|
assert!(body["error"]["message"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.contains("gpt-5-secret"));
|
||||||
|
|
||||||
println!("✅ All modular components working correctly!");
|
println!("✅ All modular components working correctly!");
|
||||||
}
|
}
|
||||||
|
|
@ -148,12 +165,21 @@ mod tests {
|
||||||
AgentSelectionError::ListenerNotFound(_)
|
AgentSelectionError::ListenerNotFound(_)
|
||||||
));
|
));
|
||||||
|
|
||||||
// Test error response creation
|
let technical_reason = "Database connection timed out";
|
||||||
let error_response = ResponseHandler::create_internal_error("Pipeline failed");
|
let err = BrightStaffError::InternalServerError(technical_reason.to_string());
|
||||||
assert_eq!(
|
|
||||||
error_response.status(),
|
let response = err.into_response();
|
||||||
hyper::StatusCode::INTERNAL_SERVER_ERROR
|
|
||||||
);
|
// --- 1. EXTRACT BYTES ---
|
||||||
|
let body_bytes = response.into_body().collect().await.unwrap().to_bytes();
|
||||||
|
|
||||||
|
// --- 2. DECLARE body_json HERE ---
|
||||||
|
let body_json: serde_json::Value =
|
||||||
|
serde_json::from_slice(&body_bytes).expect("Failed to parse JSON body");
|
||||||
|
|
||||||
|
// --- 3. USE body_json ---
|
||||||
|
assert_eq!(body_json["error"]["code"], "InternalServerError");
|
||||||
|
assert_eq!(body_json["error"]["details"]["reason"], technical_reason);
|
||||||
|
|
||||||
println!("✅ Error handling working correctly!");
|
println!("✅ Error handling working correctly!");
|
||||||
}
|
}
|
||||||
|
|
|
||||||
|
|
@ -8,9 +8,9 @@ use hermesllm::apis::openai_responses::InputParam;
|
||||||
use hermesllm::clients::{SupportedAPIsFromClient, SupportedUpstreamAPIs};
|
use hermesllm::clients::{SupportedAPIsFromClient, SupportedUpstreamAPIs};
|
||||||
use hermesllm::{ProviderRequest, ProviderRequestType};
|
use hermesllm::{ProviderRequest, ProviderRequestType};
|
||||||
use http_body_util::combinators::BoxBody;
|
use http_body_util::combinators::BoxBody;
|
||||||
use http_body_util::{BodyExt, Full};
|
use http_body_util::BodyExt;
|
||||||
use hyper::header::{self};
|
use hyper::header::{self};
|
||||||
use hyper::{Request, Response, StatusCode};
|
use hyper::{Request, Response};
|
||||||
use opentelemetry::global;
|
use opentelemetry::global;
|
||||||
use opentelemetry::trace::get_active_span;
|
use opentelemetry::trace::get_active_span;
|
||||||
use opentelemetry_http::HeaderInjector;
|
use opentelemetry_http::HeaderInjector;
|
||||||
|
|
@ -32,11 +32,7 @@ use crate::tracing::{
|
||||||
collect_custom_trace_attributes, llm as tracing_llm, operation_component, set_service_name,
|
collect_custom_trace_attributes, llm as tracing_llm, operation_component, set_service_name,
|
||||||
};
|
};
|
||||||
|
|
||||||
fn full<T: Into<Bytes>>(chunk: T) -> BoxBody<Bytes, hyper::Error> {
|
use common::errors::BrightStaffError;
|
||||||
Full::new(chunk.into())
|
|
||||||
.map_err(|never| match never {})
|
|
||||||
.boxed()
|
|
||||||
}
|
|
||||||
|
|
||||||
pub async fn llm_chat(
|
pub async fn llm_chat(
|
||||||
request: Request<hyper::body::Incoming>,
|
request: Request<hyper::body::Incoming>,
|
||||||
|
|
@ -147,10 +143,11 @@ async fn llm_chat_inner(
|
||||||
error = %err,
|
error = %err,
|
||||||
"failed to parse request as ProviderRequestType"
|
"failed to parse request as ProviderRequestType"
|
||||||
);
|
);
|
||||||
let err_msg = format!("Failed to parse request: {}", err);
|
return Ok(BrightStaffError::InvalidRequest(format!(
|
||||||
let mut bad_request = Response::new(full(err_msg));
|
"Failed to parse request: {}",
|
||||||
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
err
|
||||||
return Ok(bad_request);
|
))
|
||||||
|
.into_response());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
@ -175,9 +172,7 @@ async fn llm_chat_inner(
|
||||||
None => {
|
None => {
|
||||||
let err_msg = "No model specified in request and no default provider configured";
|
let err_msg = "No model specified in request and no default provider configured";
|
||||||
warn!("{}", err_msg);
|
warn!("{}", err_msg);
|
||||||
let mut bad_request = Response::new(full(err_msg.to_string()));
|
return Ok(BrightStaffError::NoModelSpecified.into_response());
|
||||||
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
|
||||||
return Ok(bad_request);
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
} else {
|
} else {
|
||||||
|
|
@ -198,14 +193,8 @@ async fn llm_chat_inner(
|
||||||
.get(&alias_resolved_model)
|
.get(&alias_resolved_model)
|
||||||
.is_none()
|
.is_none()
|
||||||
{
|
{
|
||||||
let err_msg = format!(
|
|
||||||
"Model '{}' not found in configured providers",
|
|
||||||
alias_resolved_model
|
|
||||||
);
|
|
||||||
warn!(model = %alias_resolved_model, "model not found in configured providers");
|
warn!(model = %alias_resolved_model, "model not found in configured providers");
|
||||||
let mut bad_request = Response::new(full(err_msg));
|
return Ok(BrightStaffError::ModelNotFound(alias_resolved_model).into_response());
|
||||||
*bad_request.status_mut() = StatusCode::BAD_REQUEST;
|
|
||||||
return Ok(bad_request);
|
|
||||||
}
|
}
|
||||||
|
|
||||||
// Handle provider/model slug format (e.g., "openai/gpt-4")
|
// Handle provider/model slug format (e.g., "openai/gpt-4")
|
||||||
|
|
@ -300,13 +289,10 @@ async fn llm_chat_inner(
|
||||||
Err(StateStorageError::NotFound(_)) => {
|
Err(StateStorageError::NotFound(_)) => {
|
||||||
// Return 409 Conflict when previous_response_id not found
|
// Return 409 Conflict when previous_response_id not found
|
||||||
warn!(previous_response_id = %prev_resp_id, "previous response_id not found");
|
warn!(previous_response_id = %prev_resp_id, "previous response_id not found");
|
||||||
let err_msg = format!(
|
return Ok(BrightStaffError::ConversationStateNotFound(
|
||||||
"Conversation state not found for previous_response_id: {}",
|
prev_resp_id.to_string(),
|
||||||
prev_resp_id
|
)
|
||||||
);
|
.into_response());
|
||||||
let mut conflict_response = Response::new(full(err_msg));
|
|
||||||
*conflict_response.status_mut() = StatusCode::CONFLICT;
|
|
||||||
return Ok(conflict_response);
|
|
||||||
}
|
}
|
||||||
Err(e) => {
|
Err(e) => {
|
||||||
// Log warning but continue on other storage errors
|
// Log warning but continue on other storage errors
|
||||||
|
|
@ -357,9 +343,11 @@ async fn llm_chat_inner(
|
||||||
{
|
{
|
||||||
Ok(result) => result,
|
Ok(result) => result,
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
let mut internal_error = Response::new(full(err.message));
|
return Ok(BrightStaffError::ForwardedError {
|
||||||
*internal_error.status_mut() = err.status_code;
|
status_code: err.status_code,
|
||||||
return Ok(internal_error);
|
message: err.message,
|
||||||
|
}
|
||||||
|
.into_response());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
@ -425,10 +413,11 @@ async fn llm_chat_inner(
|
||||||
{
|
{
|
||||||
Ok(res) => res,
|
Ok(res) => res,
|
||||||
Err(err) => {
|
Err(err) => {
|
||||||
let err_msg = format!("Failed to send request: {}", err);
|
return Ok(BrightStaffError::InternalServerError(format!(
|
||||||
let mut internal_error = Response::new(full(err_msg));
|
"Failed to send request: {}",
|
||||||
*internal_error.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
|
err
|
||||||
return Ok(internal_error);
|
))
|
||||||
|
.into_response());
|
||||||
}
|
}
|
||||||
};
|
};
|
||||||
|
|
||||||
|
|
@ -485,12 +474,11 @@ async fn llm_chat_inner(
|
||||||
|
|
||||||
match response.body(streaming_response.body) {
|
match response.body(streaming_response.body) {
|
||||||
Ok(response) => Ok(response),
|
Ok(response) => Ok(response),
|
||||||
Err(err) => {
|
Err(err) => Ok(BrightStaffError::InternalServerError(format!(
|
||||||
let err_msg = format!("Failed to create response: {}", err);
|
"Failed to create response: {}",
|
||||||
let mut internal_error = Response::new(full(err_msg));
|
err
|
||||||
*internal_error.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
|
))
|
||||||
Ok(internal_error)
|
.into_response()),
|
||||||
}
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -1,25 +1,17 @@
|
||||||
use bytes::Bytes;
|
use bytes::Bytes;
|
||||||
|
use common::errors::BrightStaffError;
|
||||||
use hermesllm::apis::OpenAIApi;
|
use hermesllm::apis::OpenAIApi;
|
||||||
use hermesllm::clients::{SupportedAPIsFromClient, SupportedUpstreamAPIs};
|
use hermesllm::clients::{SupportedAPIsFromClient, SupportedUpstreamAPIs};
|
||||||
use hermesllm::SseEvent;
|
use hermesllm::SseEvent;
|
||||||
use http_body_util::combinators::BoxBody;
|
use http_body_util::combinators::BoxBody;
|
||||||
use http_body_util::{BodyExt, Full, StreamBody};
|
use http_body_util::{BodyExt, Full, StreamBody};
|
||||||
use hyper::body::Frame;
|
use hyper::body::Frame;
|
||||||
use hyper::{Response, StatusCode};
|
use hyper::Response;
|
||||||
use tokio::sync::mpsc;
|
use tokio::sync::mpsc;
|
||||||
use tokio_stream::wrappers::ReceiverStream;
|
use tokio_stream::wrappers::ReceiverStream;
|
||||||
use tokio_stream::StreamExt;
|
use tokio_stream::StreamExt;
|
||||||
use tracing::{info, warn, Instrument};
|
use tracing::{info, warn, Instrument};
|
||||||
|
|
||||||
/// Errors that can occur during response handling
|
|
||||||
#[derive(Debug, thiserror::Error)]
|
|
||||||
pub enum ResponseError {
|
|
||||||
#[error("Failed to create response: {0}")]
|
|
||||||
ResponseCreationFailed(#[from] hyper::http::Error),
|
|
||||||
#[error("Stream error: {0}")]
|
|
||||||
StreamError(String),
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Service for handling HTTP responses and streaming
|
/// Service for handling HTTP responses and streaming
|
||||||
pub struct ResponseHandler;
|
pub struct ResponseHandler;
|
||||||
|
|
||||||
|
|
@ -35,40 +27,6 @@ impl ResponseHandler {
|
||||||
.boxed()
|
.boxed()
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Create an error response with a given status code and message
|
|
||||||
pub fn create_error_response(
|
|
||||||
status: StatusCode,
|
|
||||||
message: &str,
|
|
||||||
) -> Response<BoxBody<Bytes, hyper::Error>> {
|
|
||||||
let mut response = Response::new(Self::create_full_body(message.to_string()));
|
|
||||||
*response.status_mut() = status;
|
|
||||||
response
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Create a bad request response
|
|
||||||
pub fn create_bad_request(message: &str) -> Response<BoxBody<Bytes, hyper::Error>> {
|
|
||||||
Self::create_error_response(StatusCode::BAD_REQUEST, message)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Create an internal server error response
|
|
||||||
pub fn create_internal_error(message: &str) -> Response<BoxBody<Bytes, hyper::Error>> {
|
|
||||||
Self::create_error_response(StatusCode::INTERNAL_SERVER_ERROR, message)
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Create a JSON error response
|
|
||||||
pub fn create_json_error_response(
|
|
||||||
error_json: &serde_json::Value,
|
|
||||||
) -> Response<BoxBody<Bytes, hyper::Error>> {
|
|
||||||
let json_string = error_json.to_string();
|
|
||||||
let mut response = Response::new(Self::create_full_body(json_string));
|
|
||||||
*response.status_mut() = StatusCode::INTERNAL_SERVER_ERROR;
|
|
||||||
response.headers_mut().insert(
|
|
||||||
hyper::header::CONTENT_TYPE,
|
|
||||||
"application/json".parse().unwrap(),
|
|
||||||
);
|
|
||||||
response
|
|
||||||
}
|
|
||||||
|
|
||||||
/// Create a streaming response from a reqwest response.
|
/// Create a streaming response from a reqwest response.
|
||||||
/// The spawned streaming task is instrumented with both `agent_span` and `orchestrator_span`
|
/// The spawned streaming task is instrumented with both `agent_span` and `orchestrator_span`
|
||||||
/// so their durations reflect the actual time spent streaming to the client.
|
/// so their durations reflect the actual time spent streaming to the client.
|
||||||
|
|
@ -77,13 +35,13 @@ impl ResponseHandler {
|
||||||
llm_response: reqwest::Response,
|
llm_response: reqwest::Response,
|
||||||
agent_span: tracing::Span,
|
agent_span: tracing::Span,
|
||||||
orchestrator_span: tracing::Span,
|
orchestrator_span: tracing::Span,
|
||||||
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, ResponseError> {
|
) -> Result<Response<BoxBody<Bytes, hyper::Error>>, BrightStaffError> {
|
||||||
// Copy headers from the original response
|
// Copy headers from the original response
|
||||||
let response_headers = llm_response.headers();
|
let response_headers = llm_response.headers();
|
||||||
let mut response_builder = Response::builder();
|
let mut response_builder = Response::builder();
|
||||||
|
|
||||||
let headers = response_builder.headers_mut().ok_or_else(|| {
|
let headers = response_builder.headers_mut().ok_or_else(|| {
|
||||||
ResponseError::StreamError("Failed to get mutable headers".to_string())
|
BrightStaffError::StreamError("Failed to get mutable headers".to_string())
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
for (header_name, header_value) in response_headers.iter() {
|
for (header_name, header_value) in response_headers.iter() {
|
||||||
|
|
@ -123,7 +81,7 @@ impl ResponseHandler {
|
||||||
|
|
||||||
response_builder
|
response_builder
|
||||||
.body(stream_body)
|
.body(stream_body)
|
||||||
.map_err(ResponseError::from)
|
.map_err(BrightStaffError::from)
|
||||||
}
|
}
|
||||||
|
|
||||||
/// Collect the full response body as a string
|
/// Collect the full response body as a string
|
||||||
|
|
@ -136,7 +94,7 @@ impl ResponseHandler {
|
||||||
pub async fn collect_full_response(
|
pub async fn collect_full_response(
|
||||||
&self,
|
&self,
|
||||||
llm_response: reqwest::Response,
|
llm_response: reqwest::Response,
|
||||||
) -> Result<String, ResponseError> {
|
) -> Result<String, BrightStaffError> {
|
||||||
use hermesllm::apis::streaming_shapes::sse::SseStreamIter;
|
use hermesllm::apis::streaming_shapes::sse::SseStreamIter;
|
||||||
|
|
||||||
let response_headers = llm_response.headers();
|
let response_headers = llm_response.headers();
|
||||||
|
|
@ -144,10 +102,9 @@ impl ResponseHandler {
|
||||||
.get(hyper::header::CONTENT_TYPE)
|
.get(hyper::header::CONTENT_TYPE)
|
||||||
.is_some_and(|v| v.to_str().unwrap_or("").contains("text/event-stream"));
|
.is_some_and(|v| v.to_str().unwrap_or("").contains("text/event-stream"));
|
||||||
|
|
||||||
let response_bytes = llm_response
|
let response_bytes = llm_response.bytes().await.map_err(|e| {
|
||||||
.bytes()
|
BrightStaffError::StreamError(format!("Failed to read response: {}", e))
|
||||||
.await
|
})?;
|
||||||
.map_err(|e| ResponseError::StreamError(format!("Failed to read response: {}", e)))?;
|
|
||||||
|
|
||||||
if is_sse_streaming {
|
if is_sse_streaming {
|
||||||
let client_api =
|
let client_api =
|
||||||
|
|
@ -185,7 +142,7 @@ impl ResponseHandler {
|
||||||
} else {
|
} else {
|
||||||
// If not SSE, treat as regular text response
|
// If not SSE, treat as regular text response
|
||||||
let response_text = String::from_utf8(response_bytes.to_vec()).map_err(|e| {
|
let response_text = String::from_utf8(response_bytes.to_vec()).map_err(|e| {
|
||||||
ResponseError::StreamError(format!("Failed to decode response: {}", e))
|
BrightStaffError::StreamError(format!("Failed to decode response: {}", e))
|
||||||
})?;
|
})?;
|
||||||
|
|
||||||
Ok(response_text)
|
Ok(response_text)
|
||||||
|
|
@ -204,42 +161,6 @@ mod tests {
|
||||||
use super::*;
|
use super::*;
|
||||||
use hyper::StatusCode;
|
use hyper::StatusCode;
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_create_bad_request() {
|
|
||||||
let response = ResponseHandler::create_bad_request("Invalid request");
|
|
||||||
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_create_internal_error() {
|
|
||||||
let response = ResponseHandler::create_internal_error("Server error");
|
|
||||||
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_create_error_response() {
|
|
||||||
let response =
|
|
||||||
ResponseHandler::create_error_response(StatusCode::NOT_FOUND, "Resource not found");
|
|
||||||
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[test]
|
|
||||||
fn test_create_json_error_response() {
|
|
||||||
let error_json = serde_json::json!({
|
|
||||||
"error": {
|
|
||||||
"type": "TestError",
|
|
||||||
"message": "Test error message"
|
|
||||||
}
|
|
||||||
});
|
|
||||||
|
|
||||||
let response = ResponseHandler::create_json_error_response(&error_json);
|
|
||||||
assert_eq!(response.status(), StatusCode::INTERNAL_SERVER_ERROR);
|
|
||||||
assert_eq!(
|
|
||||||
response.headers().get("content-type").unwrap(),
|
|
||||||
"application/json"
|
|
||||||
);
|
|
||||||
}
|
|
||||||
|
|
||||||
#[tokio::test]
|
#[tokio::test]
|
||||||
async fn test_create_streaming_response_with_mock() {
|
async fn test_create_streaming_response_with_mock() {
|
||||||
use mockito::Server;
|
use mockito::Server;
|
||||||
|
|
|
||||||
|
|
@ -20,6 +20,9 @@ urlencoding = "2.1.3"
|
||||||
url = "2.5.4"
|
url = "2.5.4"
|
||||||
hermesllm = { version = "0.1.0", path = "../hermesllm" }
|
hermesllm = { version = "0.1.0", path = "../hermesllm" }
|
||||||
serde_with = "3.13.0"
|
serde_with = "3.13.0"
|
||||||
|
hyper = "1.0"
|
||||||
|
bytes = "1.0"
|
||||||
|
http-body-util = "0.1"
|
||||||
|
|
||||||
[features]
|
[features]
|
||||||
default = []
|
default = []
|
||||||
|
|
@ -30,3 +33,6 @@ serde_json = "1.0.64"
|
||||||
serial_test = "3.2"
|
serial_test = "3.2"
|
||||||
axum = "0.7"
|
axum = "0.7"
|
||||||
tokio = { version = "1.44", features = ["sync", "time", "macros", "rt"] }
|
tokio = { version = "1.44", features = ["sync", "time", "macros", "rt"] }
|
||||||
|
hyper = { version = "1.0", features = ["full"] }
|
||||||
|
bytes = "1.0"
|
||||||
|
http-body-util = "0.1"
|
||||||
|
|
|
||||||
|
|
@ -1,9 +1,13 @@
|
||||||
use proxy_wasm::types::Status;
|
|
||||||
|
|
||||||
use crate::{api::open_ai::ChatCompletionChunkResponseError, ratelimit};
|
use crate::{api::open_ai::ChatCompletionChunkResponseError, ratelimit};
|
||||||
|
use bytes::Bytes;
|
||||||
use hermesllm::apis::openai::OpenAIError;
|
use hermesllm::apis::openai::OpenAIError;
|
||||||
|
use http_body_util::{combinators::BoxBody, BodyExt, Full};
|
||||||
|
use hyper::{Error as HyperError, Response, StatusCode};
|
||||||
|
use proxy_wasm::types::Status;
|
||||||
|
use serde_json::json;
|
||||||
|
use thiserror::Error;
|
||||||
|
|
||||||
#[derive(thiserror::Error, Debug)]
|
#[derive(Error, Debug)]
|
||||||
pub enum ClientError {
|
pub enum ClientError {
|
||||||
#[error("Error dispatching HTTP call to `{upstream_name}/{path}`, error: {internal_status:?}")]
|
#[error("Error dispatching HTTP call to `{upstream_name}/{path}`, error: {internal_status:?}")]
|
||||||
DispatchError {
|
DispatchError {
|
||||||
|
|
@ -13,7 +17,7 @@ pub enum ClientError {
|
||||||
},
|
},
|
||||||
}
|
}
|
||||||
|
|
||||||
#[derive(thiserror::Error, Debug)]
|
#[derive(Error, Debug)]
|
||||||
pub enum ServerError {
|
pub enum ServerError {
|
||||||
#[error(transparent)]
|
#[error(transparent)]
|
||||||
HttpDispatch(ClientError),
|
HttpDispatch(ClientError),
|
||||||
|
|
@ -43,3 +47,174 @@ pub enum ServerError {
|
||||||
#[error("error parsing openai message: {0}")]
|
#[error("error parsing openai message: {0}")]
|
||||||
OpenAIPError(#[from] OpenAIError),
|
OpenAIPError(#[from] OpenAIError),
|
||||||
}
|
}
|
||||||
|
// -----------------------------------------------------------------------------
|
||||||
|
// BrightStaff Errors (Standardized)
|
||||||
|
// -----------------------------------------------------------------------------
|
||||||
|
#[derive(Debug, Error)]
|
||||||
|
pub enum BrightStaffError {
|
||||||
|
#[error("The requested model '{0}' does not exist")]
|
||||||
|
ModelNotFound(String),
|
||||||
|
|
||||||
|
#[error("No model specified in request and no default provider configured")]
|
||||||
|
NoModelSpecified,
|
||||||
|
|
||||||
|
#[error("Conversation state not found for previous_response_id: {0}")]
|
||||||
|
ConversationStateNotFound(String),
|
||||||
|
|
||||||
|
#[error("Internal server error")]
|
||||||
|
InternalServerError(String),
|
||||||
|
|
||||||
|
#[error("Invalid request")]
|
||||||
|
InvalidRequest(String),
|
||||||
|
|
||||||
|
#[error("{message}")]
|
||||||
|
ForwardedError {
|
||||||
|
status_code: StatusCode,
|
||||||
|
message: String,
|
||||||
|
},
|
||||||
|
|
||||||
|
#[error("Stream error: {0}")]
|
||||||
|
StreamError(String),
|
||||||
|
|
||||||
|
#[error("Failed to create response: {0}")]
|
||||||
|
ResponseCreationFailed(#[from] hyper::http::Error),
|
||||||
|
}
|
||||||
|
|
||||||
|
impl BrightStaffError {
|
||||||
|
pub fn into_response(self) -> Response<BoxBody<Bytes, HyperError>> {
|
||||||
|
let (status, code, details) = match &self {
|
||||||
|
BrightStaffError::ModelNotFound(model_name) => (
|
||||||
|
StatusCode::NOT_FOUND,
|
||||||
|
"ModelNotFound",
|
||||||
|
json!({ "rejected_model_id": model_name }),
|
||||||
|
),
|
||||||
|
|
||||||
|
BrightStaffError::NoModelSpecified => {
|
||||||
|
(StatusCode::BAD_REQUEST, "NoModelSpecified", json!({}))
|
||||||
|
}
|
||||||
|
|
||||||
|
BrightStaffError::ConversationStateNotFound(prev_resp_id) => (
|
||||||
|
StatusCode::CONFLICT,
|
||||||
|
"ConversationStateNotFound",
|
||||||
|
json!({ "previous_response_id": prev_resp_id }),
|
||||||
|
),
|
||||||
|
|
||||||
|
BrightStaffError::InternalServerError(reason) => (
|
||||||
|
StatusCode::INTERNAL_SERVER_ERROR,
|
||||||
|
"InternalServerError",
|
||||||
|
// Passing the reason into details for easier debugging
|
||||||
|
json!({ "reason": reason }),
|
||||||
|
),
|
||||||
|
|
||||||
|
BrightStaffError::InvalidRequest(reason) => (
|
||||||
|
StatusCode::BAD_REQUEST,
|
||||||
|
"InvalidRequest",
|
||||||
|
json!({ "reason": reason }),
|
||||||
|
),
|
||||||
|
|
||||||
|
BrightStaffError::ForwardedError {
|
||||||
|
status_code,
|
||||||
|
message,
|
||||||
|
} => (*status_code, "ForwardedError", json!({ "reason": message })),
|
||||||
|
|
||||||
|
BrightStaffError::StreamError(reason) => (
|
||||||
|
StatusCode::BAD_REQUEST,
|
||||||
|
"StreamError",
|
||||||
|
json!({ "reason": reason }),
|
||||||
|
),
|
||||||
|
|
||||||
|
BrightStaffError::ResponseCreationFailed(reason) => (
|
||||||
|
StatusCode::BAD_REQUEST,
|
||||||
|
"ResponseCreationFailed",
|
||||||
|
json!({ "reason": reason.to_string() }),
|
||||||
|
),
|
||||||
|
};
|
||||||
|
|
||||||
|
let body_json = json!({
|
||||||
|
"error": {
|
||||||
|
"code": code,
|
||||||
|
"message": self.to_string(),
|
||||||
|
"details": details
|
||||||
|
}
|
||||||
|
});
|
||||||
|
|
||||||
|
// 1. Create the concrete body
|
||||||
|
let full_body = Full::new(Bytes::from(body_json.to_string()));
|
||||||
|
|
||||||
|
// 2. Convert it to BoxBody
|
||||||
|
// We map_err because Full never fails, but BoxBody expects a HyperError
|
||||||
|
let boxed_body = full_body
|
||||||
|
.map_err(|never| match never {}) // This handles the "Infallible" error type
|
||||||
|
.boxed();
|
||||||
|
|
||||||
|
Response::builder()
|
||||||
|
.status(status)
|
||||||
|
.header("content-type", "application/json")
|
||||||
|
.body(boxed_body)
|
||||||
|
.unwrap_or_else(|_| {
|
||||||
|
Response::new(
|
||||||
|
Full::new(Bytes::from("Internal Error"))
|
||||||
|
.map_err(|never| match never {})
|
||||||
|
.boxed(),
|
||||||
|
)
|
||||||
|
})
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
#[cfg(test)]
|
||||||
|
mod tests {
|
||||||
|
use super::*;
|
||||||
|
use http_body_util::BodyExt; // For .collect().await
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_model_not_found_format() {
|
||||||
|
let err = BrightStaffError::ModelNotFound("gpt-5-secret".to_string());
|
||||||
|
let response = err.into_response();
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::NOT_FOUND);
|
||||||
|
|
||||||
|
// Helper to extract body as JSON
|
||||||
|
let body_bytes = response.into_body().collect().await.unwrap().to_bytes();
|
||||||
|
let body: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(body["error"]["code"], "ModelNotFound");
|
||||||
|
assert_eq!(
|
||||||
|
body["error"]["details"]["rejected_model_id"],
|
||||||
|
"gpt-5-secret"
|
||||||
|
);
|
||||||
|
assert!(body["error"]["message"]
|
||||||
|
.as_str()
|
||||||
|
.unwrap()
|
||||||
|
.contains("gpt-5-secret"));
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_forwarded_error_preserves_status() {
|
||||||
|
let err = BrightStaffError::ForwardedError {
|
||||||
|
status_code: StatusCode::TOO_MANY_REQUESTS,
|
||||||
|
message: "Rate limit exceeded on agent side".to_string(),
|
||||||
|
};
|
||||||
|
|
||||||
|
let response = err.into_response();
|
||||||
|
assert_eq!(response.status(), StatusCode::TOO_MANY_REQUESTS);
|
||||||
|
|
||||||
|
let body_bytes = response.into_body().collect().await.unwrap().to_bytes();
|
||||||
|
let body: serde_json::Value = serde_json::from_slice(&body_bytes).unwrap();
|
||||||
|
|
||||||
|
assert_eq!(body["error"]["code"], "ForwardedError");
|
||||||
|
}
|
||||||
|
|
||||||
|
#[tokio::test]
|
||||||
|
async fn test_hyper_error_wrapping() {
|
||||||
|
// Manually trigger a hyper error by creating an invalid URI/Header
|
||||||
|
let hyper_err = hyper::http::Response::builder()
|
||||||
|
.status(1000) // Invalid status
|
||||||
|
.body(())
|
||||||
|
.unwrap_err();
|
||||||
|
|
||||||
|
let err = BrightStaffError::ResponseCreationFailed(hyper_err);
|
||||||
|
let response = err.into_response();
|
||||||
|
|
||||||
|
assert_eq!(response.status(), StatusCode::BAD_REQUEST);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
|
||||||
Loading…
Add table
Add a link
Reference in a new issue