Update: 将子项目从 submodule 转为完整内容
- 移除 GovAI, nomifun-tauri, 算力盒子 的 submodule 引用 - 添加所有子项目的完整源代码 - 保留原始 .git 为 .git.bak 备份
This commit is contained in:
freedak
@@ -0,0 +1,542 @@
|
||||
// Integration tests for AnthropicProvider using wiremock to mock the Anthropic API.
|
||||
|
||||
use wiremock::matchers::{header, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
use nomi_config::compat::ProviderCompat;
|
||||
use nomi_providers::anthropic::AnthropicProvider;
|
||||
use nomi_providers::{LlmProvider, ProviderError};
|
||||
use nomi_types::llm::{LlmEvent, LlmRequest, ThinkingConfig};
|
||||
use nomi_types::message::{ContentBlock, Message, Role, StopReason};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
fn minimal_request() -> LlmRequest {
|
||||
LlmRequest {
|
||||
model: "claude-3-5-sonnet-20241022".to_string(),
|
||||
system: "You are helpful.".to_string(),
|
||||
messages: vec![Message::new(
|
||||
Role::User,
|
||||
vec![ContentBlock::Text {
|
||||
text: "Hello".to_string(),
|
||||
}],
|
||||
)],
|
||||
tools: vec![],
|
||||
max_tokens: 1024,
|
||||
thinking: None,
|
||||
reasoning_effort: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Build a complete SSE body for a simple text response.
|
||||
fn text_sse_body(text: &str) -> String {
|
||||
format!(
|
||||
"event: message_start\n\
|
||||
data: {{\"type\":\"message_start\",\"message\":{{\"id\":\"msg_test\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-sonnet-20241022\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{{\"input_tokens\":100,\"output_tokens\":1}}}}}}\n\n\
|
||||
event: content_block_start\n\
|
||||
data: {{\"type\":\"content_block_start\",\"index\":0,\"content_block\":{{\"type\":\"text\",\"text\":\"\"}}}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {{\"type\":\"content_block_delta\",\"index\":0,\"delta\":{{\"type\":\"text_delta\",\"text\":\"{text}\"}}}}\n\n\
|
||||
event: content_block_stop\n\
|
||||
data: {{\"type\":\"content_block_stop\",\"index\":0}}\n\n\
|
||||
event: message_delta\n\
|
||||
data: {{\"type\":\"message_delta\",\"delta\":{{\"stop_reason\":\"end_turn\",\"stop_sequence\":null}},\"usage\":{{\"output_tokens\":50}}}}\n\n\
|
||||
event: message_stop\n\
|
||||
data: {{\"type\":\"message_stop\"}}\n\n"
|
||||
)
|
||||
}
|
||||
|
||||
/// Collect all events from a receiver into a Vec, draining until closed.
|
||||
async fn collect_events(mut rx: tokio::sync::mpsc::Receiver<LlmEvent>) -> Vec<LlmEvent> {
|
||||
let mut events = Vec::new();
|
||||
while let Some(ev) = rx.recv().await {
|
||||
events.push(ev);
|
||||
}
|
||||
events
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_stream_text_response
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A normal text SSE stream produces TextDelta events followed by a Done event.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_stream_text_response() {
|
||||
// Arrange: start a mock server
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200)
|
||||
.set_body_raw(text_sse_body("Hello, world!"), "text/event-stream"),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
let request = minimal_request();
|
||||
|
||||
// Act
|
||||
let rx = provider
|
||||
.stream(&request)
|
||||
.await
|
||||
.expect("stream should succeed");
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Assert: at least one TextDelta and exactly one Done
|
||||
let text_deltas: Vec<&LlmEvent> = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::TextDelta(_)))
|
||||
.collect();
|
||||
assert!(!text_deltas.is_empty(), "expected at least one TextDelta");
|
||||
|
||||
match &text_deltas[0] {
|
||||
LlmEvent::TextDelta(text) => assert_eq!(text, "Hello, world!"),
|
||||
_ => panic!("expected TextDelta"),
|
||||
}
|
||||
|
||||
let done_events: Vec<&LlmEvent> = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::Done { .. }))
|
||||
.collect();
|
||||
assert_eq!(done_events.len(), 1, "expected exactly one Done event");
|
||||
|
||||
match done_events[0] {
|
||||
LlmEvent::Done { stop_reason, usage } => {
|
||||
assert_eq!(*stop_reason, StopReason::EndTurn);
|
||||
assert_eq!(usage.input_tokens, 100);
|
||||
assert_eq!(usage.output_tokens, 50);
|
||||
}
|
||||
_ => panic!("expected Done"),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_stream_text_response_crlf_framed
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Some Anthropic-compatible gateways (e.g. new-api / one-api proxies) frame SSE
|
||||
/// events with CRLF ("\r\n\r\n") instead of the Anthropic API's "\n\n". The
|
||||
/// stream must still parse to TextDelta + Done rather than silently yielding an
|
||||
/// empty response. Regression test for the stepfun-proxy / new-api breakage.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_stream_text_response_crlf_framed() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Convert the canonical LF-framed body into CRLF framing.
|
||||
let crlf_body = text_sse_body("Hello, world!").replace('\n', "\r\n");
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(crlf_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
|
||||
let rx = provider
|
||||
.stream(&minimal_request())
|
||||
.await
|
||||
.expect("stream should succeed");
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
let text: String = events
|
||||
.iter()
|
||||
.filter_map(|e| match e {
|
||||
LlmEvent::TextDelta(t) => Some(t.clone()),
|
||||
_ => None,
|
||||
})
|
||||
.collect();
|
||||
assert_eq!(text, "Hello, world!", "CRLF-framed SSE must yield text");
|
||||
|
||||
let done = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::Done { .. }))
|
||||
.count();
|
||||
assert_eq!(done, 1, "expected exactly one Done event from CRLF stream");
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_stream_tool_use
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// An SSE stream containing a tool_use block produces a ToolUse event with
|
||||
/// accumulated JSON input.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_stream_tool_use() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let sse_body = "\
|
||||
event: message_start\n\
|
||||
data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_tool\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-sonnet-20241022\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":80,\"output_tokens\":1}}}\n\n\
|
||||
event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"tool_use\",\"id\":\"toolu_abc\",\"name\":\"Read\",\"input\":{}}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"{\\\"file\"}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"input_json_delta\",\"partial_json\":\"_path\\\":\\\"/tmp/test\\\"}\"}}\n\n\
|
||||
event: content_block_stop\n\
|
||||
data: {\"type\":\"content_block_stop\",\"index\":0}\n\n\
|
||||
event: message_delta\n\
|
||||
data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"tool_use\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":30}}\n\n\
|
||||
event: message_stop\n\
|
||||
data: {\"type\":\"message_stop\"}\n\n";
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
let request = minimal_request();
|
||||
|
||||
// Act
|
||||
let rx = provider
|
||||
.stream(&request)
|
||||
.await
|
||||
.expect("stream should succeed");
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Assert: one ToolUse event with correct fields
|
||||
let tool_events: Vec<&LlmEvent> = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::ToolUse { .. }))
|
||||
.collect();
|
||||
assert_eq!(tool_events.len(), 1, "expected exactly one ToolUse event");
|
||||
|
||||
match tool_events[0] {
|
||||
LlmEvent::ToolUse {
|
||||
id, name, input, ..
|
||||
} => {
|
||||
assert_eq!(id, "toolu_abc");
|
||||
assert_eq!(name, "Read");
|
||||
assert_eq!(input["file_path"], "/tmp/test");
|
||||
}
|
||||
_ => panic!("expected ToolUse"),
|
||||
}
|
||||
|
||||
// Done event should reflect tool_use stop reason
|
||||
let done_events: Vec<&LlmEvent> = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::Done { .. }))
|
||||
.collect();
|
||||
assert_eq!(done_events.len(), 1);
|
||||
match done_events[0] {
|
||||
LlmEvent::Done { stop_reason, .. } => {
|
||||
assert_eq!(*stop_reason, StopReason::ToolUse);
|
||||
}
|
||||
_ => panic!("expected Done"),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_stream_with_thinking
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// An SSE stream containing a thinking block produces ThinkingDelta events.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_stream_with_thinking() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let sse_body = "\
|
||||
event: message_start\n\
|
||||
data: {\"type\":\"message_start\",\"message\":{\"id\":\"msg_think\",\"type\":\"message\",\"role\":\"assistant\",\"content\":[],\"model\":\"claude-3-5-sonnet-20241022\",\"stop_reason\":null,\"stop_sequence\":null,\"usage\":{\"input_tokens\":90,\"output_tokens\":1}}}\n\n\
|
||||
event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":0,\"content_block\":{\"type\":\"thinking\",\"thinking\":\"\"}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":0,\"delta\":{\"type\":\"thinking_delta\",\"thinking\":\"Let me think...\"}}\n\n\
|
||||
event: content_block_stop\n\
|
||||
data: {\"type\":\"content_block_stop\",\"index\":0}\n\n\
|
||||
event: content_block_start\n\
|
||||
data: {\"type\":\"content_block_start\",\"index\":1,\"content_block\":{\"type\":\"text\",\"text\":\"\"}}\n\n\
|
||||
event: content_block_delta\n\
|
||||
data: {\"type\":\"content_block_delta\",\"index\":1,\"delta\":{\"type\":\"text_delta\",\"text\":\"Answer.\"}}\n\n\
|
||||
event: content_block_stop\n\
|
||||
data: {\"type\":\"content_block_stop\",\"index\":1}\n\n\
|
||||
event: message_delta\n\
|
||||
data: {\"type\":\"message_delta\",\"delta\":{\"stop_reason\":\"end_turn\",\"stop_sequence\":null},\"usage\":{\"output_tokens\":20}}\n\n\
|
||||
event: message_stop\n\
|
||||
data: {\"type\":\"message_stop\"}\n\n";
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
// Enable thinking in the request
|
||||
let mut request = minimal_request();
|
||||
request.thinking = Some(ThinkingConfig::Enabled {
|
||||
budget_tokens: 5000,
|
||||
});
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
|
||||
// Act
|
||||
let rx = provider
|
||||
.stream(&request)
|
||||
.await
|
||||
.expect("stream should succeed");
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Assert: ThinkingDelta event present with expected content
|
||||
let thinking_events: Vec<&LlmEvent> = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::ThinkingDelta(_)))
|
||||
.collect();
|
||||
assert!(
|
||||
!thinking_events.is_empty(),
|
||||
"expected at least one ThinkingDelta"
|
||||
);
|
||||
|
||||
match thinking_events[0] {
|
||||
LlmEvent::ThinkingDelta(text) => assert_eq!(text, "Let me think..."),
|
||||
_ => panic!("expected ThinkingDelta"),
|
||||
}
|
||||
|
||||
// TextDelta should also be present
|
||||
let text_events: Vec<&LlmEvent> = events
|
||||
.iter()
|
||||
.filter(|e| matches!(e, LlmEvent::TextDelta(_)))
|
||||
.collect();
|
||||
assert!(
|
||||
!text_events.is_empty(),
|
||||
"expected at least one TextDelta after thinking"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_auth_error
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A 401 response from the API should produce a ProviderError::Api with status 401.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_auth_error() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let error_body =
|
||||
r#"{"type":"error","error":{"type":"authentication_error","message":"invalid x-api-key"}}"#;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_string(error_body))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"bad-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
let request = minimal_request();
|
||||
|
||||
// Act
|
||||
let result = provider.stream(&request).await;
|
||||
|
||||
// Assert: returns an Api error with status 401
|
||||
match result {
|
||||
Err(ProviderError::Api { status, message }) => {
|
||||
assert_eq!(status, 401);
|
||||
assert!(
|
||||
message.contains("authentication_error") || message.contains("invalid x-api-key"),
|
||||
"unexpected error message: {message}"
|
||||
);
|
||||
}
|
||||
Err(other) => panic!("expected Api error, got: {other:?}"),
|
||||
Ok(_) => panic!("expected an error but stream succeeded"),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_rate_limit_retryable
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// A 429 response from the API should produce a ProviderError::RateLimited.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_rate_limit_retryable() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(ResponseTemplate::new(429).set_body_string(
|
||||
r#"{"type":"error","error":{"type":"rate_limit_error","message":"rate limit exceeded"}}"#,
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
let request = minimal_request();
|
||||
|
||||
// Act
|
||||
let result = provider.stream(&request).await;
|
||||
|
||||
// Assert: RateLimited error, which is retryable
|
||||
match result {
|
||||
Err(ProviderError::RateLimited { retry_after_ms, .. }) => {
|
||||
assert!(retry_after_ms > 0, "retry_after_ms should be positive");
|
||||
}
|
||||
Err(other) => panic!("expected RateLimited error, got: {other:?}"),
|
||||
Ok(_) => panic!("expected an error but stream succeeded"),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_request_headers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// The provider must send the correct HTTP headers: x-api-key, anthropic-version,
|
||||
/// and content-type. This test uses wiremock header matchers to verify them.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_request_headers() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Register the mock with header matchers; only requests carrying the
|
||||
// correct headers will match and receive a 200 response.
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.and(header("x-api-key", "my-secret-key"))
|
||||
.and(header("anthropic-version", "2023-06-01"))
|
||||
.and(header("content-type", "application/json"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_raw(text_sse_body("ok"), "text/event-stream"),
|
||||
)
|
||||
.expect(1) // exactly one matching request must arrive
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"my-secret-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
let request = minimal_request();
|
||||
|
||||
// Act — should succeed because the headers are correct
|
||||
let result = provider.stream(&request).await;
|
||||
assert!(result.is_ok(), "stream failed: {:?}", result.err());
|
||||
|
||||
// Drain the channel so the spawned task finishes
|
||||
if let Ok(rx) = result {
|
||||
collect_events(rx).await;
|
||||
}
|
||||
|
||||
// wiremock verifies the `expect(1)` assertion when MockServer is dropped;
|
||||
// if the header matcher was not satisfied the test will panic here.
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_prompt_caching_header
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// When cache is enabled the provider must include the anthropic-beta header
|
||||
/// for prompt caching.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_prompt_caching_header() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.and(header("anthropic-beta", "prompt-caching-2024-07-31"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_raw(text_sse_body("cached"), "text/event-stream"),
|
||||
)
|
||||
.expect(1)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
// with_cache(true) — default, but explicit here for clarity
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(true);
|
||||
let request = minimal_request();
|
||||
|
||||
let result = provider.stream(&request).await;
|
||||
assert!(result.is_ok(), "stream failed: {:?}", result.err());
|
||||
|
||||
if let Ok(rx) = result {
|
||||
collect_events(rx).await;
|
||||
}
|
||||
|
||||
server.verify().await;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_anthropic_no_prompt_caching_header_when_disabled
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// When cache is disabled the anthropic-beta header must NOT be present.
|
||||
/// We verify this by mounting a mock that matches only without that header and
|
||||
/// checking it receives exactly one request.
|
||||
#[tokio::test]
|
||||
async fn test_anthropic_no_prompt_caching_header_when_disabled() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// This mock matches any POST to /v1/messages (no anthropic-beta requirement).
|
||||
// We then confirm via received_requests that the header is absent.
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/messages"))
|
||||
.respond_with(
|
||||
ResponseTemplate::new(200).set_body_raw(text_sse_body("no cache"), "text/event-stream"),
|
||||
)
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = AnthropicProvider::new(
|
||||
"test-api-key",
|
||||
&server.uri(),
|
||||
ProviderCompat::anthropic_defaults(),
|
||||
)
|
||||
.with_cache(false);
|
||||
let request = minimal_request();
|
||||
|
||||
let result = provider.stream(&request).await;
|
||||
assert!(result.is_ok(), "stream failed: {:?}", result.err());
|
||||
|
||||
if let Ok(rx) = result {
|
||||
collect_events(rx).await;
|
||||
}
|
||||
|
||||
// Inspect the captured request to assert that anthropic-beta is absent
|
||||
let received = server.received_requests().await.unwrap();
|
||||
assert_eq!(received.len(), 1, "expected exactly one request");
|
||||
let has_beta = received[0].headers.contains_key("anthropic-beta");
|
||||
assert!(
|
||||
!has_beta,
|
||||
"anthropic-beta header should not be present when cache is disabled"
|
||||
);
|
||||
}
|
||||
@@ -0,0 +1,720 @@
|
||||
use std::time::Duration;
|
||||
|
||||
use nomi_config::compat::ProviderCompat;
|
||||
use nomi_providers::LlmProvider;
|
||||
use nomi_providers::openai::OpenAIProvider;
|
||||
use nomi_types::llm::{LlmEvent, LlmRequest};
|
||||
use nomi_types::message::{ContentBlock, Message, Role, StopReason};
|
||||
use serde_json::json;
|
||||
use tokio::io::{AsyncReadExt, AsyncWriteExt};
|
||||
use tokio::net::TcpListener;
|
||||
use wiremock::matchers::{header, method, path};
|
||||
use wiremock::{Mock, MockServer, ResponseTemplate};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Build a minimal LlmRequest suitable for all tests.
|
||||
fn make_request() -> LlmRequest {
|
||||
LlmRequest {
|
||||
model: "gpt-4o".to_string(),
|
||||
system: "You are a test assistant.".to_string(),
|
||||
messages: vec![Message::new(
|
||||
Role::User,
|
||||
vec![ContentBlock::Text {
|
||||
text: "Hello".to_string(),
|
||||
}],
|
||||
)],
|
||||
tools: vec![],
|
||||
max_tokens: 512,
|
||||
thinking: None,
|
||||
reasoning_effort: None,
|
||||
}
|
||||
}
|
||||
|
||||
/// Collect all events from the receiver until the channel closes.
|
||||
async fn collect_events(mut rx: tokio::sync::mpsc::Receiver<LlmEvent>) -> Vec<LlmEvent> {
|
||||
let mut events = Vec::new();
|
||||
while let Some(event) = rx.recv().await {
|
||||
events.push(event);
|
||||
}
|
||||
events
|
||||
}
|
||||
|
||||
/// Build a raw SSE body string from a slice of JSON lines.
|
||||
/// Each line is wrapped in `data: ...\n\n` and a final `data: [DONE]\n\n` is appended.
|
||||
fn build_sse_body(data_lines: &[&str]) -> String {
|
||||
let mut body = String::new();
|
||||
for line in data_lines {
|
||||
body.push_str("data: ");
|
||||
body.push_str(line);
|
||||
body.push_str("\n\n");
|
||||
}
|
||||
body.push_str("data: [DONE]\n\n");
|
||||
body
|
||||
}
|
||||
|
||||
async fn start_server_after_initial_connect_refusal(sse_body: String) -> String {
|
||||
let probe = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = probe.local_addr().unwrap();
|
||||
drop(probe);
|
||||
|
||||
tokio::spawn(async move {
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
|
||||
let listener = TcpListener::bind(addr).await.unwrap();
|
||||
let (mut second, _) = listener.accept().await.unwrap();
|
||||
let mut buf = [0_u8; 4096];
|
||||
let _ = second.read(&mut buf).await.unwrap();
|
||||
|
||||
let response = format!(
|
||||
"HTTP/1.1 200 OK\r\ncontent-type: text/event-stream\r\ncontent-length: {}\r\nconnection: close\r\n\r\n{}",
|
||||
sse_body.len(),
|
||||
sse_body
|
||||
);
|
||||
second.write_all(response.as_bytes()).await.unwrap();
|
||||
});
|
||||
|
||||
format!("http://{addr}")
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_stream_text_response
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that a normal text response (multiple content deltas followed by a
|
||||
/// stop chunk with usage) is parsed into the correct sequence of TextDelta
|
||||
/// and Done events.
|
||||
#[tokio::test]
|
||||
async fn test_openai_stream_text_response() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Chunk 1: first text delta
|
||||
let chunk1 = json!({
|
||||
"id": "chatcmpl-001",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "role": "assistant", "content": "Hello" },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Chunk 2: second text delta
|
||||
let chunk2 = json!({
|
||||
"id": "chatcmpl-001",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "content": ", world!" },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Chunk 3: finish_reason = "stop" with usage
|
||||
let chunk3 = json!({
|
||||
"id": "chatcmpl-001",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 25,
|
||||
"completion_tokens": 10
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let sse_body = build_sse_body(&[&chunk1, &chunk2, &chunk3]);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.and(header("authorization", "Bearer test-key"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Expect: TextDelta("Hello"), TextDelta(", world!"), Done{EndTurn}
|
||||
assert_eq!(events.len(), 3, "expected 3 events, got: {:?}", events);
|
||||
|
||||
match &events[0] {
|
||||
LlmEvent::TextDelta(text) => assert_eq!(text, "Hello"),
|
||||
e => panic!("expected TextDelta, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[1] {
|
||||
LlmEvent::TextDelta(text) => assert_eq!(text, ", world!"),
|
||||
e => panic!("expected TextDelta, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[2] {
|
||||
LlmEvent::Done { stop_reason, usage } => {
|
||||
assert_eq!(*stop_reason, StopReason::EndTurn);
|
||||
assert_eq!(usage.input_tokens, 25);
|
||||
assert_eq!(usage.output_tokens, 10);
|
||||
}
|
||||
e => panic!("expected Done, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_initial_connect_error_is_retried
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that the provider retries when the initial HTTP request fails before
|
||||
/// receiving any response. This covers transient connect/TLS failures where no
|
||||
/// model output has been emitted yet.
|
||||
#[tokio::test]
|
||||
async fn test_openai_initial_connect_error_is_retried() {
|
||||
let chunk = json!({
|
||||
"id": "chatcmpl-retry",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "role": "assistant", "content": "Recovered" },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
let finish = json!({
|
||||
"id": "chatcmpl-retry",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop"
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
let sse_body = build_sse_body(&[&chunk, &finish]);
|
||||
let base_url = start_server_after_initial_connect_refusal(sse_body).await;
|
||||
|
||||
let provider = OpenAIProvider::new("test-key", &base_url, ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
assert_eq!(
|
||||
events.len(),
|
||||
2,
|
||||
"expected retry success events, got: {:?}",
|
||||
events
|
||||
);
|
||||
match &events[0] {
|
||||
LlmEvent::TextDelta(text) => assert_eq!(text, "Recovered"),
|
||||
e => panic!("expected TextDelta, got: {:?}", e),
|
||||
}
|
||||
match &events[1] {
|
||||
LlmEvent::Done { stop_reason, .. } => assert_eq!(*stop_reason, StopReason::EndTurn),
|
||||
e => panic!("expected Done, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_stream_tool_call_aggregation
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that a tool call streamed in multiple delta chunks (id in first chunk,
|
||||
/// name in first chunk, arguments split across chunks) is correctly aggregated
|
||||
/// into a single ToolUse event.
|
||||
#[tokio::test]
|
||||
async fn test_openai_stream_tool_call_aggregation() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Chunk 1: tool call header — id and function name arrive first
|
||||
let chunk1 = json!({
|
||||
"id": "chatcmpl-002",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 0,
|
||||
"id": "call_abc123",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"arguments": "{\"path\":"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Chunk 2: arguments continuation
|
||||
let chunk2 = json!({
|
||||
"id": "chatcmpl-002",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 0,
|
||||
"function": {
|
||||
"arguments": "\"/tmp/test.txt\"}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Chunk 3: finish_reason = "tool_calls" with usage
|
||||
let chunk3 = json!({
|
||||
"id": "chatcmpl-002",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 40,
|
||||
"completion_tokens": 15
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let sse_body = build_sse_body(&[&chunk1, &chunk2, &chunk3]);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Expect: ToolUse, Done{ToolUse}
|
||||
assert_eq!(events.len(), 2, "expected 2 events, got: {:?}", events);
|
||||
|
||||
match &events[0] {
|
||||
LlmEvent::ToolUse {
|
||||
id, name, input, ..
|
||||
} => {
|
||||
assert_eq!(id, "call_abc123");
|
||||
assert_eq!(name, "read_file");
|
||||
assert_eq!(input["path"], "/tmp/test.txt");
|
||||
}
|
||||
e => panic!("expected ToolUse, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[1] {
|
||||
LlmEvent::Done { stop_reason, usage } => {
|
||||
assert_eq!(*stop_reason, StopReason::ToolUse);
|
||||
assert_eq!(usage.input_tokens, 40);
|
||||
assert_eq!(usage.output_tokens, 15);
|
||||
}
|
||||
e => panic!("expected Done, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_multiple_tool_calls
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that when the API streams multiple parallel tool calls (different
|
||||
/// indices) they are all emitted as separate ToolUse events.
|
||||
#[tokio::test]
|
||||
async fn test_openai_multiple_tool_calls() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Chunk 1: first tool call (index 0)
|
||||
let chunk1 = json!({
|
||||
"id": "chatcmpl-003",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 0,
|
||||
"id": "call_tool0",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "list_files",
|
||||
"arguments": "{\"dir\": \"/tmp\"}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Chunk 2: second tool call (index 1)
|
||||
let chunk2 = json!({
|
||||
"id": "chatcmpl-003",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {
|
||||
"tool_calls": [{
|
||||
"index": 1,
|
||||
"id": "call_tool1",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": "read_file",
|
||||
"arguments": "{\"path\": \"/etc/hosts\"}"
|
||||
}
|
||||
}]
|
||||
},
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Chunk 3: finish_reason = "tool_calls"
|
||||
let chunk3 = json!({
|
||||
"id": "chatcmpl-003",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "tool_calls"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 60,
|
||||
"completion_tokens": 20
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let sse_body = build_sse_body(&[&chunk1, &chunk2, &chunk3]);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Expect: ToolUse (index 0), ToolUse (index 1), Done{ToolUse}
|
||||
assert_eq!(events.len(), 3, "expected 3 events, got: {:?}", events);
|
||||
|
||||
match &events[0] {
|
||||
LlmEvent::ToolUse {
|
||||
id, name, input, ..
|
||||
} => {
|
||||
assert_eq!(id, "call_tool0");
|
||||
assert_eq!(name, "list_files");
|
||||
assert_eq!(input["dir"], "/tmp");
|
||||
}
|
||||
e => panic!("expected first ToolUse, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[1] {
|
||||
LlmEvent::ToolUse {
|
||||
id, name, input, ..
|
||||
} => {
|
||||
assert_eq!(id, "call_tool1");
|
||||
assert_eq!(name, "read_file");
|
||||
assert_eq!(input["path"], "/etc/hosts");
|
||||
}
|
||||
e => panic!("expected second ToolUse, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[2] {
|
||||
LlmEvent::Done { stop_reason, .. } => {
|
||||
assert_eq!(*stop_reason, StopReason::ToolUse);
|
||||
}
|
||||
e => panic!("expected Done, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_stream_state_transitions
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that the stream correctly stops processing events once it encounters
|
||||
/// the `[DONE]` sentinel — any data after [DONE] is ignored and the receiver
|
||||
/// channel closes cleanly.
|
||||
#[tokio::test]
|
||||
async fn test_openai_stream_state_transitions() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// A single text delta followed by a stop chunk, then the [DONE] sentinel.
|
||||
let chunk1 = json!({
|
||||
"id": "chatcmpl-004",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "content": "Transition test." },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let chunk2 = json!({
|
||||
"id": "chatcmpl-004",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 10,
|
||||
"completion_tokens": 5
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
// Build SSE body manually: two data lines, then [DONE], then a stray line
|
||||
// that must NOT produce any events.
|
||||
let mut sse_body = String::new();
|
||||
sse_body.push_str("data: ");
|
||||
sse_body.push_str(&chunk1);
|
||||
sse_body.push_str("\n\n");
|
||||
sse_body.push_str("data: ");
|
||||
sse_body.push_str(&chunk2);
|
||||
sse_body.push_str("\n\n");
|
||||
sse_body.push_str("data: [DONE]\n\n");
|
||||
// Stray chunk after [DONE] — must be ignored
|
||||
sse_body.push_str("data: {\"choices\":[{\"index\":0,\"delta\":{\"content\":\"ignored\"},\"finish_reason\":null}]}\n\n");
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Expect exactly: TextDelta, Done — the trailing chunk after [DONE] is discarded.
|
||||
assert_eq!(events.len(), 2, "expected 2 events, got: {:?}", events);
|
||||
|
||||
match &events[0] {
|
||||
LlmEvent::TextDelta(text) => assert_eq!(text, "Transition test."),
|
||||
e => panic!("expected TextDelta, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[1] {
|
||||
LlmEvent::Done { stop_reason, usage } => {
|
||||
assert_eq!(*stop_reason, StopReason::EndTurn);
|
||||
assert_eq!(usage.input_tokens, 10);
|
||||
assert_eq!(usage.output_tokens, 5);
|
||||
assert_eq!(usage.cache_creation_tokens, 0);
|
||||
assert_eq!(usage.cache_read_tokens, 0);
|
||||
}
|
||||
e => panic!("expected Done, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_api_error_non_success_status
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that a non-2xx HTTP response is surfaced as a ProviderError::Api.
|
||||
#[tokio::test]
|
||||
async fn test_openai_api_error_non_success_status() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(401).set_body_string(
|
||||
r#"{"error":{"message":"Invalid API key","type":"invalid_request_error"}}"#,
|
||||
))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider = OpenAIProvider::new("bad-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let result = provider.stream(&make_request()).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
match result.unwrap_err() {
|
||||
nomi_providers::ProviderError::Api { status, .. } => {
|
||||
assert_eq!(status, 401);
|
||||
}
|
||||
e => panic!("expected Api error, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_rate_limited
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that a 429 response is surfaced as ProviderError::RateLimited.
|
||||
#[tokio::test]
|
||||
async fn test_openai_rate_limited() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let body = r#"{"error":{"message":"You exceeded your current quota","type":"insufficient_quota","code":"insufficient_quota"}}"#;
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(429).set_body_string(body))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let result = provider.stream(&make_request()).await;
|
||||
|
||||
assert!(result.is_err());
|
||||
let err = result.unwrap_err();
|
||||
let display = err.to_string();
|
||||
match err {
|
||||
nomi_providers::ProviderError::RateLimited { retry_after_ms, .. } => {
|
||||
assert_eq!(retry_after_ms, 5000);
|
||||
}
|
||||
e => panic!("expected RateLimited error, got: {:?}", e),
|
||||
}
|
||||
|
||||
assert!(
|
||||
display.contains("insufficient_quota"),
|
||||
"rate limit error should preserve provider body, got: {display}"
|
||||
);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_stream_max_tokens_stop_reason
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that finish_reason "length" maps to StopReason::MaxTokens.
|
||||
#[tokio::test]
|
||||
async fn test_openai_stream_max_tokens_stop_reason() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
let chunk1 = json!({
|
||||
"id": "chatcmpl-005",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "content": "Truncated" },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let chunk2 = json!({
|
||||
"id": "chatcmpl-005",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "length"
|
||||
}],
|
||||
"usage": {
|
||||
"prompt_tokens": 100,
|
||||
"completion_tokens": 512
|
||||
}
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let sse_body = build_sse_body(&[&chunk1, &chunk2]);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
assert_eq!(events.len(), 2);
|
||||
|
||||
match &events[1] {
|
||||
LlmEvent::Done { stop_reason, usage } => {
|
||||
assert_eq!(*stop_reason, StopReason::MaxTokens);
|
||||
assert_eq!(usage.input_tokens, 100);
|
||||
assert_eq!(usage.output_tokens, 512);
|
||||
}
|
||||
e => panic!("expected Done with MaxTokens, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// test_openai_stream_empty_content_delta_skipped
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
/// Verify that empty content strings in deltas do NOT produce TextDelta events
|
||||
/// (the provider filters them out).
|
||||
#[tokio::test]
|
||||
async fn test_openai_stream_empty_content_delta_skipped() {
|
||||
let server = MockServer::start().await;
|
||||
|
||||
// Chunk with empty content — should be silently skipped
|
||||
let chunk_empty = json!({
|
||||
"id": "chatcmpl-006",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "content": "" },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let chunk_text = json!({
|
||||
"id": "chatcmpl-006",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": { "content": "actual content" },
|
||||
"finish_reason": null
|
||||
}]
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let chunk_done = json!({
|
||||
"id": "chatcmpl-006",
|
||||
"object": "chat.completion.chunk",
|
||||
"choices": [{
|
||||
"index": 0,
|
||||
"delta": {},
|
||||
"finish_reason": "stop"
|
||||
}],
|
||||
"usage": { "prompt_tokens": 5, "completion_tokens": 3 }
|
||||
})
|
||||
.to_string();
|
||||
|
||||
let sse_body = build_sse_body(&[&chunk_empty, &chunk_text, &chunk_done]);
|
||||
|
||||
Mock::given(method("POST"))
|
||||
.and(path("/v1/chat/completions"))
|
||||
.respond_with(ResponseTemplate::new(200).set_body_raw(sse_body, "text/event-stream"))
|
||||
.mount(&server)
|
||||
.await;
|
||||
|
||||
let provider =
|
||||
OpenAIProvider::new("test-key", &server.uri(), ProviderCompat::openai_defaults());
|
||||
let rx = provider.stream(&make_request()).await.unwrap();
|
||||
let events = collect_events(rx).await;
|
||||
|
||||
// Expect only TextDelta("actual content") and Done — no empty TextDelta
|
||||
assert_eq!(events.len(), 2, "expected 2 events, got: {:?}", events);
|
||||
|
||||
match &events[0] {
|
||||
LlmEvent::TextDelta(text) => assert_eq!(text, "actual content"),
|
||||
e => panic!("expected TextDelta with actual content, got: {:?}", e),
|
||||
}
|
||||
|
||||
match &events[1] {
|
||||
LlmEvent::Done { stop_reason, .. } => assert_eq!(*stop_reason, StopReason::EndTurn),
|
||||
e => panic!("expected Done, got: {:?}", e),
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user