Skip to main content

ddb_core/mcp/
prompts.rs

1//! MCP Prompt implementations for DDB
2
3use anyhow::{anyhow, Result};
4use serde_json::{json, Value};
5
6pub struct PromptHandler;
7
8impl PromptHandler {
9    pub fn new() -> Self {
10        Self
11    }
12
13    pub fn list_prompts(&self) -> Vec<Value> {
14        vec![
15            json!({
16                "name": "query_data",
17                "description": "Generate a SQL query to retrieve data from a table",
18                "arguments": [
19                    { "name": "table", "description": "Table name", "required": true },
20                    { "name": "columns", "description": "Columns to select", "required": false }
21                ]
22            }),
23            json!({
24                "name": "analyze_data",
25                "description": "Generate queries for data analysis",
26                "arguments": [
27                    { "name": "table", "description": "Table name", "required": true },
28                    { "name": "metric", "description": "Metric to analyze", "required": true }
29                ]
30            }),
31        ]
32    }
33
34    pub async fn get_prompt(&self, name: &str, arguments: Option<Value>) -> Result<Value> {
35        let args = arguments.unwrap_or(json!({}));
36
37        match name {
38            "query_data" => {
39                let table = args["table"].as_str().ok_or_else(|| anyhow!("Missing table"))?;
40                let columns = args["columns"].as_str().unwrap_or("*");
41
42                let query = format!("SELECT {} FROM {}", columns, table);
43                Ok(json!({
44                    "messages": [{
45                        "role": "assistant",
46                        "content": {
47                            "type": "text",
48                            "text": format!("Here's a query for {}:\n\n```sql\n{}\n```", table, query)
49                        }
50                    }]
51                }))
52            }
53            "analyze_data" => {
54                let table = args["table"].as_str().ok_or_else(|| anyhow!("Missing table"))?;
55                let metric = args["metric"].as_str().unwrap_or("COUNT");
56
57                let query = if metric.to_uppercase() == "COUNT" {
58                    format!("SELECT COUNT(*) FROM {}", table)
59                } else {
60                    format!("SELECT {}(*) FROM {}", metric.to_uppercase(), table)
61                };
62
63                Ok(json!({
64                    "messages": [{
65                        "role": "assistant",
66                        "content": {
67                            "type": "text",
68                            "text": format!("Here's an analysis query:\n\n```sql\n{}\n```", query)
69                        }
70                    }]
71                }))
72            }
73            _ => Err(anyhow!("Unknown prompt: {}", name)),
74        }
75    }
76}
77
78impl Default for PromptHandler {
79    fn default() -> Self {
80        Self::new()
81    }
82}