246 lines
9 KiB
Rust
246 lines
9 KiB
Rust
//! Sampler configuration types.
|
|
//!
|
|
//! [`SamplerConfig`] is the per-request configuration handed to the
|
|
//! sampler. It deliberately does **not** alias
|
|
//! `xai_grok_sampling_types::SamplingConfig` so that the sampler crate
|
|
//! avoids transitive dependencies on shell-specific types
|
|
//! (`xai-grok-tools`, etc.).
|
|
|
|
use indexmap::IndexMap;
|
|
use serde::{Deserialize, Serialize};
|
|
use xai_grok_sampling_types::{
|
|
ApiBackend, CompactionAtTokens, CompactionsRemaining, DoomLoopRecoveryPolicy, ReasoningEffort,
|
|
};
|
|
|
|
use crate::attribution::SharedAttributionCallback;
|
|
use crate::retry::{DEFAULT_MAX_RETRIES, RATE_LIMIT_RETRY_THRESHOLD};
|
|
|
|
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
|
|
#[serde(rename_all = "snake_case")]
|
|
pub enum AuthScheme {
|
|
#[default]
|
|
Bearer,
|
|
XApiKey,
|
|
}
|
|
|
|
/// All knobs that control a single sampling request.
|
|
///
|
|
/// The session typically owns one `SamplerConfig` per active model
|
|
/// and passes it (or a per-request override) to the actor on every
|
|
/// submit.
|
|
///
|
|
/// # Construction in `xai-grok-shell`
|
|
///
|
|
/// `SamplerConfig` is the single source of truth for sampler
|
|
/// configuration. The shell builds it directly (see
|
|
/// `agent::config::resolve_model_to_sampling_config` and
|
|
/// `session::acp_session::SessionActor::reconstruct_full_config`) by
|
|
/// composing chat-state's `xai_grok_sampling_types::SamplingConfig`
|
|
/// with `Credentials` (api key, client version).
|
|
///
|
|
/// URL-derived request headers (e.g. `X-XAI-Token-Auth` for the
|
|
/// cli-chat-proxy) are
|
|
/// folded into [`Self::extra_headers`] by
|
|
/// `agent::config::inject_url_derived_headers` before the
|
|
/// `SamplerConfig` is handed to the actor. Auth is selected separately
|
|
/// via `auth_scheme`, while `api_backend` controls only the request/response
|
|
/// protocol shape.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct SamplerConfig {
|
|
pub api_key: Option<String>,
|
|
pub base_url: String,
|
|
pub model: String,
|
|
pub max_completion_tokens: Option<u32>,
|
|
pub temperature: Option<f32>,
|
|
pub top_p: Option<f32>,
|
|
pub api_backend: ApiBackend,
|
|
#[serde(default)]
|
|
pub auth_scheme: AuthScheme,
|
|
/// Extra request headers applied verbatim. The sampler never inspects
|
|
/// the URL to derive headers; callers (the session) inject proxy auth
|
|
/// and other access headers here before constructing the config.
|
|
pub extra_headers: IndexMap<String, String>,
|
|
/// Total context window size in tokens. The sampler does not enforce
|
|
/// it; it is informational metadata used by the session for compaction
|
|
/// decisions.
|
|
pub context_window: u64,
|
|
pub force_http1: bool,
|
|
pub max_retries: Option<u32>,
|
|
pub stream_tool_calls: bool,
|
|
pub idle_timeout_secs: Option<u64>,
|
|
|
|
// Reasoning effort
|
|
pub reasoning_effort: Option<ReasoningEffort>,
|
|
|
|
// Client identity
|
|
pub origin_client: Option<OriginClientInfo>,
|
|
pub client_identifier: Option<String>,
|
|
pub deployment_id: Option<String>,
|
|
pub user_id: Option<String>,
|
|
pub client_version: Option<String>,
|
|
|
|
/// Optional hook invoked at every UNAUTHORIZED (401) response
|
|
/// site. The sampler passes the bearer that was actually sent on
|
|
/// the wire to the callback; the implementation is free to do
|
|
/// whatever it wants with it (typically: join it with a live
|
|
/// credential source and emit an attribution event for diagnosis
|
|
/// of stale-token vs. server-rejected-live-token 401s). `None`
|
|
/// (default) is a no-op -- the 401 arm returns the same
|
|
/// `SamplingError::Auth` it always did.
|
|
///
|
|
/// `Arc<dyn Trait>` is not serializable, so the field is skipped
|
|
/// in (de)serialization. Round-tripping a config through serde
|
|
/// drops the callback; callers that deserialize a `SamplerConfig`
|
|
/// from disk must re-attach the callback before passing it to
|
|
/// [`crate::SamplingClient::new`] or 401 attribution will be
|
|
/// silently disabled for the rebuilt client.
|
|
#[serde(skip)]
|
|
pub attribution_callback: Option<SharedAttributionCallback>,
|
|
|
|
/// Live bearer resolve per request. `None` uses construction-time `api_key`.
|
|
#[serde(skip)]
|
|
pub bearer_resolver: Option<SharedBearerResolver>,
|
|
|
|
#[serde(default)]
|
|
pub supports_backend_search: bool,
|
|
|
|
/// Per-model config for the `x-compactions-remaining` header; `None` disables it.
|
|
#[serde(default)]
|
|
pub compactions_remaining: Option<CompactionsRemaining>,
|
|
|
|
/// Per-model config for the `x-compaction-at` header; `None` disables it.
|
|
#[serde(default)]
|
|
pub compaction_at_tokens: Option<CompactionAtTokens>,
|
|
|
|
/// Server-side doom-loop check policy; `None` disables it. When set, the
|
|
/// client itself sends the opt-in `x-grok-doom-loop-check` header on
|
|
/// streaming Responses API requests and absorbs the reported trigger
|
|
/// events (unlike the environment headers in [`Self::extra_headers`],
|
|
/// this header gates the client's own decode behavior, so it lives with
|
|
/// the decoder).
|
|
#[serde(default)]
|
|
pub doom_loop_recovery: Option<DoomLoopRecoveryPolicy>,
|
|
|
|
/// Per-request header injector (e.g. OTel traceparent). Called in `post()`.
|
|
#[serde(skip)]
|
|
pub header_injector: Option<SharedHeaderInjector>,
|
|
}
|
|
|
|
impl Default for SamplerConfig {
|
|
/// Empty defaults so callers can use `..Default::default()` and
|
|
/// new fields don't ripple through every literal site.
|
|
fn default() -> Self {
|
|
Self {
|
|
api_key: None,
|
|
base_url: String::new(),
|
|
model: String::new(),
|
|
max_completion_tokens: None,
|
|
temperature: None,
|
|
top_p: None,
|
|
api_backend: ApiBackend::default(),
|
|
auth_scheme: AuthScheme::default(),
|
|
extra_headers: IndexMap::new(),
|
|
context_window: 0,
|
|
force_http1: false,
|
|
max_retries: None,
|
|
stream_tool_calls: false,
|
|
idle_timeout_secs: None,
|
|
reasoning_effort: None,
|
|
origin_client: None,
|
|
client_identifier: None,
|
|
deployment_id: None,
|
|
user_id: None,
|
|
client_version: None,
|
|
attribution_callback: None,
|
|
bearer_resolver: None,
|
|
supports_backend_search: false,
|
|
compactions_remaining: None,
|
|
compaction_at_tokens: None,
|
|
doom_loop_recovery: None,
|
|
header_injector: None,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Cheap sync read of the current bearer for [`SamplerConfig::bearer_resolver`].
|
|
pub trait BearerResolver: Send + Sync + std::fmt::Debug {
|
|
fn current_bearer(&self) -> Option<String>;
|
|
}
|
|
|
|
pub type SharedBearerResolver = std::sync::Arc<dyn BearerResolver>;
|
|
|
|
/// Per-request header injection (e.g. OTel `traceparent`).
|
|
pub trait HeaderInjector: Send + Sync + std::fmt::Debug {
|
|
fn inject(&self, headers: &mut reqwest::header::HeaderMap);
|
|
}
|
|
|
|
pub type SharedHeaderInjector = std::sync::Arc<dyn HeaderInjector>;
|
|
|
|
/// Retry knobs for the sampler's internal transport-error retry loop.
|
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
|
pub struct RetryPolicy {
|
|
/// Maximum number of retries before giving up.
|
|
pub max_retries: u32,
|
|
/// After this many rate-limit (429) retries, escalate to the caller.
|
|
/// Lower than `max_retries` because rate-limit waits can be long.
|
|
pub rate_limit_retry_threshold: u32,
|
|
}
|
|
|
|
impl Default for RetryPolicy {
|
|
fn default() -> Self {
|
|
Self {
|
|
max_retries: DEFAULT_MAX_RETRIES,
|
|
rate_limit_retry_threshold: RATE_LIMIT_RETRY_THRESHOLD,
|
|
}
|
|
}
|
|
}
|
|
|
|
/// Identity of the client that originated the request, used for
|
|
/// User-Agent rendering. The shell layer composes this with platform
|
|
/// info into a final UA string.
|
|
#[derive(Clone, Debug, Eq, PartialEq, Serialize, Deserialize)]
|
|
pub struct OriginClientInfo {
|
|
pub product: String,
|
|
pub version: Option<String>,
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn retry_policy_defaults() {
|
|
let policy = RetryPolicy::default();
|
|
assert_eq!(policy.max_retries, DEFAULT_MAX_RETRIES);
|
|
assert_eq!(
|
|
policy.rate_limit_retry_threshold,
|
|
RATE_LIMIT_RETRY_THRESHOLD
|
|
);
|
|
}
|
|
|
|
/// Configs serialized before the field existed must keep deserializing.
|
|
#[test]
|
|
fn config_without_doom_loop_recovery_deserializes_to_none() {
|
|
let mut stripped = serde_json::to_value(SamplerConfig::default()).unwrap();
|
|
stripped
|
|
.as_object_mut()
|
|
.unwrap()
|
|
.remove("doom_loop_recovery");
|
|
let config: SamplerConfig = serde_json::from_value(stripped).unwrap();
|
|
assert!(config.doom_loop_recovery.is_none());
|
|
|
|
let with_policy = SamplerConfig {
|
|
doom_loop_recovery: Some(DoomLoopRecoveryPolicy {
|
|
max_threshold: 8,
|
|
max_retries: 2,
|
|
}),
|
|
..Default::default()
|
|
};
|
|
let round_tripped: SamplerConfig =
|
|
serde_json::from_value(serde_json::to_value(&with_policy).unwrap()).unwrap();
|
|
assert_eq!(
|
|
round_tripped.doom_loop_recovery,
|
|
with_policy.doom_loop_recovery
|
|
);
|
|
}
|
|
}
|