goose-context-management 0.1.0-alpha.13

Conversation compaction for Goose
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143

//! Provider-level compaction: retry a completion that overflowed the context
//! window against a compacted history instead of surfacing the error.

use std::sync::Arc;

use async_trait::async_trait;
use goose_provider_types::base::{MessageStream, Provider};
use goose_provider_types::conversation::message::Message;
use goose_provider_types::conversation::token_usage::ProviderUsage;
use goose_provider_types::errors::ProviderError;
use goose_provider_types::model::ModelConfig;

use crate::model::ProviderModel;
use crate::templates::Templates;

/// Wraps a provider so a `ContextLengthExceeded` response triggers compaction
/// and one retry with the summary standing in for the prior history.
pub struct CompactingProvider {
    inner: Arc<dyn Provider>,
    templates: Templates,
}

impl CompactingProvider {
    pub fn new(inner: Arc<dyn Provider>) -> Self {
        Self {
            inner,
            templates: Templates::default(),
        }
    }

    pub fn with_templates(mut self, templates: Templates) -> Self {
        self.templates = templates;
        self
    }

    async fn compacted_messages(
        &self,
        model_config: &ModelConfig,
        messages: &[Message],
    ) -> Result<Vec<Message>, ProviderError> {
        let model = ProviderModel::new(self.inner.clone(), model_config.clone());
        let summary = crate::summarize(&model, None, &self.templates, messages)
            .await
            .map_err(|error| ProviderError::ExecutionError(error.to_string()))?;
        Ok(vec![summary.message])
    }
}

#[async_trait]
impl Provider for CompactingProvider {
    fn get_name(&self) -> &str {
        self.inner.get_name()
    }

    async fn stream(
        &self,
        model_config: &ModelConfig,
        system: &str,
        messages: &[Message],
        tools: &[rmcp::model::Tool],
    ) -> Result<MessageStream, ProviderError> {
        match self
            .inner
            .stream(model_config, system, messages, tools)
            .await
        {
            Err(ProviderError::ContextLengthExceeded(_)) => {
                let compacted = self.compacted_messages(model_config, messages).await?;
                self.inner
                    .stream(model_config, system, &compacted, tools)
                    .await
            }
            other => other,
        }
    }

    async fn complete(
        &self,
        model_config: &ModelConfig,
        system: &str,
        messages: &[Message],
        tools: &[rmcp::model::Tool],
    ) -> Result<(Message, ProviderUsage), ProviderError> {
        match self
            .inner
            .complete(model_config, system, messages, tools)
            .await
        {
            Err(ProviderError::ContextLengthExceeded(_)) => {
                let compacted = self.compacted_messages(model_config, messages).await?;
                self.inner
                    .complete(model_config, system, &compacted, tools)
                    .await
            }
            other => other,
        }
    }

    async fn get_context_limit(&self, model: &str, override_limit: Option<usize>) -> usize {
        self.inner.get_context_limit(model, override_limit).await
    }

    fn manages_own_context(&self) -> bool {
        true
    }

    fn supports_builtin_tools(&self) -> bool {
        self.inner.supports_builtin_tools()
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    struct TestProvider;

    #[async_trait]
    impl Provider for TestProvider {
        fn get_name(&self) -> &str {
            "test"
        }

        async fn stream(
            &self,
            _model_config: &ModelConfig,
            _system: &str,
            _messages: &[Message],
            _tools: &[rmcp::model::Tool],
        ) -> Result<MessageStream, ProviderError> {
            unreachable!()
        }
    }

    #[test]
    fn builtin_tool_support_is_delegated_to_inner_provider() {
        let inner: Arc<dyn Provider> = Arc::new(TestProvider);
        let provider = CompactingProvider::new(inner);

        assert!(provider.manages_own_context());
        assert!(provider.supports_builtin_tools());
    }
}