feat: introduce Rig agent runtime migration
This commit is contained in:
@@ -0,0 +1,201 @@
|
||||
use futures::StreamExt;
|
||||
use galaxy_agent_core::{
|
||||
AgentEvent, AgentRuntime, ConversationMessage, MessageContent, MessageRole,
|
||||
};
|
||||
use rig_core::client::CompletionClient;
|
||||
use rig_core::providers::openai;
|
||||
use rig_core::test_utils::MockStreamingClient;
|
||||
|
||||
use super::*;
|
||||
|
||||
fn text_request() -> TurnRequest {
|
||||
TurnRequest::new(
|
||||
"test-model",
|
||||
vec![ConversationMessage {
|
||||
role: MessageRole::User,
|
||||
content: MessageContent::Text("Hello".to_string()),
|
||||
}],
|
||||
)
|
||||
}
|
||||
|
||||
fn sse(lines: &[&str]) -> bytes::Bytes {
|
||||
lines
|
||||
.iter()
|
||||
.map(|line| format!("data: {line}\n\n"))
|
||||
.collect::<String>()
|
||||
.into()
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn rig_stream_maps_reasoning_text_usage_and_stop() {
|
||||
let http_client = MockStreamingClient {
|
||||
sse_bytes: sse(&[
|
||||
r#"{"id":"cmpl-1","model":"test-model","choices":[{"delta":{"reasoning_content":"thinking ","tool_calls":[]},"finish_reason":null}],"usage":null}"#,
|
||||
r#"{"id":"cmpl-1","model":"test-model","choices":[{"delta":{"content":"Hello ","tool_calls":[]},"finish_reason":null}],"usage":null}"#,
|
||||
r#"{"id":"cmpl-1","model":"test-model","choices":[{"delta":{"content":"world","tool_calls":[]},"finish_reason":"stop"}],"usage":null}"#,
|
||||
r#"{"choices":[],"usage":{"prompt_tokens":4,"completion_tokens":6,"total_tokens":10,"prompt_tokens_details":{"cached_tokens":2}}}"#,
|
||||
"[DONE]",
|
||||
]),
|
||||
};
|
||||
let client = openai::CompletionsClient::builder()
|
||||
.api_key("test-key")
|
||||
.base_url("http://localhost/v1")
|
||||
.http_client(http_client)
|
||||
.build()
|
||||
.unwrap();
|
||||
let model = client.completion_model("test-model");
|
||||
let (_, control) = galaxy_agent_core::turn_control();
|
||||
|
||||
let events = start_model_turn(model, text_request(), control, None, true)
|
||||
.await
|
||||
.unwrap()
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(events[0], AgentEvent::TurnStarted { .. }));
|
||||
assert_eq!(
|
||||
events[1..],
|
||||
[
|
||||
AgentEvent::ReasoningDelta {
|
||||
text: "thinking ".to_string(),
|
||||
},
|
||||
AgentEvent::TextDelta {
|
||||
text: "Hello ".to_string(),
|
||||
},
|
||||
AgentEvent::TextDelta {
|
||||
text: "world".to_string(),
|
||||
},
|
||||
AgentEvent::UsageUpdated {
|
||||
usage: Usage {
|
||||
input_tokens: 4,
|
||||
output_tokens: 6,
|
||||
cached_input_tokens: 2,
|
||||
cache_creation_input_tokens: 0,
|
||||
},
|
||||
},
|
||||
AgentEvent::TurnStopped {
|
||||
reason: StopReason::Completed,
|
||||
},
|
||||
]
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn cancellation_before_stream_start_is_a_normal_stop() {
|
||||
let runtime = OpenAICompatibleRuntime::new(OpenAICompatibleRuntimeConfig {
|
||||
base_url: "http://localhost/v1".to_string(),
|
||||
api_key: None,
|
||||
model: "test-model".to_string(),
|
||||
max_output_tokens: None,
|
||||
supports_system_messages: true,
|
||||
});
|
||||
let (sender, control) = galaxy_agent_core::turn_control();
|
||||
sender.send(TurnCommand::Cancel).await.unwrap();
|
||||
|
||||
let events = runtime
|
||||
.start_turn(text_request(), control)
|
||||
.await
|
||||
.unwrap()
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap();
|
||||
|
||||
assert!(matches!(events[0], AgentEvent::TurnStarted { .. }));
|
||||
assert_eq!(
|
||||
events[1],
|
||||
AgentEvent::TurnStopped {
|
||||
reason: StopReason::Cancelled,
|
||||
}
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn usage_at_the_requested_limit_maps_to_max_tokens() {
|
||||
let http_client = MockStreamingClient {
|
||||
sse_bytes: sse(&[
|
||||
r#"{"choices":[{"delta":{"content":"cut off","tool_calls":[]},"finish_reason":"length"}],"usage":null}"#,
|
||||
r#"{"choices":[],"usage":{"prompt_tokens":2,"completion_tokens":6,"total_tokens":8}}"#,
|
||||
"[DONE]",
|
||||
]),
|
||||
};
|
||||
let client = openai::CompletionsClient::builder()
|
||||
.api_key("test-key")
|
||||
.base_url("http://localhost/v1")
|
||||
.http_client(http_client)
|
||||
.build()
|
||||
.unwrap();
|
||||
let model = client.completion_model("test-model");
|
||||
let (sender, control) = galaxy_agent_core::turn_control();
|
||||
let mut request = text_request();
|
||||
request.max_output_tokens = Some(6);
|
||||
|
||||
let events = start_model_turn(model, request, control, None, true)
|
||||
.await
|
||||
.unwrap()
|
||||
.collect::<Vec<_>>()
|
||||
.await
|
||||
.into_iter()
|
||||
.collect::<Result<Vec<_>, _>>()
|
||||
.unwrap();
|
||||
drop(sender);
|
||||
|
||||
assert_eq!(
|
||||
events.last(),
|
||||
Some(&AgentEvent::TurnStopped {
|
||||
reason: StopReason::MaxTokens,
|
||||
})
|
||||
);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_conversion_preserves_history_tools_and_limits() {
|
||||
let mut request = text_request();
|
||||
request.system_prompt = Some("Be useful".to_string());
|
||||
request.max_output_tokens = Some(123);
|
||||
request.tools.push(galaxy_agent_core::ToolDefinition {
|
||||
name: "shell".to_string(),
|
||||
description: "Run a command".to_string(),
|
||||
input_schema: serde_json::json!({"type": "object"}),
|
||||
});
|
||||
|
||||
let converted = build_completion_request(request, Some(999), true).unwrap();
|
||||
|
||||
assert_eq!(converted.max_tokens, Some(123));
|
||||
assert_eq!(converted.tools.len(), 1);
|
||||
assert_eq!(converted.tools[0].name, "shell");
|
||||
assert_eq!(converted.chat_history.len(), 2);
|
||||
assert!(matches!(
|
||||
converted.chat_history.iter().next(),
|
||||
Some(Message::System { content }) if content == "Be useful"
|
||||
));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn request_conversion_places_system_prompt_in_user_message_when_system_role_is_unsupported() {
|
||||
let mut request = text_request();
|
||||
request.system_prompt = Some("Be useful".to_string());
|
||||
|
||||
let converted = build_completion_request(request, None, false).unwrap();
|
||||
let messages = converted.chat_history.iter().collect::<Vec<_>>();
|
||||
|
||||
assert_eq!(messages.len(), 2);
|
||||
let Message::User { content } = messages[0] else {
|
||||
panic!("expected the system prompt to use the user role");
|
||||
};
|
||||
let Some(UserContent::Text(text)) = content.iter().next() else {
|
||||
panic!("expected text instructions");
|
||||
};
|
||||
assert_eq!(text.text, "Be useful");
|
||||
assert_eq!(
|
||||
messages
|
||||
.iter()
|
||||
.filter(|message| matches!(message, Message::System { .. }))
|
||||
.count(),
|
||||
0
|
||||
);
|
||||
}
|
||||
Reference in New Issue
Block a user