111 lines
6.1 KiB
Rust
111 lines
6.1 KiB
Rust
//! Configured pi-ai API adapters use a fake server and persist catalog-priced usage.
|
|
use sausage::{
|
|
exec::{Command, Exec, Options, Stream},
|
|
world::World,
|
|
};
|
|
use std::{sync::Arc, time::Duration};
|
|
struct NoExec;
|
|
#[async_trait::async_trait]
|
|
impl Exec for NoExec {
|
|
async fn exec(
|
|
&self,
|
|
_: Command,
|
|
_: Options,
|
|
_: &mut (dyn for<'a> FnMut(Stream, &'a [u8]) + Send),
|
|
) -> anyhow::Result<i32> {
|
|
anyhow::bail!("unexpected world command")
|
|
}
|
|
}
|
|
fn event(name: &str, value: serde_json::Value) -> String {
|
|
format!("event: {name}\ndata: {value}\n\n")
|
|
}
|
|
#[tokio::test(flavor = "current_thread")]
|
|
async fn provider_transport_and_prices() -> anyhow::Result<()> {
|
|
use serde_json::json;
|
|
let _ = rustls::crypto::aws_lc_rs::default_provider().install_default();
|
|
let server = tiny_http::Server::http("127.0.0.1:0").map_err(|e| anyhow::anyhow!(e))?;
|
|
let base = format!("http://{}", server.server_addr());
|
|
let settings = json!({"main_agent_name":"main","prompt":"Shared context","providers":{
|
|
"messages":{"type":"anthropic-messages","catalog":"anthropic","base_url":format!("{base}/messages-api"),"api_key":"test-provider-key","models":["claude-opus-5-5"]},
|
|
"responses":{"type":"openai-responses","catalog":"openai","base_url":format!("{base}/responses-api"),"api_key":"test-provider-key","models":["gpt-6.1-sol"]},
|
|
"chat":{"type":"openai-chat-completions","catalog":"openai","base_url":format!("{base}/chat-api"),"api_key":"test-provider-key","models":["gpt-6.1-sol"]}
|
|
}});
|
|
let thread = std::thread::spawn(move || -> anyhow::Result<()> {
|
|
for path in [
|
|
"/messages-api/v1/messages?beta=true",
|
|
"/responses-api/responses",
|
|
"/chat-api/chat/completions",
|
|
"/responses-api/responses",
|
|
"/edited-api/responses",
|
|
] {
|
|
let mut request = server
|
|
.recv_timeout(Duration::from_secs(30))?
|
|
.ok_or_else(|| anyhow::anyhow!("request timeout: {path}"))?;
|
|
assert_eq!(request.url(), path);
|
|
let key = if path.starts_with("/edited-api") {
|
|
"next-key"
|
|
} else {
|
|
"test-provider-key"
|
|
};
|
|
assert!(
|
|
request
|
|
.headers()
|
|
.iter()
|
|
.any(|h| (h.field.equiv("Authorization")
|
|
&& h.value.as_str() == format!("Bearer {key}"))
|
|
|| (h.field.equiv("x-api-key") && h.value.as_str() == key)),
|
|
"missing provider key"
|
|
);
|
|
let mut body = String::new();
|
|
request.as_reader().read_to_string(&mut body)?;
|
|
let body: serde_json::Value = serde_json::from_str(&body)?;
|
|
assert_eq!(body["stream"], true);
|
|
let stream = if path.contains("/messages") {
|
|
assert_eq!(body["model"], "claude-opus-5-5");
|
|
assert!(body["messages"].is_array());
|
|
[
|
|
event("message_start",json!({"type":"message_start","message":{"id":"m","type":"message","role":"assistant","model":"claude-opus-5-5","content":[],"usage":{"input_tokens":10,"output_tokens":0}}})),
|
|
event("content_block_start",json!({"type":"content_block_start","index":0,"content_block":{"type":"text","text":""}})),
|
|
event("content_block_delta",json!({"type":"content_block_delta","index":0,"delta":{"type":"text_delta","text":"hello"}})),
|
|
event("content_block_stop",json!({"type":"content_block_stop","index":0})),
|
|
event("message_delta",json!({"type":"message_delta","delta":{"stop_reason":"end_turn"},"usage":{"output_tokens":1}})),
|
|
event("message_stop",json!({"type":"message_stop"})),
|
|
].concat()
|
|
} else if path.contains("/chat/") {
|
|
assert_eq!(body["model"], "gpt-6.1-sol");
|
|
assert!(body["messages"].is_array());
|
|
[json!({"id":"c","object":"chat.completion.chunk","choices":[{"index":0,"delta":{"role":"assistant","content":"hello"},"finish_reason":null}]}),
|
|
json!({"id":"c","object":"chat.completion.chunk","choices":[{"index":0,"delta":{},"finish_reason":"stop"}],"usage":{"prompt_tokens":10,"completion_tokens":1,"total_tokens":11}})]
|
|
.iter().map(|v|format!("data: {v}\n\n")).collect::<String>()+"data: [DONE]\n\n"
|
|
} else {
|
|
assert_eq!(body["model"], "gpt-6.1-sol");
|
|
assert!(body["input"].is_array());
|
|
[
|
|
event("response.output_item.added",json!({"type":"response.output_item.added","output_index":0,"item":{"type":"message","id":"m","role":"assistant","content":[]}})),
|
|
event("response.content_part.added",json!({"type":"response.content_part.added","output_index":0,"content_index":0,"part":{"type":"output_text","text":""}})),
|
|
event("response.output_text.delta",json!({"type":"response.output_text.delta","output_index":0,"content_index":0,"delta":"hello"})),
|
|
event("response.completed",json!({"type":"response.completed","response":{"id":"r","status":"completed","output":[],"usage":{"input_tokens":10,"output_tokens":1,"total_tokens":11}}})),
|
|
].concat()
|
|
};
|
|
request.respond(tiny_http::Response::from_string(stream).with_header(
|
|
tiny_http::Header::from_bytes("Content-Type", "text/event-stream").unwrap(),
|
|
))?;
|
|
}
|
|
Ok(())
|
|
});
|
|
let dir = tempfile::tempdir()?;
|
|
let config = dir.path().join("world/etc/sausage");
|
|
std::fs::create_dir_all(&config)?;
|
|
std::fs::write(config.join("config.yaml"), serde_json::to_vec(&settings)?)?;
|
|
let source = format!(
|
|
"globalThis.providerSettings={settings};globalThis.providerBase={};\n{}",
|
|
serde_json::to_string(&base)?,
|
|
include_str!("../../js/provider-models-test-bundle.js")
|
|
);
|
|
let result =
|
|
sausage::harness::run_module(World::open(dir.path())?, Arc::new(NoExec), &source).await;
|
|
let server_result = thread.join().unwrap();
|
|
result?;
|
|
server_result
|
|
}
|