use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::{Value, json};
use super::{Capability, CapabilityLocalization, CapabilityStatus, RiskLevel};
use crate::tools::{Tool, ToolExecutionResult};
use crate::traits::ToolContext;
pub const KNOWLEDGE_BASE_CAPABILITY_ID: &str = "knowledge_base";
const ENTRY_KINDS: &[&str] = &["note", "table", "business", "query", "runbook"];
#[derive(Debug, Clone, Serialize, Deserialize, Default)]
pub struct KnowledgeBaseConfig {
#[serde(default)]
pub bases: Vec<String>,
#[serde(default)]
pub kinds: Vec<String>,
}
pub fn validate_knowledge_base_config(cfg: &KnowledgeBaseConfig) -> Result<(), String> {
for base in &cfg.bases {
if !is_valid_kb_id(base) {
return Err(format!(
"knowledge_base bases[*] must be a kb_<32-hex> id, got '{base}'"
));
}
}
let mut seen = std::collections::HashSet::new();
for base in &cfg.bases {
if !seen.insert(base) {
return Err(format!(
"knowledge_base bases[*] contains duplicate '{base}'"
));
}
}
for kind in &cfg.kinds {
if !ENTRY_KINDS.contains(&kind.as_str()) {
return Err(format!(
"knowledge_base kinds[*] must be one of {:?}, got '{kind}'",
ENTRY_KINDS
));
}
}
Ok(())
}
fn is_valid_kb_id(s: &str) -> bool {
s.len() == 35
&& s.starts_with("kb_")
&& s[3..]
.chars()
.all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
}
pub struct SearchKnowledgeTool {
config: KnowledgeBaseConfig,
}
impl SearchKnowledgeTool {
pub fn new(config: KnowledgeBaseConfig) -> Self {
Self { config }
}
}
#[async_trait]
impl Tool for SearchKnowledgeTool {
fn name(&self) -> &str {
"search_knowledge"
}
fn description(&self) -> &str {
"Search curated organization Knowledge Bases (table docs, business rules, \
validated query templates, runbooks) by keyword. Consult this before \
answering data questions and cite results by their kbe_ id."
}
fn parameters_schema(&self) -> Value {
json!({
"type": "object",
"properties": {
"query": { "type": "string", "description": "Keyword search across entry title and body" },
"kind": {
"type": "string",
"enum": ["note", "table", "business", "query", "runbook"],
"description": "Optional filter by entry kind"
},
"tags": { "type": "array", "items": { "type": "string" }, "description": "Optional tag filter" },
"limit": { "type": "integer", "minimum": 1, "maximum": 25, "default": 10 }
},
"required": ["query"],
"additionalProperties": false
})
}
async fn execute(&self, _arguments: Value) -> ToolExecutionResult {
ToolExecutionResult::tool_error("search_knowledge requires execution context")
}
async fn execute_with_context(
&self,
arguments: Value,
context: &ToolContext,
) -> ToolExecutionResult {
let query = match arguments.get("query").and_then(|v| v.as_str()) {
Some(q) if !q.trim().is_empty() => q.trim().to_string(),
_ => return ToolExecutionResult::tool_error("Missing required parameter: query"),
};
if self.config.bases.is_empty() {
return ToolExecutionResult::success(json!({ "count": 0, "results": [] }));
}
let Some(store) = context.knowledge_store.as_ref() else {
return ToolExecutionResult::tool_error(
"Knowledge search is not available in this execution context",
);
};
let Some(org_id) = context.org_id else {
return ToolExecutionResult::tool_error(
"Knowledge search requires an organization context",
);
};
let kind = arguments
.get("kind")
.and_then(|v| v.as_str())
.map(|s| s.to_string())
.or_else(|| match self.config.kinds.as_slice() {
[single] => Some(single.clone()),
_ => None,
});
let tags: Vec<String> = arguments
.get("tags")
.and_then(|v| v.as_array())
.map(|a| {
a.iter()
.filter_map(|v| v.as_str().map(|s| s.to_lowercase()))
.collect()
})
.unwrap_or_default();
let limit = arguments
.get("limit")
.and_then(|v| v.as_u64())
.map(|v| (v as usize).clamp(1, 25))
.unwrap_or(10);
match store
.search_knowledge(
org_id,
&self.config.bases,
&query,
kind.as_deref(),
&tags,
limit,
)
.await
{
Ok(hits) => ToolExecutionResult::success(json!({
"count": hits.len(),
"results": hits,
})),
Err(e) => {
ToolExecutionResult::internal_error_msg(format!("knowledge search failed: {e}"))
}
}
}
fn requires_context(&self) -> bool {
true
}
fn deferrable_policy(&self) -> crate::tool_types::DeferrablePolicy {
crate::tool_types::DeferrablePolicy::Never
}
}
pub struct KnowledgeBaseCapability;
impl Capability for KnowledgeBaseCapability {
fn id(&self) -> &str {
KNOWLEDGE_BASE_CAPABILITY_ID
}
fn name(&self) -> &str {
"Knowledge Base"
}
fn description(&self) -> &str {
"Bind an agent to curated org Knowledge Bases and give it the \
`search_knowledge` tool to ground answers in human-edited table docs, \
business rules, validated SQL templates, and runbooks. \
See `specs/knowledge-bases.md`."
}
fn status(&self) -> CapabilityStatus {
CapabilityStatus::Available
}
fn icon(&self) -> Option<&str> {
Some("library")
}
fn category(&self) -> Option<&str> {
Some("Knowledge")
}
fn features(&self) -> Vec<&'static str> {
vec!["knowledge"]
}
fn risk_level(&self) -> RiskLevel {
RiskLevel::Low
}
fn system_prompt_addition(&self) -> Option<&str> {
Some(
"You can search curated organization knowledge with the `search_knowledge` tool \
(table docs, business rules, validated queries, runbooks). Consult it before \
answering data questions, and cite the entries you use by their kbe_ id.",
)
}
fn tools(&self) -> Vec<Box<dyn Tool>> {
self.tools_with_config(&Value::Null)
}
fn tools_with_config(&self, config: &Value) -> Vec<Box<dyn Tool>> {
let cfg: KnowledgeBaseConfig = serde_json::from_value(config.clone()).unwrap_or_default();
vec![Box::new(SearchKnowledgeTool::new(cfg))]
}
fn config_schema(&self) -> Option<Value> {
Some(json!({
"type": "object",
"properties": {
"bases": {
"type": "array",
"title": "Knowledge Bases",
"description": "Knowledge Base IDs the agent can search.",
"items": {
"type": "string",
"title": "Knowledge Base ID",
"description": "Knowledge Base ID (kb_<32-hex>).",
"pattern": "^kb_[0-9a-f]{32}$"
}
},
"kinds": {
"type": "array",
"title": "Default kind filter",
"description": "Optional kind filter applied when the agent does not pass `kind`.",
"items": {
"type": "string",
"title": "Entry kind",
"description": "Knowledge Base entry kind to include.",
"oneOf": [
{ "const": "note", "title": "Note" },
{ "const": "table", "title": "Table doc" },
{ "const": "business", "title": "Business rule" },
{ "const": "query", "title": "Query template" },
{ "const": "runbook", "title": "Runbook" }
]
}
}
}
}))
}
fn localizations(&self) -> Vec<CapabilityLocalization> {
vec![
CapabilityLocalization {
locale: "en",
name: None,
description: None,
config_description: Some(
"Selects which Knowledge Bases the agent can search and an optional \
default entry-kind filter.",
),
config_overlay: None,
},
CapabilityLocalization {
locale: "uk",
name: Some("База знань"),
description: Some(
"Прив'язує агента до курованих Баз знань організації, щоб відповіді \
спиралися на редаговані людьми описи таблиць, бізнес-правила, \
перевірені SQL-шаблони та runbook-и.",
),
config_description: Some(
"Визначає, у яких Базах знань агент може шукати, та необов'язковий \
типовий фільтр за видом записів.",
),
config_overlay: Some(json!({
"properties": {
"bases": {
"title": "Бази знань",
"description": "Ідентифікатори Баз знань, у яких агент може шукати.",
"items": {
"title": "Ідентифікатор Бази знань",
"description": "Ідентифікатор Бази знань (kb_<32-hex>)."
}
},
"kinds": {
"title": "Типовий фільтр виду",
"description": "Необов'язковий фільтр за видом записів, що застосовується, коли агент не передає kind явно.",
"items": {
"title": "Вид запису",
"description": "Вид записів Бази знань, який потрібно включити.",
"enum_labels": {
"note": "Нотатка",
"table": "Опис таблиці",
"business": "Бізнес-правило",
"query": "Шаблон запиту",
"runbook": "Runbook"
}
}
}
}
})),
},
]
}
fn validate_config(&self, config: &Value) -> Result<(), String> {
if config.is_null() {
return Ok(());
}
let typed: KnowledgeBaseConfig = serde_json::from_value(config.clone())
.map_err(|e| format!("invalid knowledge_base config: {e}"))?;
validate_knowledge_base_config(&typed)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn id_and_name() {
let cap = KnowledgeBaseCapability;
assert_eq!(cap.id(), "knowledge_base");
assert_eq!(cap.name(), "Knowledge Base");
}
#[test]
fn validate_accepts_empty_config() {
let cap = KnowledgeBaseCapability;
assert!(cap.validate_config(&json!({})).is_ok());
assert!(cap.validate_config(&json!({ "bases": [] })).is_ok());
assert!(cap.validate_config(&Value::Null).is_ok());
}
#[test]
fn validate_accepts_well_formed_bases_and_kinds() {
let cap = KnowledgeBaseCapability;
let cfg = json!({
"bases": ["kb_00000000000000000000000000000001"],
"kinds": ["table", "business"]
});
assert!(cap.validate_config(&cfg).is_ok());
}
#[test]
fn validate_rejects_malformed_kb_id() {
let cap = KnowledgeBaseCapability;
let cfg = json!({ "bases": ["mem_00000000000000000000000000000001"] });
let err = cap.validate_config(&cfg).unwrap_err();
assert!(err.contains("kb_"));
}
#[test]
fn validate_rejects_duplicate_bases() {
let cap = KnowledgeBaseCapability;
let cfg = json!({
"bases": [
"kb_00000000000000000000000000000001",
"kb_00000000000000000000000000000001"
]
});
let err = cap.validate_config(&cfg).unwrap_err();
assert!(err.contains("duplicate"));
}
#[test]
fn validate_rejects_unknown_kind() {
let cap = KnowledgeBaseCapability;
let cfg = json!({ "kinds": ["nope"] });
let err = cap.validate_config(&cfg).unwrap_err();
assert!(err.contains("kinds"));
}
#[test]
fn uk_localization_and_schema_one_of_match_validation() {
let cap = KnowledgeBaseCapability;
assert_eq!(cap.localized_name(Some("uk-UA")), "База знань");
assert!(
cap.localized_description(Some("uk-UA"))
.contains("Баз знань")
);
assert!(cap.describe_schema(Some("uk")).is_some());
assert!(cap.describe_schema(None).is_some());
let schema = cap.config_schema().expect("config schema");
let consts: Vec<&str> = schema["properties"]["kinds"]["items"]["oneOf"]
.as_array()
.expect("oneOf")
.iter()
.map(|v| v["const"].as_str().expect("const"))
.collect();
assert_eq!(consts, ENTRY_KINDS);
for kind in consts {
assert!(cap.validate_config(&json!({ "kinds": [kind] })).is_ok());
}
}
use crate::traits::{KnowledgeSearchHit, KnowledgeStore, ToolContext};
use crate::typed_id::{DEFAULT_ORG_ID, SessionId};
use std::sync::Arc;
struct MockKnowledgeStore {
hits: Vec<KnowledgeSearchHit>,
}
#[async_trait]
impl KnowledgeStore for MockKnowledgeStore {
async fn search_knowledge(
&self,
_org_id: crate::typed_id::OrgId,
kb_public_ids: &[String],
_query: &str,
_kind: Option<&str>,
_tags: &[String],
_limit: usize,
) -> crate::error::Result<Vec<KnowledgeSearchHit>> {
if kb_public_ids.is_empty() {
Ok(Vec::new())
} else {
Ok(self.hits.clone())
}
}
}
fn hit() -> KnowledgeSearchHit {
KnowledgeSearchHit {
id: "kbe_00000000000000000000000000000001".into(),
kb_id: "kb_00000000000000000000000000000001".into(),
title: "Orders".into(),
kind: "table".into(),
tags: vec!["sales".into()],
snippet: "One row per order.".into(),
resource: None,
}
}
#[tokio::test]
async fn search_tool_returns_results_from_store() {
let tool = SearchKnowledgeTool::new(KnowledgeBaseConfig {
bases: vec!["kb_00000000000000000000000000000001".into()],
kinds: vec![],
});
let mut ctx = ToolContext::new(SessionId::new());
ctx.knowledge_store = Some(Arc::new(MockKnowledgeStore { hits: vec![hit()] }));
ctx.org_id = Some(DEFAULT_ORG_ID);
let result = tool
.execute_with_context(json!({ "query": "orders" }), &ctx)
.await;
match result {
ToolExecutionResult::Success(v) => {
assert_eq!(v["count"], 1);
assert_eq!(
v["results"][0]["id"],
"kbe_00000000000000000000000000000001"
);
}
other => panic!("expected success, got {other:?}"),
}
}
#[tokio::test]
async fn search_tool_with_no_bases_returns_empty() {
let tool = SearchKnowledgeTool::new(KnowledgeBaseConfig::default());
let mut ctx = ToolContext::new(SessionId::new());
ctx.knowledge_store = Some(Arc::new(MockKnowledgeStore { hits: vec![hit()] }));
ctx.org_id = Some(DEFAULT_ORG_ID);
let result = tool
.execute_with_context(json!({ "query": "orders" }), &ctx)
.await;
match result {
ToolExecutionResult::Success(v) => assert_eq!(v["count"], 0),
other => panic!("expected success, got {other:?}"),
}
}
#[tokio::test]
async fn search_tool_requires_query() {
let tool = SearchKnowledgeTool::new(KnowledgeBaseConfig {
bases: vec!["kb_00000000000000000000000000000001".into()],
kinds: vec![],
});
let ctx = ToolContext::new(SessionId::new());
let result = tool.execute_with_context(json!({}), &ctx).await;
assert!(matches!(result, ToolExecutionResult::ToolError(_)));
}
}