goose_context_management/
model.rs

1use std::sync::Arc;
2
3use async_trait::async_trait;
4use goose_provider_types::base::Provider;
5use goose_provider_types::conversation::message::Message;
6use goose_provider_types::conversation::token_usage::ProviderUsage;
7use goose_provider_types::errors::ProviderError;
8use goose_provider_types::model::ModelConfig;
9
10/// The single completion call compaction needs. Implementations decide model
11/// selection, fallbacks and session plumbing.
12#[async_trait]
13pub trait CompactionModel: Send + Sync {
14    async fn complete(
15        &self,
16        system: &str,
17        messages: &[Message],
18    ) -> Result<(Message, ProviderUsage), ProviderError>;
19}
20
21/// Counts tokens for usage estimation and retained-context reporting.
22#[async_trait]
23pub trait TokenEstimator: Send + Sync {
24    async fn count_chat_tokens(&self, system: &str, messages: &[Message]) -> usize;
25    async fn count_text_tokens(&self, text: &str) -> usize;
26}
27
28pub struct ProviderModel {
29    provider: Arc<dyn Provider>,
30    model_config: ModelConfig,
31}
32
33impl ProviderModel {
34    pub fn new(provider: Arc<dyn Provider>, model_config: ModelConfig) -> Self {
35        Self {
36            provider,
37            model_config,
38        }
39    }
40}
41
42#[async_trait]
43impl CompactionModel for ProviderModel {
44    async fn complete(
45        &self,
46        system: &str,
47        messages: &[Message],
48    ) -> Result<(Message, ProviderUsage), ProviderError> {
49        self.provider
50            .complete(&self.model_config, system, messages, &[])
51            .await
52    }
53}
goose_context_management/model.rs

goose_context_management/
model.rs
