Merge branch 'main' into adil/agent_format

crates/brightstaff/src/handlers/chat_completions.rs

@@ -211,9 +211,10 @@ pub async fn chat(
         }
     };
 
-    // copy over the headers from the original response
+    // copy over the headers and status code from the original response
     let response_headers = llm_response.headers().clone();
-    let mut response = Response::builder();
+    let upstream_status = llm_response.status();
+    let mut response = Response::builder().status(upstream_status);
     let headers = response.headers_mut().unwrap();
     for (header_name, header_value) in response_headers.iter() {
         headers.insert(header_name, header_value.clone());

crates/hermesllm/src/clients/transformer.rs

@@ -102,7 +102,7 @@ impl TryFrom<AnthropicMessagesRequest> for ChatCompletionsRequest {
             messages: openai_messages,
             temperature: req.temperature,
             top_p: req.top_p,
-            max_tokens: Some(req.max_tokens),
+            max_completion_tokens: Some(req.max_tokens),
             stream: req.stream,
             stop: req.stop_sequences,
             tools: openai_tools,
@@ -142,7 +142,9 @@ impl TryFrom<ChatCompletionsRequest> for AnthropicMessagesRequest {
             model: req.model,
             system: system_prompt,
             messages,
-            max_tokens: req.max_tokens.unwrap_or(DEFAULT_MAX_TOKENS),
+            max_tokens: req.max_completion_tokens
+                .or(req.max_tokens)
+                .unwrap_or(DEFAULT_MAX_TOKENS),
             container: None,
             mcp_servers: None,
             service_tier: None,
@@ -1079,7 +1081,7 @@ mod tests {
 
         assert_eq!(openai_req.model, "claude-3-sonnet-20240229");
         assert_eq!(openai_req.messages.len(), 2); // system + user message
-        assert_eq!(openai_req.max_tokens, Some(1024));
+        assert_eq!(openai_req.max_completion_tokens, Some(1024));
         assert_eq!(openai_req.temperature, Some(0.7));
         assert_eq!(openai_req.top_p, Some(0.9));
         assert_eq!(openai_req.stream, Some(false));

crates/hermesllm/src/providers/request.rs

@@ -360,6 +360,9 @@ mod tests {
             assert_eq!(openai_req.model, openai_req2.model);
             assert_eq!(openai_req.messages[0].role, openai_req2.messages[0].role);
             assert_eq!(openai_req.messages[0].content.extract_text(), openai_req2.messages[0].content.extract_text());
-            assert_eq!(openai_req.max_tokens, openai_req2.max_tokens);
+            // After roundtrip, deprecated max_tokens should be converted to max_completion_tokens
+            let original_max_tokens = openai_req.max_completion_tokens.or(openai_req.max_tokens);
+            let roundtrip_max_tokens = openai_req2.max_completion_tokens.or(openai_req2.max_tokens);
+            assert_eq!(original_max_tokens, roundtrip_max_tokens);
         }
 }

crates/llm_gateway/src/stream_context.rs

@@ -45,6 +45,8 @@ pub struct StreamContext {
     traces_queue: Arc<Mutex<VecDeque<TraceData>>>,
     overrides: Rc<Option<Overrides>>,
     user_message: Option<String>,
+    /// Store upstream response status code to handle error responses gracefully
+    upstream_status_code: Option<StatusCode>,
 }
 
 impl StreamContext {
@@ -72,6 +74,7 @@ impl StreamContext {
             traces_queue,
             request_body_sent_time: None,
             user_message: None,
+            upstream_status_code: None,
         }
     }
 
@@ -871,6 +874,19 @@ impl HttpContext for StreamContext {
     }
 
     fn on_http_response_headers(&mut self, _num_headers: usize, _end_of_stream: bool) -> Action {
+        // Capture the upstream response status code to handle errors appropriately
+        if let Some(status_str) = self.get_http_response_header(":status") {
+            if let Ok(status_code) = status_str.parse::<u16>() {
+                self.upstream_status_code = StatusCode::from_u16(status_code).ok();
+
+                info!(
+                    "[ARCHGW_REQ_ID:{}] UPSTREAM_RESPONSE_STATUS: {}",
+                    self.request_identifier(),
+                    status_code
+                );
+            }
+        }
+
         self.remove_http_response_header("content-length");
         self.remove_http_response_header("content-encoding");
 
@@ -888,6 +904,32 @@ impl HttpContext for StreamContext {
             return Action::Continue;
         }
 
+        // Check if this is an error response from upstream
+        if let Some(status_code) = &self.upstream_status_code {
+            if status_code.is_client_error() || status_code.is_server_error() {
+                info!(
+                    "[ARCHGW_REQ_ID:{}] UPSTREAM_ERROR_RESPONSE: status={} body_size={}",
+                    self.request_identifier(),
+                    status_code.as_u16(),
+                    body_size
+                );
+
+                // For error responses, forward the upstream error directly without parsing
+                if body_size > 0 {
+                    if let Ok(body) = self.read_raw_response_body(body_size) {
+                        debug!(
+                            "[ARCHGW_REQ_ID:{}] UPSTREAM_ERROR_BODY: {}",
+                            self.request_identifier(),
+                            String::from_utf8_lossy(&body)
+                        );
+                        // Forward the error response as-is
+                        self.set_http_response_body(0, body_size, &body);
+                    }
+                }
+                return Action::Continue;
+            }
+        }
+
         match self.client_api {
             Some(SupportedAPIs::OpenAIChatCompletions(_)) => {}
             Some(SupportedAPIs::AnthropicMessagesAPI(_)) => {}