feat: update UI

This commit is contained in:
leokun
2026-08-26 22:01:14 +08:00
parent 5450fc76e2
commit ded534fecf
27 changed files with 701 additions and 205 deletions
+2 -2
View File
@@ -130,8 +130,8 @@ pub fn api_router(service: ControlService) -> Router {
put(models::update).delete(models::remove),
)
.route(
"/__byok-api__/api/models/{model_hash}/test",
post(models::test),
"/__byok-api__/api/models/{model_hash}/test/{test_id}",
post(models::test).delete(models::cancel),
)
.route("/__byok-api__/api/llm-calls", get(calls::list))
.route("/__byok-api__/api/llm-calls/{call_id}", get(calls::detail))
+10 -2
View File
@@ -64,9 +64,17 @@ pub async fn update(
pub async fn test(
State(service): State<ControlService>,
Path(model_hash): Path<String>,
Path((model_hash, test_id)): Path<(String, String)>,
) -> Result<Json<ModelConnectivityResult>> {
Ok(Json(service.test_model(&model_hash).await?))
Ok(Json(service.test_model(&model_hash, &test_id).await?))
}
pub async fn cancel(
State(service): State<ControlService>,
Path((_model_hash, test_id)): Path<(String, String)>,
) -> Result<StatusCode> {
service.cancel_model_test(&test_id);
Ok(StatusCode::NO_CONTENT)
}
pub async fn discover(
+133 -22
View File
@@ -1,4 +1,8 @@
use std::{collections::BTreeSet, sync::Arc, time::Instant};
use std::{
collections::{BTreeMap, BTreeSet},
sync::{Arc, Mutex},
time::Instant,
};
use base64::{engine::general_purpose::STANDARD, Engine};
use futures_util::StreamExt;
@@ -33,6 +37,7 @@ pub struct ControlService {
store: Store,
cursor_harness: CursorHarness,
provider: Arc<dyn Provider>,
model_tests: Arc<Mutex<BTreeMap<String, CancellationToken>>>,
}
#[derive(Clone, Debug, Serialize)]
@@ -142,6 +147,7 @@ impl ControlService {
cursor_harness: CursorHarness::new(store.clone())?,
store,
provider,
model_tests: Arc::new(Mutex::new(BTreeMap::new())),
})
}
@@ -243,7 +249,49 @@ impl ControlService {
self.store.update_model(model_hash, input).await
}
pub async fn test_model(&self, model_hash: &str) -> Result<ModelConnectivityResult> {
pub async fn test_model(
&self,
model_hash: &str,
test_id: &str,
) -> Result<ModelConnectivityResult> {
let cancellation = CancellationToken::new();
let cancellation = {
let mut tests = self
.model_tests
.lock()
.expect("model test registry mutex poisoned");
tests
.entry(test_id.to_owned())
.or_insert_with(|| cancellation.clone())
.clone()
};
let result = self.run_model_test(model_hash, cancellation).await;
self.model_tests
.lock()
.expect("model test registry mutex poisoned")
.remove(test_id);
result
}
pub fn cancel_model_test(&self, test_id: &str) {
let cancellation = {
let mut tests = self
.model_tests
.lock()
.expect("model test registry mutex poisoned");
tests
.entry(test_id.to_owned())
.or_insert_with(CancellationToken::new)
.clone()
};
cancellation.cancel();
}
async fn run_model_test(
&self,
model_hash: &str,
cancellation: CancellationToken,
) -> Result<ModelConnectivityResult> {
const TEST_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(45);
const TEST_PROMPT: &str = "Output the numbers 1 through 120 separated by a single space. No commas, no newlines, no explanation.";
@@ -255,12 +303,11 @@ impl ControlService {
let mut model = ModelSpec::new(model_hash);
configured.configure(&mut model);
model.max_output_tokens = Some(configured.max_output_tokens().unwrap_or(65_536));
let test_id = format!("model-test-{}", uuid::Uuid::new_v4());
let call_id = test_id.clone();
let call_id = format!("model-test-{}", uuid::Uuid::new_v4());
let invocation = ModelInvocation {
call_id: test_id.clone(),
run_id: test_id.clone(),
conversation_id: test_id,
call_id: call_id.clone(),
run_id: call_id.clone(),
conversation_id: call_id.clone(),
provider_call_index: 0,
request: ModelRequest {
prompt: PromptSpec {
@@ -277,7 +324,6 @@ impl ControlService {
}],
},
};
let cancellation = CancellationToken::new();
let started = Instant::now();
let mut first_text_at = None;
let mut output_tokens = None;
@@ -305,6 +351,9 @@ impl ControlService {
_ => {}
}
}
if cancellation.is_cancelled() {
return Err(Error::Cancelled);
}
if !finished {
return Err(Error::Protocol(
"provider stream ended without Done during connectivity test".into(),
@@ -792,7 +841,7 @@ mod tests {
use tokio_util::sync::CancellationToken;
use crate::{
model::{ModelConfigInput, ModelInvocation, ModelType, ProjectedContent},
model::{ModelConfig, ModelConfigInput, ModelInvocation, ModelType, ProjectedContent},
provider::{FinishReason, ModelEvent, Provider, ProviderStream},
store::Store,
};
@@ -803,6 +852,10 @@ mod tests {
invocation: Arc<Mutex<Option<ModelInvocation>>>,
}
struct CancellationProvider {
started: Arc<tokio::sync::Notify>,
}
impl Provider for TestProvider {
fn stream(
&self,
@@ -826,17 +879,23 @@ mod tests {
}
}
#[tokio::test]
async fn connectivity_test_uses_the_configured_llm_provider() {
let directory = tempfile::tempdir().unwrap();
let store = Store::connect(&format!(
"sqlite://{}",
directory.path().join("control.db").display()
))
.await
.unwrap();
let invocation = Arc::new(Mutex::new(None));
let model = store
impl Provider for CancellationProvider {
fn stream(
&self,
_invocation: ModelInvocation,
cancellation: CancellationToken,
) -> ProviderStream {
let started = self.started.clone();
Box::pin(async_stream::try_stream! {
started.notify_one();
cancellation.cancelled().await;
if false { yield ModelEvent::TextStart; }
})
}
}
async fn create_test_model(store: &Store) -> ModelConfig {
store
.create_model(&ModelConfigInput {
model_id: "reasoning-model".into(),
display_name: "Reasoning Model".into(),
@@ -861,7 +920,20 @@ mod tests {
thinking_budget_tokens: None,
})
.await
.unwrap();
.unwrap()
}
#[tokio::test]
async fn connectivity_test_uses_the_configured_llm_provider() {
let directory = tempfile::tempdir().unwrap();
let store = Store::connect(&format!(
"sqlite://{}",
directory.path().join("control.db").display()
))
.await
.unwrap();
let invocation = Arc::new(Mutex::new(None));
let model = create_test_model(&store).await;
let service = ControlService::new(
store,
Arc::new(TestProvider {
@@ -870,7 +942,10 @@ mod tests {
)
.unwrap();
let result = service.test_model(&model.model_hash).await.unwrap();
let result = service
.test_model(&model.model_hash, "test-id")
.await
.unwrap();
assert_eq!(result.output, "OK");
assert_eq!(result.output_tokens, 2);
@@ -892,6 +967,42 @@ mod tests {
));
}
#[tokio::test]
async fn connectivity_test_can_be_cancelled() {
let directory = tempfile::tempdir().unwrap();
let store = Store::connect(&format!(
"sqlite://{}",
directory.path().join("cancel.db").display()
))
.await
.unwrap();
let model = create_test_model(&store).await;
let started = Arc::new(tokio::sync::Notify::new());
let service = ControlService::new(
store,
Arc::new(CancellationProvider {
started: started.clone(),
}),
)
.unwrap();
let running_service = service.clone();
let model_hash = model.model_hash.clone();
let task =
tokio::spawn(
async move { running_service.test_model(&model_hash, "cancel-test").await },
);
started.notified().await;
service.cancel_model_test("cancel-test");
assert!(matches!(task.await.unwrap(), Err(crate::Error::Cancelled)));
assert!(!service
.model_tests
.lock()
.unwrap()
.contains_key("cancel-test"));
}
#[test]
fn connectivity_output_token_estimate_handles_words_and_empty_text() {
assert_eq!(super::estimate_output_tokens("1 2 3"), 3);