1use 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}