cel_brief/sources/
tool_catalog.rs

1//! [`ToolCatalogSource`] — static [`ToolSchema`] list at [`Priority::High`].
2//!
3//! Owns a `Vec<ToolSchema>` and emits each as a [`Contribution::tool`] every
4//! turn. High priority because tool schemas usually steer model behaviour more
5//! than any single memory or history item; the builder is free to drop tools
6//! under budget pressure, but they outrank `Normal` / `Low` content.
7
8use async_trait::async_trait;
9
10use crate::source::{Contribution, ContributionContent, Source, SourceError};
11use crate::types::{BriefContext, Priority, SourceId, ToolSchema};
12
13/// A [`Source`] that exposes a fixed catalog of [`ToolSchema`]s to the model.
14///
15/// Each contributed schema inherits this source's [`SourceId`] via the
16/// builder's admission step (the input schemas keep whatever `source` they
17/// were constructed with — see [`Self::new`] for details).
18#[derive(Debug, Clone)]
19pub struct ToolCatalogSource {
20    id: SourceId,
21    tools: Vec<ToolSchema>,
22}
23
24impl ToolCatalogSource {
25    /// Construct with the default ID `"tool_catalog"`. The input schemas'
26    /// `source` field is rewritten to this source's ID so attribution in the
27    /// final brief is correct.
28    pub fn new(tools: impl IntoIterator<Item = ToolSchema>) -> Self {
29        let id = SourceId::new("tool_catalog");
30        let tools = tools
31            .into_iter()
32            .map(|mut t| {
33                t.source = id.clone();
34                t
35            })
36            .collect();
37        ToolCatalogSource { id, tools }
38    }
39
40    /// Override the default [`SourceId`]. The source rewrites every contained
41    /// schema's `source` field to match.
42    pub fn with_id(mut self, id: impl Into<SourceId>) -> Self {
43        self.id = id.into();
44        for tool in &mut self.tools {
45            tool.source = self.id.clone();
46        }
47        self
48    }
49
50    /// Read-only access to the contained schemas.
51    pub fn tools(&self) -> &[ToolSchema] {
52        &self.tools
53    }
54
55    /// Default per-tool token estimate (`description.len() / 4`).
56    fn estimate_tokens(tool: &ToolSchema) -> usize {
57        // Cheap heuristic: description + serialised schema length / 4.
58        let schema_len = serde_json::to_string(&tool.input_schema)
59            .map(|s| s.len())
60            .unwrap_or(0);
61        (tool.description.len() + schema_len + tool.name.len()).div_ceil(4)
62    }
63}
64
65#[async_trait]
66impl Source for ToolCatalogSource {
67    fn id(&self) -> SourceId {
68        self.id.clone()
69    }
70
71    fn priority(&self) -> Priority {
72        Priority::High
73    }
74
75    async fn contribute(&self, _ctx: &BriefContext) -> Result<Vec<Contribution>, SourceError> {
76        if self.tools.is_empty() {
77            return Err(SourceError::Skipped("empty tool catalog".into()));
78        }
79        Ok(self
80            .tools
81            .iter()
82            .map(|tool| {
83                let est = Self::estimate_tokens(tool);
84                Contribution {
85                    content: ContributionContent::Tool {
86                        schema: tool.clone(),
87                    },
88                    estimated_tokens: est,
89                    importance: 0.9,
90                    redactable: false,
91                    tags: vec!["tool".into()],
92                }
93            })
94            .collect())
95    }
96}
97
98#[cfg(test)]
99mod tests {
100    use super::*;
101    use crate::types::TokenBudget;
102    use serde_json::json;
103
104    fn schema(name: &str) -> ToolSchema {
105        ToolSchema {
106            name: name.into(),
107            description: format!("Test tool {name}"),
108            input_schema: json!({"type":"object","properties":{}}),
109            source: SourceId::new("__unset__"),
110        }
111    }
112
113    #[tokio::test]
114    async fn emits_one_contribution_per_tool() {
115        let src = ToolCatalogSource::new(vec![schema("alpha"), schema("beta")]);
116        let ctx = BriefContext::new(TokenBudget::default());
117        let contributions = src.contribute(&ctx).await.expect("ok");
118        assert_eq!(contributions.len(), 2);
119
120        for c in &contributions {
121            assert_eq!(c.importance, 0.9);
122            assert!(!c.redactable);
123            match &c.content {
124                ContributionContent::Tool { schema } => {
125                    assert_eq!(schema.source, SourceId::new("tool_catalog"));
126                }
127                other => panic!("expected Tool, got {other:?}"),
128            }
129        }
130    }
131
132    #[tokio::test]
133    async fn empty_catalog_is_skipped() {
134        let src = ToolCatalogSource::new(Vec::<ToolSchema>::new());
135        let ctx = BriefContext::new(TokenBudget::default());
136        match src.contribute(&ctx).await {
137            Err(SourceError::Skipped(_)) => {}
138            other => panic!("expected Skipped, got {other:?}"),
139        }
140    }
141
142    #[tokio::test]
143    async fn with_id_rewrites_tool_source() {
144        let src = ToolCatalogSource::new(vec![schema("alpha")]).with_id("example_tools");
145        assert_eq!(src.id(), SourceId::new("example_tools"));
146        let ctx = BriefContext::new(TokenBudget::default());
147        let contributions = src.contribute(&ctx).await.expect("ok");
148        match &contributions[0].content {
149            ContributionContent::Tool { schema } => {
150                assert_eq!(schema.source, SourceId::new("example_tools"));
151            }
152            other => panic!("expected Tool, got {other:?}"),
153        }
154    }
155
156    #[tokio::test]
157    async fn priority_is_high() {
158        let src = ToolCatalogSource::new(vec![schema("alpha")]);
159        assert_eq!(src.priority(), Priority::High);
160    }
161}
cel_brief/sources/tool_catalog.rs

cel_brief/sources/
tool_catalog.rs
