skippy-server 0.76.1

Embedded Skippy staged runtime server
use crate::cli::OpenAiGuardrailsCliMode;
use openai_frontend::CompactingOpenAiBackend;
use openai_frontend::CompactionConfig;
use openai_frontend::GuardedOpenAiBackend;
use openai_frontend::GuardrailMode;
use openai_frontend::GuardrailPolicy;
use openai_frontend::GuardrailPolicyHandle;
use openai_frontend::OpenAiBackend;
use openai_frontend::RetryExhaustionMode;
use openai_frontend::StreamingGuardrailMode;
use serde::Serialize;
use std::sync::Arc;

#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum OpenAiGuardrailsTarget {
    Skippy,
}

impl OpenAiGuardrailsTarget {
    const fn as_status_label(self) -> &'static str {
        match self {
            Self::Skippy => "skippy",
        }
    }
}

#[derive(Clone, Debug, PartialEq)]
pub struct OpenAiGuardrailsConfig {
    pub target: OpenAiGuardrailsTarget,
    pub policy: GuardrailPolicyHandle,
    pub compaction: Option<CompactionConfig>,
}

impl OpenAiGuardrailsConfig {
    pub fn disabled_for_skippy() -> Self {
        Self {
            target: OpenAiGuardrailsTarget::Skippy,
            policy: GuardrailPolicyHandle::default(),
            compaction: None,
        }
    }

    pub fn compatibility_for_skippy() -> Self {
        Self {
            target: OpenAiGuardrailsTarget::Skippy,
            policy: GuardrailPolicy {
                mode: GuardrailMode::MetricsOnly,
                apply_to_all_models: true,
                retry_exhaustion_mode: RetryExhaustionMode::PassLastText,
                ..GuardrailPolicy::default()
            }
            .into(),
            compaction: None,
        }
    }

    pub fn for_standalone_mode(mode: OpenAiGuardrailsCliMode) -> Self {
        match mode {
            OpenAiGuardrailsCliMode::Disabled => Self::disabled_for_skippy(),
            OpenAiGuardrailsCliMode::Metrics => Self::compatibility_for_skippy(),
            OpenAiGuardrailsCliMode::Enforce => Self {
                target: OpenAiGuardrailsTarget::Skippy,
                policy: GuardrailPolicy {
                    mode: GuardrailMode::Enforce,
                    apply_to_all_models: true,
                    ..GuardrailPolicy::default()
                }
                .into(),
                compaction: None,
            },
        }
    }

    pub fn status(&self) -> OpenAiGuardrailsStatus {
        let policy = self.policy.snapshot();
        OpenAiGuardrailsStatus {
            mode: guardrail_mode_label(policy.mode),
            target: self.target.as_status_label(),
            streaming: streaming_mode_label(policy.streaming_mode),
            retry_exhaustion: retry_exhaustion_label(&policy),
            small_model_policy: small_model_policy_label(&policy),
            small_param_threshold_b: policy.small_param_threshold_b,
            max_tool_retries: policy.max_tool_retries,
            max_structured_retries: policy.max_structured_retries,
        }
    }

    fn should_wrap_guardrail_backend(&self) -> bool {
        matches!(self.target, OpenAiGuardrailsTarget::Skippy)
    }

    #[cfg(test)]
    pub(super) fn wrap_backend(&self, backend: Arc<dyn OpenAiBackend>) -> Arc<dyn OpenAiBackend> {
        self.wrap_backend_with_context_limit(backend, None)
    }

    pub(super) fn wrap_backend_with_context_limit(
        &self,
        backend: Arc<dyn OpenAiBackend>,
        context_limit_tokens: Option<usize>,
    ) -> Arc<dyn OpenAiBackend> {
        let backend = self.wrap_compacting_backend(backend, context_limit_tokens);
        if self.should_wrap_guardrail_backend() {
            Arc::new(GuardedOpenAiBackend::with_policy_handle(
                backend,
                self.policy.clone(),
            ))
        } else {
            backend
        }
    }

    fn wrap_compacting_backend(
        &self,
        backend: Arc<dyn OpenAiBackend>,
        context_limit_tokens: Option<usize>,
    ) -> Arc<dyn OpenAiBackend> {
        let Some(mut compaction) = self.compaction else {
            return backend;
        };
        if compaction.context_limit_tokens.is_none() {
            compaction.context_limit_tokens = context_limit_tokens;
        }
        Arc::new(CompactingOpenAiBackend::new(backend, compaction))
    }
}

#[derive(Clone, Debug, PartialEq, Serialize)]
pub struct OpenAiGuardrailsStatus {
    pub mode: &'static str,
    pub target: &'static str,
    pub streaming: &'static str,
    pub retry_exhaustion: &'static str,
    pub small_model_policy: &'static str,
    pub small_param_threshold_b: f32,
    pub max_tool_retries: u8,
    pub max_structured_retries: u8,
}

fn guardrail_mode_label(mode: GuardrailMode) -> &'static str {
    match mode {
        GuardrailMode::Disabled => "disabled",
        GuardrailMode::MetricsOnly => "metrics",
        GuardrailMode::Enforce => "enforce",
    }
}

fn streaming_mode_label(mode: StreamingGuardrailMode) -> &'static str {
    match mode {
        StreamingGuardrailMode::PassThrough => "pass_through",
    }
}

fn retry_exhaustion_label(policy: &GuardrailPolicy) -> &'static str {
    match policy.retry_exhaustion_mode {
        RetryExhaustionMode::Error => "error",
        RetryExhaustionMode::PassLastText => "pass_last_text",
    }
}

fn small_model_policy_label(policy: &GuardrailPolicy) -> &'static str {
    if policy.apply_to_all_models {
        "all"
    } else {
        "small_models_only"
    }
}