synwire_agent/sampling.rs
1//! Direct model sampling provider.
2//!
3//! Uses the configured chat model directly for standalone (non-MCP) mode.
4//! When running the agent without an MCP host this provider is used in place
5//! of `McpSampling` from the MCP server crate.
6
7use std::sync::Arc;
8
9use synwire_core::{BoxFuture, SamplingError, SamplingProvider, SamplingRequest, SamplingResponse};
10
11/// Type alias for the callback used to invoke a language model.
12///
13/// The function receives a [`SamplingRequest`] and returns a boxed future that
14/// resolves to a [`SamplingResponse`] or a [`SamplingError`].
15pub type SamplingFn = Arc<
16 dyn Fn(SamplingRequest) -> BoxFuture<'static, Result<SamplingResponse, SamplingError>>
17 + Send
18 + Sync,
19>;
20
21/// A sampling provider backed by direct model invocation.
22///
23/// Used when running the agent without an MCP host. Wraps a caller-supplied
24/// callback that performs the actual model call. When no callback is provided
25/// (via [`DirectModelSampling::unavailable`]), all sampling calls return
26/// [`SamplingError::NotAvailable`] to trigger graceful degradation in callers.
27pub struct DirectModelSampling {
28 invoke: Option<SamplingFn>,
29}
30
31impl DirectModelSampling {
32 /// Create a new provider with the given model invocation callback.
33 ///
34 /// The callback is wrapped in an `Arc` so the provider remains `Clone`-friendly
35 /// and can be shared across tasks.
36 #[must_use]
37 pub fn new(
38 invoke: impl Fn(SamplingRequest) -> BoxFuture<'static, Result<SamplingResponse, SamplingError>>
39 + Send
40 + Sync
41 + 'static,
42 ) -> Self {
43 Self {
44 invoke: Some(Arc::new(invoke)),
45 }
46 }
47
48 /// Create a provider that always reports sampling as unavailable.
49 ///
50 /// [`SamplingProvider::is_available`] will return `false` and all calls to
51 /// [`SamplingProvider::sample`] will return [`SamplingError::NotAvailable`].
52 #[must_use]
53 pub const fn unavailable() -> Self {
54 Self { invoke: None }
55 }
56}
57
58impl SamplingProvider for DirectModelSampling {
59 fn is_available(&self) -> bool {
60 self.invoke.is_some()
61 }
62
63 fn sample(
64 &self,
65 request: SamplingRequest,
66 ) -> BoxFuture<'_, Result<SamplingResponse, SamplingError>> {
67 match &self.invoke {
68 Some(invoke) => invoke(request),
69 None => Box::pin(async { Err(SamplingError::NotAvailable) }),
70 }
71 }
72}