gigi-cli 1.0.0

Gigi — A Claude Code-like AI coding assistant CLI in Rust
use anyhow::Result;
use async_trait::async_trait;
use reqwest::Client;
use serde_json::Value;

use super::{Tool, ToolOutput};

// =============================================================================
// TechQueryTool — Bridge to your PyTorch-based documentation search system
// =============================================================================

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"})
                    });

                    // Format the results nicely
                    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 {
                        // If the response isn't in the expected format, just dump it
                        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) => {
                // Connection errors are expected when the service isn't running
                Ok(ToolOutput::error(format!(
                    "Could not connect to tech query service at {}. \
                     Is your PyTorch documentation search system running? Error: {}",
                    self.endpoint, e
                )))
            }
        }
    }
}