cascade_agent/tools/
knowledge_tool.rs

1//! Knowledge query tool — queries a local vector DB for relevant knowledge.
2
3use std::sync::Arc;
4
5use async_trait::async_trait;
6use serde_json::{json, Value};
7
8use super::{Tool, ToolResult};
9
10// ---------------------------------------------------------------------------
11// KnowledgeProvider trait
12// ---------------------------------------------------------------------------
13// The knowledge module hasn't been built yet, so we define a minimal trait
14// that the real KnowledgeBase will implement.  This keeps this tool decoupled
15// from the (future) knowledge implementation.
16
17/// A single knowledge result returned from a vector search.
18#[derive(Debug, Clone, serde::Serialize, serde::Deserialize)]
19pub struct KnowledgeHit {
20    pub text: String,
21    pub score: f32,
22    #[serde(skip_serializing_if = "Option::is_none")]
23    pub metadata: Option<serde_json::Map<String, Value>>,
24}
25
26/// Trait that any knowledge backend must implement.
27#[async_trait]
28pub trait KnowledgeProvider: Send + Sync {
29    /// Search the knowledge base for chunks similar to `query`.
30    async fn query(
31        &self,
32        query: &str,
33        collection: &str,
34        limit: usize,
35    ) -> Result<Vec<KnowledgeHit>, String>;
36}
37
38// ---------------------------------------------------------------------------
39// KnowledgeQueryTool
40// ---------------------------------------------------------------------------
41
42/// Queries the local vector DB for relevant knowledge.
43///
44/// Parameters:
45///   - `query`      (string, required): the query text.
46///   - `collection` (string, optional): the collection to search (default "general").
47///   - `limit`      (int, optional):    max results (default 5).
48pub struct KnowledgeQueryTool {
49    /// The knowledge backend. `None` when no knowledge base is configured.
50    provider: Option<Arc<dyn KnowledgeProvider>>,
51    default_collection: String,
52    default_limit: usize,
53}
54
55impl KnowledgeQueryTool {
56    /// Create with a knowledge provider backend.
57    pub fn new(provider: Arc<dyn KnowledgeProvider>) -> Self {
58        Self {
59            provider: Some(provider),
60            default_collection: "general".to_owned(),
61            default_limit: 5,
62        }
63    }
64
65    /// Create with a knowledge provider and explicit defaults.
66    pub fn with_defaults(
67        provider: Arc<dyn KnowledgeProvider>,
68        default_collection: String,
69        default_limit: usize,
70    ) -> Self {
71        Self {
72            provider: Some(provider),
73            default_collection,
74            default_limit,
75        }
76    }
77
78    /// Create a *disabled* knowledge tool (no backend configured).
79    pub fn disabled() -> Self {
80        Self {
81            provider: None,
82            default_collection: "general".to_owned(),
83            default_limit: 5,
84        }
85    }
86}
87
88#[async_trait]
89impl Tool for KnowledgeQueryTool {
90    fn name(&self) -> &str {
91        "knowledge_query"
92    }
93
94    fn description(&self) -> &str {
95        "Query the local knowledge base for relevant text snippets. Returns passages ranked by semantic similarity."
96    }
97
98    fn parameters_schema(&self) -> Value {
99        json!({
100            "type": "object",
101            "properties": {
102                "query": {
103                    "type": "string",
104                    "description": "The query text to search for."
105                },
106                "collection": {
107                    "type": "string",
108                    "description": "The knowledge collection to search.",
109                    "default": self.default_collection
110                },
111                "limit": {
112                    "type": "integer",
113                    "description": "Maximum number of results to return.",
114                    "default": self.default_limit
115                }
116            },
117            "required": ["query"]
118        })
119    }
120
121    async fn execute(&self, args: Value) -> ToolResult {
122        let query = match args.get("query").and_then(|v| v.as_str()) {
123            Some(q) => q,
124            None => return ToolResult::err("Missing required parameter 'query'"),
125        };
126
127        let collection = args
128            .get("collection")
129            .and_then(|v| v.as_str())
130            .unwrap_or(&self.default_collection);
131
132        let limit = args
133            .get("limit")
134            .and_then(|v| v.as_u64())
135            .unwrap_or(self.default_limit as u64) as usize;
136
137        let provider = match &self.provider {
138            Some(p) => p,
139            None => {
140                return ToolResult::err(
141                    "Knowledge base is not configured. Set up the knowledge module first.",
142                );
143            }
144        };
145
146        match provider.query(query, collection, limit).await {
147            Ok(hits) => {
148                let results: Vec<Value> = hits
149                    .iter()
150                    .map(|h| {
151                        let mut obj = json!({
152                            "text": h.text,
153                            "score": h.score,
154                        });
155                        if let Some(ref meta) = h.metadata {
156                            obj.as_object_mut()
157                                .unwrap()
158                                .insert("metadata".into(), Value::Object(meta.clone()));
159                        }
160                        obj
161                    })
162                    .collect();
163
164                ToolResult::ok(json!({
165                    "query": query,
166                    "collection": collection,
167                    "count": results.len(),
168                    "results": results,
169                }))
170            }
171            Err(e) => ToolResult::err(format!("Knowledge query failed: {}", e)),
172        }
173    }
174}
175
176// ---------------------------------------------------------------------------
177// No-op knowledge provider for testing
178// ---------------------------------------------------------------------------
179
180#[cfg(test)]
181struct MockKnowledgeProvider;
182
183#[cfg(test)]
184#[async_trait]
185impl KnowledgeProvider for MockKnowledgeProvider {
186    async fn query(
187        &self,
188        query: &str,
189        _collection: &str,
190        limit: usize,
191    ) -> Result<Vec<KnowledgeHit>, String> {
192        Ok(vec![KnowledgeHit {
193            text: format!("Mock result for: {}", query),
194            score: 0.99,
195            metadata: None,
196        }]
197        .into_iter()
198        .cycle()
199        .take(limit)
200        .collect())
201    }
202}
203
204#[cfg(test)]
205mod tests {
206    use super::*;
207
208    #[tokio::test]
209    async fn knowledge_query_works() {
210        let provider = Arc::new(MockKnowledgeProvider);
211        let tool = KnowledgeQueryTool::new(provider);
212        let result = tool.execute(json!({"query": "rust async"})).await;
213        assert_eq!(result.status, super::super::ToolStatus::Success);
214        let data = result.data.as_object().unwrap();
215        assert_eq!(data["count"].as_u64().unwrap(), 5);
216        assert!(data["results"].as_array().unwrap()[0]["text"]
217            .as_str()
218            .unwrap()
219            .contains("Mock result for: rust async"));
220    }
221
222    #[tokio::test]
223    async fn knowledge_query_disabled() {
224        let tool = KnowledgeQueryTool::disabled();
225        let result = tool.execute(json!({"query": "test"})).await;
226        assert_eq!(result.status, super::super::ToolStatus::Error);
227        assert!(result.error.unwrap().contains("not configured"));
228    }
229
230    #[tokio::test]
231    async fn knowledge_query_custom_params() {
232        let provider = Arc::new(MockKnowledgeProvider);
233        let tool = KnowledgeQueryTool::with_defaults(provider, "custom".to_owned(), 3);
234        let result = tool
235            .execute(json!({
236                "query": "test",
237                "collection": "special",
238                "limit": 2
239            }))
240            .await;
241        assert_eq!(result.status, super::super::ToolStatus::Success);
242        assert_eq!(result.data["count"].as_u64().unwrap(), 2);
243    }
244}
cascade_agent/tools/knowledge_tool.rs

cascade_agent/tools/
knowledge_tool.rs
