use anyhow::Result;
use async_trait::async_trait;
use reqwest::Client;
use serde_json::Value;
use super::{Tool, ToolOutput};
pub struct TechQueryTool {
client: Client,
endpoint: String,
}
impl TechQueryTool {
pub fn new(endpoint: Option<String>) -> Self {
Self {
client: Client::new(),
endpoint: endpoint.unwrap_or_else(|| "http://localhost:5000/search".to_string()),
}
}
}
#[async_trait]
impl Tool for TechQueryTool {
fn name(&self) -> &str {
"tech_query"
}
fn description(&self) -> &str {
"Search through technical documentation using a semantic search system. \
This connects to an external PyTorch-based documentation search service. \
Use this when you need to look up API references, library documentation, \
or technical specifications."
}
fn parameters_schema(&self) -> Value {
serde_json::json!({
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "The technical query to search for"
},
"top_k": {
"type": "integer",
"description": "Number of results to return (default: 5)"
}
},
"required": ["query"]
})
}
async fn execute(&self, input: Value) -> Result<ToolOutput> {
let query = input["query"]
.as_str()
.ok_or_else(|| anyhow::anyhow!("Missing required parameter: query"))?;
let top_k = input["top_k"].as_u64().unwrap_or(5);
let body = serde_json::json!({
"query": query,
"top_k": top_k,
});
let result = self
.client
.post(&self.endpoint)
.header("Content-Type", "application/json")
.json(&body)
.send()
.await;
match result {
Ok(response) => {
if response.status().is_success() {
let data: Value = response.json().await.unwrap_or_else(|_| {
serde_json::json!({"error": "Failed to parse response"})
});
let mut output = format!("Search results for: \"{}\"\n\n", query);
if let Some(results) = data["results"].as_array() {
for (i, result) in results.iter().enumerate() {
let title = result["title"].as_str().unwrap_or("Untitled");
let content = result["content"].as_str().unwrap_or("");
let score = result["score"].as_f64().unwrap_or(0.0);
output.push_str(&format!(
"{}. {} (relevance: {:.2})\n{}\n\n",
i + 1,
title,
score,
content
));
}
} else {
output.push_str(&serde_json::to_string_pretty(&data).unwrap_or_default());
}
Ok(ToolOutput::success(output))
} else {
let status = response.status();
let error_text = response.text().await.unwrap_or_default();
Ok(ToolOutput::error(format!(
"Tech query service returned error ({}): {}",
status, error_text
)))
}
}
Err(e) => {
Ok(ToolOutput::error(format!(
"Could not connect to tech query service at {}. \
Is your PyTorch documentation search system running? Error: {}",
self.endpoint, e
)))
}
}
}
}