use crate::client::KibanaClient;
use crate::etl::Extractor;
use crate::{Error, Result, ResultContext};
use serde_json::Value;
use tokio::task::JoinSet;
pub struct AgentsExtractor {
client: KibanaClient,
manifest: Option<super::AgentsManifest>,
}
impl AgentsExtractor {
pub fn new(client: KibanaClient, manifest: Option<super::AgentsManifest>) -> Self {
Self { client, manifest }
}
pub async fn search_agents(&self, _query: Option<&str>) -> Result<Vec<Value>> {
let path = "api/agent_builder/agents";
tracing::debug!(
"Fetching agents from {} in space '{}'",
path,
self.client.space_id()
);
let response = self
.client
.get(path)
.await
.context("Failed to fetch agents")?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(Error::api_response(status, body));
}
let search_result: Value = response
.json()
.await
.context("Failed to parse agents response")?;
let agents: Vec<Value> = search_result
.get("results")
.and_then(|v| v.as_array())
.map(|arr| arr.to_vec())
.unwrap_or_default();
tracing::info!("Found {} agent(s) via search", agents.len());
Ok(agents)
}
async fn fetch_manifest_agents(&self, manifest: &super::AgentsManifest) -> Result<Vec<Value>> {
let mut agents = Vec::new();
let mut set = JoinSet::new();
for entry in &manifest.agents {
let client = self.client.clone();
let agent_id = entry.id.clone();
let agent_name = entry.name.clone();
set.spawn(async move {
let path = format!("api/agent_builder/agents/{}", agent_id);
tracing::debug!(
"Fetching agent '{}' from space '{}'",
agent_id,
client.space_id()
);
let response = client.get(&path).await.with_context(|| {
format!("Failed to fetch agent '{}' ({})", agent_name, agent_id)
})?;
if !response.status().is_success() {
let status = response.status();
let body = response.text().await.unwrap_or_default();
return Err(Error::api_response(status, body));
}
let agent: Value = response
.json()
.await
.with_context(|| format!("Failed to parse agent '{}' response", agent_id))?;
tracing::debug!("Fetched agent: {}", agent_id);
Ok::<Value, Error>(agent)
});
}
while let Some(res) = set.join_next().await {
match res {
Ok(Ok(agent)) => agents.push(agent),
Ok(Err(e)) => tracing::warn!("{}", e),
Err(e) => tracing::error!("Task panicked: {}", e),
}
}
tracing::info!("Fetched {} agent(s) from manifest", agents.len());
Ok(agents)
}
}
impl Extractor for AgentsExtractor {
type Item = Value;
async fn extract(&self) -> Result<Vec<Self::Item>> {
let agents = if let Some(manifest) = &self.manifest {
self.fetch_manifest_agents(manifest).await?
} else {
tracing::warn!("No manifest provided - use search API to discover agents");
Vec::new()
};
tracing::info!(
"Extracted {} agent(s){}",
agents.len(),
if self.manifest.is_some() {
" (from manifest)"
} else {
""
}
);
Ok(agents)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::client::{Auth, KibanaClient};
use url::Url;
#[test]
fn test_extractor_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 _extractor = AgentsExtractor::new(space_client, None);
}
}