use crate::client::KibanaClient;
use crate::etl::Loader;
use crate::{Error, Result};
use serde_json::Value;
use tokio::task::JoinSet;
pub struct ToolsLoader {
client: KibanaClient,
}
impl ToolsLoader {
pub fn new(client: KibanaClient) -> Self {
Self { client }
}
}
impl Loader for ToolsLoader {
type Item = Value;
async fn load(&self, items: Vec<Self::Item>) -> Result<usize> {
let mut count = 0;
let mut set = JoinSet::new();
for tool in items {
let client = self.client.clone();
set.spawn(async move {
let tool_id = tool
.get("id")
.and_then(|v| v.as_str())
.ok_or(Error::MissingResourceId { resource: "tool" })?;
if tool
.get("readonly")
.and_then(|v| v.as_bool())
.unwrap_or(false)
{
tracing::debug!("Skipping readonly tool: {}", tool_id);
return Ok::<bool, Error>(false);
}
let path = format!("api/agent_builder/tools/{}", tool_id);
let exists = match client.head(&path).await?.status().as_u16() {
200 => true,
404 => false,
status => {
return Err(Error::message(format!(
"Failed to check tool existence ({tool_id}): {status}"
)));
}
};
if exists {
let mut tool_body = tool.clone();
if let Some(obj) = tool_body.as_object_mut() {
obj.remove("id");
obj.remove("readonly");
obj.remove("schema");
obj.remove("type");
}
let path = format!("api/agent_builder/tools/{}", tool_id);
let response = client.put_json_value(&path, &tool_body).await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(Error::api_response(status, body));
}
tracing::info!("Updated tool: {}", tool_id);
} else {
let mut tool_body = tool.clone();
if let Some(obj) = tool_body.as_object_mut() {
obj.remove("readonly");
obj.remove("schema");
}
let path = "api/agent_builder/tools";
let response = client.post_json_value(path, &tool_body).await?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(Error::api_response(status, body));
}
tracing::info!("Created tool: {}", tool_id);
}
Ok::<bool, Error>(true)
});
}
while let Some(res) = set.join_next().await {
match res {
Ok(Ok(loaded)) => {
if loaded {
count += 1;
}
}
Ok(Err(e)) => tracing::error!("Failed to load tool: {}", e),
Err(e) => tracing::error!("Task panicked: {}", e),
}
}
Ok(count)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::{Auth, KibanaClient};
use serde_json::json;
use url::Url;
#[test]
fn test_loader_creation() {
let url = Url::parse("http://localhost:5601").unwrap();
let client = KibanaClient::new(url, Auth::None).unwrap();
let space_client = client.space("default").unwrap();
let _loader = ToolsLoader::new(space_client);
}
#[tokio::test]
async fn test_missing_id_fails() {
let url = Url::parse("http://localhost:5601").unwrap();
let client = KibanaClient::new(url, Auth::None).unwrap();
let space_client = client.space("default").unwrap();
let loader = ToolsLoader::new(space_client);
let tool = json!({"description": "No ID"});
let result = loader.load(vec![tool]).await;
assert!(result.is_ok());
assert_eq!(result.unwrap(), 0);
}
}