Skip to content

feat(libsy): Pass models at runtime not construction time #630

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
grahamking wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

feat(libsy): Pass models at runtime not construction time #630

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
71 changes: 45 additions & 26 deletions crates/libsy-llm-client/src/run.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ use http::StatusCode;
use parking_lot::Mutex;
use switchyard_libsy::{Algorithm, CallModel, LibsyError, Result, RoutingOutcome, drive};
use switchyard_protocol::{
LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason,
Category, LlmClientError, ModelId, Request, Response, RoutedLlmClient, RoutingFallbackReason,
};
use switchyard_translation::prepare_request_for_target;

Expand All @@ -44,6 +44,7 @@ pub async fn run(
algorithm: Arc<dyn Algorithm>,
clients: ClientRouter,
request: Request,
models: HashMap<Category, Vec<ModelId>>,
observer: Option<RunObserver>,
) -> Result<(ModelId, Response)> {
let algorithm_name = algorithm.name().to_string();
Expand All @@ -52,7 +53,7 @@ pub async fn run(
// This says if we have an observer, put Some(..) in routing_observations.
// No observer means we don't want any routing_observations.
let routing_observations = observer.as_ref().map(|_| Arc::new(Mutex::new(Vec::new())));
let outcome = drive(algorithm, request, {
let outcome = drive(algorithm, request, models, {
let routing_observations = routing_observations.clone();
move |call| serve(routing_clients.clone(), call, routing_observations.clone())
})
Expand Down Expand Up @@ -110,9 +111,10 @@ pub async fn decide(
algorithm: Arc<dyn Algorithm>,
clients: ClientRouter,
request: Request,
models: HashMap<Category, Vec<ModelId>>,
) -> Result<RoutingOutcome> {
let routing_clients = clients.clone();
let mut outcome = drive(algorithm, request, move |call| {
let mut outcome = drive(algorithm, request, models, move |call| {
serve(routing_clients.clone(), call, None)
})
.await?;
Expand Down Expand Up @@ -455,9 +457,7 @@ mod tests {

use crate::{Backend, HttpBackendConfig, ModelConfig, TranslatingLlmClient};

struct CandidateAlgorithm {
models: Vec<ModelId>,
}
struct CandidateAlgorithm {}

struct AnsweredAlgorithm {
model: ModelId,
Expand All @@ -471,13 +471,14 @@ mod tests {

async fn route(
self: Arc<Self>,
_driver: Driver,
driver: Driver,
request: Request,
) -> Result<RoutingOutcome> {
let selected_model = self.models.first().cloned().ok_or(LibsyError::NoTargets)?;
let models = driver.models_for(Category::Any);
let selected_model = models.first().cloned().ok_or(LibsyError::NoTargets)?;
Ok(RoutingOutcome::route_to(
selected_model,
self.models.iter().skip(1).cloned().collect(),
models.iter().skip(1).cloned().collect(),
request,
))
}
Expand Down Expand Up @@ -605,19 +606,27 @@ mod tests {
requests: Mutex::new(Vec::new()),
first,
});
let algorithm = Arc::new(CandidateAlgorithm {
models: vec!["weak".into(), "strong".into()],
});
let algorithm = Arc::new(CandidateAlgorithm {});
let models = to_category_map(&["weak", "strong"]);
let result = run(
algorithm,
ClientRouter::single(client.clone()),
request(),
models,
None,
)
.await;
(client, result)
}

fn to_category_map(names: &[&str]) -> HashMap<Category, Vec<ModelId>> {
[(
Category::Any,
names.iter().map(|name| ModelId::from(*name)).collect(),
)]
.into()
}

#[tokio::test]
async fn answered_outcome_does_not_make_a_second_model_call() -> Result<()> {
let client = Arc::new(CandidateClient {
Expand All @@ -635,6 +644,7 @@ mod tests {
}),
ClientRouter::single(client.clone()),
request(),
HashMap::new(),
Some(observer),
)
.await?;
Expand Down Expand Up @@ -673,11 +683,10 @@ mod tests {
);

run(
Arc::new(CandidateAlgorithm {
models: vec!["weak".into(), "strong".into()],
}),
Arc::new(CandidateAlgorithm {}),
clients,
request(),
to_category_map(&["weak", "strong"]),
None,
)
.await?;
Expand Down Expand Up @@ -709,6 +718,7 @@ mod tests {
}),
clients,
request(),
HashMap::new(),
)
.await?;

Expand Down Expand Up @@ -739,11 +749,10 @@ mod tests {
);

let outcome = decide(
Arc::new(CandidateAlgorithm {
models: vec!["weak".into(), "strong".into()],
}),
Arc::new(CandidateAlgorithm {}),
clients,
request(),
to_category_map(&["weak", "strong"]),
)
.await?;

Expand Down Expand Up @@ -880,10 +889,15 @@ mod tests {
])
.map_err(|error| LibsyError::external("building test client", error))?,
);
let algorithm = Arc::new(CandidateAlgorithm {
models: vec!["weak".into(), "strong".into()],
});
run(algorithm, ClientRouter::single(client), request(), None).await?;
let algorithm = Arc::new(CandidateAlgorithm {});
run(
algorithm,
ClientRouter::single(client),
request(),
to_category_map(&["weak", "strong"]),
None,
)
.await?;

assert_eq!(&*calls.lock(), &["weak", "weak", "weak", "strong"]);
Ok(())
Expand Down Expand Up @@ -971,17 +985,22 @@ mod tests {
])
.expect("building test client"),
);
let algorithm = Arc::new(CandidateAlgorithm {
models: vec!["weak".into(), "strong".into()],
});
let algorithm = Arc::new(CandidateAlgorithm {});
let mut llm_request = text_request(Some("auto".to_string()), "hello".to_string());
llm_request.stream = true;
let request = Request {
llm_request,
raw_request: None,
metadata: None,
};
let result = run(algorithm, ClientRouter::single(client), request, None).await;
let result = run(
algorithm,
ClientRouter::single(client),
request,
to_category_map(&["weak", "strong"]),
None,
)
.await;
(server, calls, result)
}

Expand Down
48 changes: 39 additions & 9 deletions crates/libsy-llm-client/tests/observability.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@
//! and model name. Counters are cumulative across flushes; the helpers take the
//! latest (max) matching data point.

use std::collections::BTreeMap;
use std::collections::{BTreeMap, HashMap};
use std::fmt;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, OnceLock};
Expand Down Expand Up @@ -39,7 +39,7 @@ use switchyard_libsy::{
StageRouterConfig, Step, TaskClassifierConfig,
};
use switchyard_llm_client::{ClientRouter, RunObservation, RunObserver};
use switchyard_protocol::ModelId;
use switchyard_protocol::{Category, ModelId};
use switchyard_protocol::{
ContentBlock, LlmRequest, LlmResponse, Message, Metadata, Request, Response, Role,
RoutedLlmClient, ToolCall, ToolResult, Usage, WireFormat,
Expand Down Expand Up @@ -568,7 +568,14 @@ async fn run(
client: Arc<dyn RoutedLlmClient>,
request: Request,
) -> switchyard_libsy::Result<(ModelId, Response)> {
switchyard_llm_client::run(algorithm, ClientRouter::single(client), request, None).await
switchyard_llm_client::run(
algorithm,
ClientRouter::single(client),
request,
HashMap::new(),
None,
)
.await
}

fn classifier_router(
Expand Down Expand Up @@ -714,6 +721,14 @@ async fn affinity_keeps_the_algorithm_selection_after_client_fallback()
Arc::clone(&router),
ClientRouter::single(client.clone()),
request.clone(),
Category::to_map(
Category::Any,
&[
"affinity-fallback-strong",
"affinity-fallback-weak",
"affinity-fallback-judge",
],
),
None,
)
.await?;
Expand All @@ -724,9 +739,14 @@ async fn affinity_keeps_the_algorithm_selection_after_client_fallback()
);

client.efficient_available.store(true, Ordering::Relaxed);
let (selected, second_response) =
switchyard_llm_client::run(router, ClientRouter::single(client.clone()), request, None)
.await?;
let (selected, second_response) = switchyard_llm_client::run(
router,
ClientRouter::single(client.clone()),
request,
HashMap::new(),
None,
)
.await?;
assert_eq!(selected, "affinity-fallback-weak");
assert_eq!(
second_response.served_model().map(ModelId::as_str),
Expand Down Expand Up @@ -1050,6 +1070,7 @@ async fn observed_run_reports_one_successful_routed_call() -> switchyard_libsy::
algo(ALGO, MODEL),
ClientRouter::single(client),
request_with_metadata("observed-session", "observed-correlation"),
HashMap::new(),
Some(observer),
)
.await?;
Expand Down Expand Up @@ -1224,7 +1245,10 @@ async fn failed_call_records_error_outcome_and_warn_logs() -> switchyard_libsy::
name: ALGO.to_string(),
target: MODEL.into(),
});
let stream = algorithm.run_stream(request_with_metadata("obs-session-2", "obs-corr-2"));
let stream = algorithm.run_stream(
request_with_metadata("obs-session-2", "obs-corr-2"),
HashMap::new(),
);
tokio::pin!(stream);

let mut saw_error_step = false;
Expand Down Expand Up @@ -1476,7 +1500,10 @@ async fn in_flight_gauge_reads_a_run_parked_on_an_unanswered_routing_call()
name: ALGO.to_string(),
target: MODEL.into(),
});
let stream = algorithm.run_stream(request_with_metadata("obs-session-if", "obs-corr-if"));
let stream = algorithm.run_stream(
request_with_metadata("obs-session-if", "obs-corr-if"),
HashMap::new(),
);
tokio::pin!(stream);
let attributes = [("algorithm", ALGO)];

Expand Down Expand Up @@ -1531,7 +1558,10 @@ async fn in_flight_gauge_clears_when_a_run_is_abandoned() -> switchyard_libsy::R
// disconnected client abandons a run. The run task is aborted mid-await and never
// reaches the code that follows it, so only a drop can return the count.
{
let stream = algorithm.run_stream(request_with_metadata("obs-session-ab", "obs-corr-ab"));
let stream = algorithm.run_stream(
request_with_metadata("obs-session-ab", "obs-corr-ab"),
HashMap::new(),
);
tokio::pin!(stream);
let Some(Ok(Step::CallModel(_call))) = stream.next().await else {
return Err(test_error("expected an offloaded routing call"));
Expand Down
Loading
Loading