forge-guardrails 0.1.2

Foundation types for an LLM-agent workflow framework
Documentation
use crate::context::manager::ContextManager;
use crate::core::message::Message;

pub(super) enum ContextAccess<'a> {
    Direct(&'a mut ContextManager),
    Shared(&'a tokio::sync::Mutex<ContextManager>),
}

impl ContextAccess<'_> {
    pub(super) async fn maybe_compact(
        &mut self,
        messages: &[Message],
        step_index: i64,
        step_hint: Option<&str>,
    ) -> Option<Vec<Message>> {
        match self {
            Self::Direct(context) => match context.maybe_compact(messages, step_index, step_hint) {
                std::borrow::Cow::Owned(new_msgs) => Some(new_msgs),
                std::borrow::Cow::Borrowed(_) => None,
            },
            Self::Shared(context) => {
                let mut context = context.lock().await;
                match context.maybe_compact(messages, step_index, step_hint) {
                    std::borrow::Cow::Owned(new_msgs) => Some(new_msgs),
                    std::borrow::Cow::Borrowed(_) => None,
                }
            }
        }
    }

    pub(super) async fn check_thresholds(&mut self, messages: &[Message]) -> Option<String> {
        match self {
            Self::Direct(context) => context.check_thresholds(messages),
            Self::Shared(context) => {
                let mut context = context.lock().await;
                context.check_thresholds(messages)
            }
        }
    }

    pub(super) async fn update_token_count(&mut self, count: i64) {
        match self {
            Self::Direct(context) => context.update_token_count(count),
            Self::Shared(context) => {
                let mut context = context.lock().await;
                context.update_token_count(count);
            }
        }
    }
}