Skip to main content

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}