cascade_agent/knowledge/
mod.rs

1//! Knowledge base facade — embeds, stores, and retrieves knowledge.
2
3pub mod embeddings;
4pub mod vectordb;
5
6use std::sync::Arc;
7
8use async_trait::async_trait;
9use serde_json::{Map, Value};
10
11use crate::config::KnowledgeSettings;
12use crate::error::Result;
13use crate::tools::knowledge_tool::{KnowledgeHit, KnowledgeProvider};
14
15use embeddings::Embedder;
16use vectordb::{KnowledgeEntry, SearchResult, VectorStore};
17
18// ---------------------------------------------------------------------------
19// KnowledgeBase
20// ---------------------------------------------------------------------------
21
22/// High-level knowledge base that combines embedding and vector storage.
23#[derive(Debug)]
24pub struct KnowledgeBase {
25    store: VectorStore,
26    config: KnowledgeSettings,
27}
28
29impl KnowledgeBase {
30    /// Create a new knowledge base from configuration.
31    ///
32    /// Downloads the embedding model (if needed) and connects to LanceDB.
33    pub async fn new(config: &KnowledgeSettings) -> Result<Self> {
34        // Create embedder (blocking model download + load).
35        let model_name = config.embedding_model.clone();
36        let embedder = tokio::task::spawn_blocking(move || Embedder::new(&model_name))
37            .await
38            .map_err(|e| {
39                crate::error::AgentError::KnowledgeError(format!("Embedder init panicked: {}", e))
40            })??;
41
42        let embedder = Arc::new(embedder);
43        let store = VectorStore::new(&config.db_path, embedder, &config.default_collection).await?;
44
45        // Ensure the default collection exists.
46        store.create_collection(&config.default_collection).await?;
47
48        Ok(Self {
49            store,
50            config: config.clone(),
51        })
52    }
53
54    /// Query the default collection.
55    pub async fn query_existing(&self, query: &str) -> Result<Vec<SearchResult>> {
56        let limit = self.config.max_results;
57        self.store
58            .search(&self.config.default_collection, query, limit)
59            .await
60    }
61
62    /// Store entries into a collection.
63    pub async fn store_results(
64        &self,
65        collection: &str,
66        entries: Vec<KnowledgeEntry>,
67    ) -> Result<()> {
68        self.store.insert(collection, entries).await
69    }
70
71    /// Create a new collection.
72    pub async fn create_collection(&self, name: &str) -> Result<()> {
73        self.store.create_collection(name).await
74    }
75
76    /// List all collections.
77    pub async fn list_collections(&self) -> Result<Vec<String>> {
78        self.store.list_collections().await
79    }
80}
81
82// ---------------------------------------------------------------------------
83// KnowledgeProvider implementation
84// ---------------------------------------------------------------------------
85
86#[async_trait]
87impl KnowledgeProvider for KnowledgeBase {
88    async fn query(
89        &self,
90        query: &str,
91        collection: &str,
92        limit: usize,
93    ) -> std::result::Result<Vec<KnowledgeHit>, String> {
94        let results = self
95            .store
96            .search(collection, query, limit)
97            .await
98            .map_err(|e| format!("Knowledge query failed: {}", e))?;
99
100        let hits: Vec<KnowledgeHit> = results
101            .into_iter()
102            .filter(|r| {
103                // Apply similarity threshold from config.
104                r.score >= self.config.similarity_threshold
105            })
106            .map(|r| KnowledgeHit {
107                text: r.text,
108                score: r.score,
109                metadata: if r.metadata.is_null() {
110                    None
111                } else if let Value::Object(map) = r.metadata {
112                    Some(map)
113                } else {
114                    let mut map = Map::new();
115                    map.insert("value".into(), r.metadata);
116                    Some(map)
117                },
118            })
119            .collect();
120
121        Ok(hits)
122    }
123}
cascade_agent/knowledge/mod.rs

cascade_agent/knowledge/
mod.rs
