boos/BoosAPI
2
0
Fork 2
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add workflow studio and postgres support

Browse Source
This commit is contained in:
boos
2026-05-25 16:03:07 +08:00
parent 7cdb02ef1d
commit cb35d5fc6e
96 changed files with 12414 additions and 15548 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/routes/proxy/mod.rs
+162 -25
View File
@@ -103,10 +103,12 @@ async fn messages_handler(
let model = canonical.model.clone();
if canonical.stream {
let chunks = execute_messages_stream_with_conductor(state, key_id, model, canonical).await?;
let sse = Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
let chunks =
execute_messages_stream_with_conductor(state, key_id, model, canonical).await?;
let sse =
Sse::new(stream::iter(chunks).map(|chunk| {
Ok::<_, Infallible>(axum::response::sse::Event::default().data(chunk))
}));
Ok(sse.into_response())
} else {
let result = execute_messages_with_conductor(state, key_id, model, canonical).await?;
@@ -371,10 +373,53 @@ async fn images_generations(
.ok_or_else(|| AppError::BadRequest("model is required".into()))?
.to_string();
let route = resolve_route(&state, &model).await?;
let result = executors::execute_image_generation(state.clone(), route, body).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
let started_at = Instant::now();
let selected = TokenRouter::new(state.db.clone())
.resolve_selected(&model)
.await?;
let body = body_with_selected_model(body, &selected);
let result = executors::execute_image_generation(
state.clone(),
selected_channel_to_route(selected.clone()),
body,
)
.await;
match result {
Ok(payload) => {
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&state,
key_id,
&selected,
&model,
Some(elapsed_ms_since(started_at)),
0,
0,
cost,
1,
false,
)
.await;
record_billing(&state, key_id, &model, 0, 0, cost).await;
Ok(Json(payload))
}
Err(err) => {
record_proxy_log_failure(
&state,
key_id,
Some(&selected),
&model,
upstream_status_from_error(&err),
Some(elapsed_ms_since(started_at)),
Some(&err.to_string()),
1,
false,
)
.await;
Err(err)
}
}
}
/// POST /v1/images/edits
@@ -452,21 +497,63 @@ async fn images_edits(
}
let model = model.ok_or_else(|| AppError::BadRequest("model is required".into()))?;
let route = resolve_route(&state, &model).await?;
let started_at = Instant::now();
let selected = TokenRouter::new(state.db.clone())
.resolve_selected(&model)
.await?;
// Override the multipart `model` field with the conductor-resolved
// upstream model name (route may map "dall-e-3" → upstream-specific id).
if route.model != model {
if selected.actual_model != model {
for field in fields.iter_mut() {
if field.name == "model" && field.bytes.is_none() {
field.text = Some(route.model.clone());
field.text = Some(selected.actual_model.clone());
}
}
}
let result = executors::execute_image_edit(state.clone(), route, fields).await?;
record_billing(&state, key_id, &model, 0, 0, 0.0).await;
Ok(Json(result))
let result = executors::execute_image_edit(
state.clone(),
selected_channel_to_route(selected.clone()),
fields,
)
.await;
match result {
Ok(payload) => {
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&state,
key_id,
&selected,
&model,
Some(elapsed_ms_since(started_at)),
0,
0,
cost,
1,
false,
)
.await;
record_billing(&state, key_id, &model, 0, 0, cost).await;
Ok(Json(payload))
}
Err(err) => {
record_proxy_log_failure(
&state,
key_id,
Some(&selected),
&model,
upstream_status_from_error(&err),
Some(elapsed_ms_since(started_at)),
Some(&err.to_string()),
1,
false,
)
.await;
Err(err)
}
}
}
/// GET /v1/models
@@ -564,6 +651,23 @@ async fn record_billing(
}
}
fn estimate_request_cost(
selected: &SelectedChannel,
prompt_tokens: i64,
completion_tokens: i64,
) -> f64 {
let unit_cost = selected.account_unit_cost.unwrap_or(0.0).max(0.0);
if unit_cost <= 0.0 {
return 0.0;
}
let total_tokens = (prompt_tokens.max(0) + completion_tokens.max(0)) as f64;
if total_tokens <= 0.0 {
unit_cost
} else {
total_tokens / 1_000_000.0 * unit_cost
}
}
async fn record_proxy_log_success(
state: &AppState,
key_id: Option<i64>,
@@ -699,6 +803,31 @@ fn body_with_selected_model(mut body: Value, selected: &SelectedChannel) -> Valu
body
}
fn ensure_stream_usage(mut body: Value) -> Value {
if let Some(object) = body.as_object_mut() {
match object.get_mut("stream_options") {
Some(options) => {
if let Some(map) = options.as_object_mut() {
map.entry("include_usage".to_string())
.or_insert_with(|| Value::Bool(true));
} else {
object.insert(
"stream_options".to_string(),
serde_json::json!({ "include_usage": true }),
);
}
}
None => {
object.insert(
"stream_options".to_string(),
serde_json::json!({ "include_usage": true }),
);
}
}
}
body
}
fn build_conductor(state: &AppState) -> ProxyConductor {
let db = state.db.clone();
let select_db = db.clone();
@@ -797,6 +926,7 @@ async fn execute_messages_with_conductor(
attempts,
} => {
let (prompt, completion) = extract_usage(&response);
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
@@ -805,7 +935,7 @@ async fn execute_messages_with_conductor(
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
cost,
attempts,
false,
)
@@ -816,7 +946,7 @@ async fn execute_messages_with_conductor(
&billing_model,
prompt,
completion,
0.0,
cost,
)
.await;
Ok(response)
@@ -938,6 +1068,7 @@ async fn execute_messages_stream_with_conductor(
}
};
let cost = estimate_request_cost(&selected, 0, 0);
record_proxy_log_success(
&billing_state,
key_id,
@@ -946,12 +1077,12 @@ async fn execute_messages_stream_with_conductor(
Some(elapsed_ms_since(started_at)),
0,
0,
0.0,
cost,
attempts,
true,
)
.await;
record_billing(&billing_state, key_id, &billing_model, 0, 0, 0.0).await;
record_billing(&billing_state, key_id, &billing_model, 0, 0, cost).await;
Ok(chunks)
}
@@ -1005,6 +1136,7 @@ async fn execute_chat_completions_with_conductor(
attempts,
} => {
let (prompt, completion) = extract_usage(&response);
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
@@ -1013,7 +1145,7 @@ async fn execute_chat_completions_with_conductor(
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
cost,
attempts,
false,
)
@@ -1024,7 +1156,7 @@ async fn execute_chat_completions_with_conductor(
&billing_model,
prompt,
completion,
0.0,
cost,
)
.await;
Ok(response)
@@ -1073,7 +1205,8 @@ async fn execute_chat_completions_stream_with_conductor(
downstream_policy: None,
attempt_fn: Box::new(move |ctx| {
let state = state.clone();
let body = body_with_selected_model(body.clone(), &ctx.selected);
let body =
ensure_stream_usage(body_with_selected_model(body.clone(), &ctx.selected));
async move {
match executors::execute_chat_completions_stream(
state,
@@ -1160,8 +1293,11 @@ async fn execute_chat_completions_stream_with_conductor(
.iter()
.rev()
.find_map(|chunk| {
let payload = chunk.strip_prefix("data: ")?;
if payload.trim() == "[DONE]" {
let payload = chunk
.strip_prefix("data: ")
.or_else(|| chunk.strip_prefix("data:"))?
.trim();
if payload == "[DONE]" {
return None;
}
let value: Value = serde_json::from_str(payload).ok()?;
@@ -1173,6 +1309,7 @@ async fn execute_chat_completions_stream_with_conductor(
}
})
.unwrap_or((0, 0));
let cost = estimate_request_cost(&selected, prompt, completion);
record_proxy_log_success(
&billing_state,
key_id,
@@ -1181,7 +1318,7 @@ async fn execute_chat_completions_stream_with_conductor(
Some(elapsed_ms_since(started_at)),
prompt,
completion,
0.0,
cost,
attempts,
true,
)
@@ -1192,7 +1329,7 @@ async fn execute_chat_completions_stream_with_conductor(
&billing_model,
prompt,
completion,
0.0,
cost,
)
.await;
Ok(chunks)