agent_knowledge_base/

lib.rs

/*!
agent-knowledge-base: keyword-indexed knowledge base for LLM agents.

```rust
use agent_knowledge_base::KnowledgeBase;

let mut kb = KnowledgeBase::new();
kb.add("rust_overview", "Rust is a systems language focused on safety and performance.");
kb.add("python_overview", "Python is a high-level interpreted language.");
let results = kb.search("rust safety");
assert!(!results.is_empty());
assert!(results[0].1.contains("Rust"));
```
*/

use std::collections::HashMap;

/// A knowledge base entry.
#[derive(Debug, Clone)]
pub struct Entry {
    pub id: String,
    pub text: String,
    pub tags: Vec<String>,
}

/// Simple keyword-scored knowledge base.
pub struct KnowledgeBase {
    entries: HashMap<String, Entry>,
}

impl KnowledgeBase {
    pub fn new() -> Self { Self { entries: HashMap::new() } }

    /// Add or replace an entry.
    pub fn add(&mut self, id: &str, text: &str) {
        self.entries.insert(id.to_string(), Entry { id: id.to_string(), text: text.to_string(), tags: Vec::new() });
    }

    /// Add with tags.
    pub fn add_tagged(&mut self, id: &str, text: &str, tags: &[&str]) {
        self.entries.insert(id.to_string(), Entry {
            id: id.to_string(),
            text: text.to_string(),
            tags: tags.iter().map(|t| t.to_string()).collect(),
        });
    }

    pub fn get(&self, id: &str) -> Option<&Entry> { self.entries.get(id) }

    pub fn remove(&mut self, id: &str) -> Option<Entry> { self.entries.remove(id) }

    pub fn len(&self) -> usize { self.entries.len() }
    pub fn is_empty(&self) -> bool { self.entries.is_empty() }

    /// Search by keywords (space-separated). Returns (id, text, score) tuples sorted by score desc.
    pub fn search(&self, query: &str) -> Vec<(&str, &str, f64)> {
        let keywords: Vec<String> = query.split_whitespace().map(|w| w.to_lowercase()).collect();
        if keywords.is_empty() { return Vec::new(); }

        let mut scored: Vec<(&str, &str, f64)> = self.entries.values()
            .filter_map(|e| {
                let text_lower = e.text.to_lowercase();
                let matched = keywords.iter().filter(|k| text_lower.contains(k.as_str())).count();
                if matched > 0 {
                    Some((e.id.as_str(), e.text.as_str(), matched as f64 / keywords.len() as f64))
                } else {
                    None
                }
            })
            .collect();

        scored.sort_by(|a, b| b.2.partial_cmp(&a.2).unwrap().then(a.0.cmp(b.0)));
        scored
    }

    /// Find entries by tag.
    pub fn by_tag(&self, tag: &str) -> Vec<&Entry> {
        let mut v: Vec<&Entry> = self.entries.values().filter(|e| e.tags.iter().any(|t| t == tag)).collect();
        v.sort_by_key(|e| &e.id);
        v
    }

    /// All entry ids (sorted).
    pub fn ids(&self) -> Vec<&str> {
        let mut v: Vec<&str> = self.entries.keys().map(|s| s.as_str()).collect();
        v.sort();
        v
    }
}

impl Default for KnowledgeBase {
    fn default() -> Self { Self::new() }
}

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

    #[test]
    fn add_and_get() {
        let mut kb = KnowledgeBase::new();
        kb.add("doc1", "Rust is safe");
        assert_eq!(kb.get("doc1").unwrap().text, "Rust is safe");
    }

    #[test]
    fn search_finds_match() {
        let mut kb = KnowledgeBase::new();
        kb.add("doc1", "Rust is a systems programming language");
        let r = kb.search("rust");
        assert!(!r.is_empty());
    }

    #[test]
    fn search_returns_empty_for_no_match() {
        let mut kb = KnowledgeBase::new();
        kb.add("doc1", "Rust programming");
        let r = kb.search("python");
        assert!(r.is_empty());
    }

    #[test]
    fn search_scores_multi_keyword() {
        let mut kb = KnowledgeBase::new();
        kb.add("doc1", "rust safety performance");
        kb.add("doc2", "rust language");
        let r = kb.search("rust safety");
        // doc1 matches both keywords, doc2 matches one
        assert_eq!(r[0].0, "doc1");
    }

    #[test]
    fn case_insensitive_search() {
        let mut kb = KnowledgeBase::new();
        kb.add("doc1", "Rust Is Great");
        let r = kb.search("rust");
        assert!(!r.is_empty());
    }

    #[test]
    fn add_tagged_and_by_tag() {
        let mut kb = KnowledgeBase::new();
        kb.add_tagged("doc1", "content", &["lang", "systems"]);
        kb.add_tagged("doc2", "content", &["scripting"]);
        assert_eq!(kb.by_tag("lang").len(), 1);
        assert_eq!(kb.by_tag("scripting").len(), 1);
    }

    #[test]
    fn remove() {
        let mut kb = KnowledgeBase::new();
        kb.add("x", "text");
        assert!(kb.remove("x").is_some());
        assert!(kb.get("x").is_none());
    }

    #[test]
    fn len_and_empty() {
        let mut kb = KnowledgeBase::new();
        assert!(kb.is_empty());
        kb.add("x", "y");
        assert_eq!(kb.len(), 1);
    }

    #[test]
    fn ids_sorted() {
        let mut kb = KnowledgeBase::new();
        kb.add("c", "c");
        kb.add("a", "a");
        kb.add("b", "b");
        assert_eq!(kb.ids(), vec!["a", "b", "c"]);
    }

    #[test]
    fn search_empty_query() {
        let mut kb = KnowledgeBase::new();
        kb.add("x", "text");
        assert!(kb.search("").is_empty());
    }

    #[test]
    fn add_replaces_existing() {
        let mut kb = KnowledgeBase::new();
        kb.add("x", "old");
        kb.add("x", "new");
        assert_eq!(kb.get("x").unwrap().text, "new");
        assert_eq!(kb.len(), 1);
    }

    #[test]
    fn search_score_is_fraction() {
        let mut kb = KnowledgeBase::new();
        kb.add("doc", "rust safety memory");
        let r = kb.search("rust memory");
        assert!((r[0].2 - 1.0).abs() < 1e-9); // both keywords match
    }
}