use anyhow::{Context, Result, bail};
use colored::Colorize;
use cosq_core::config::Config;
use cosq_core::stored_query::{
StoredQuery, StoredQueryMetadata, find_stored_query, list_stored_queries, query_file_path,
user_queries_dir,
};
use crate::cli::QueriesCommands;
pub async fn run(cmd: QueriesCommands, quiet: bool) -> Result<()> {
match cmd {
QueriesCommands::List => list(),
QueriesCommands::Create { name, project } => create(&name, project),
QueriesCommands::Edit { name } => edit(&name),
QueriesCommands::Delete { name, yes } => delete(&name, yes),
QueriesCommands::Show { name } => show(&name),
QueriesCommands::Generate {
description,
db,
container,
project,
} => generate(description, db, container, project, quiet).await,
}
}
fn list() -> Result<()> {
let queries = list_stored_queries().unwrap_or_default();
if queries.is_empty() {
println!("No stored queries found.");
println!(
"\n Create one with: {}",
"cosq queries create <name>".cyan()
);
return Ok(());
}
println!(
"{} ({}):\n",
"Stored queries".bold(),
"~/.cosq/queries/".dimmed()
);
let max_name_len = queries.iter().map(|q| q.name.len()).max().unwrap_or(0);
for query in &queries {
let ai_badge = if query.metadata.generated_by.is_some() {
" (AI)".dimmed().to_string()
} else {
String::new()
};
println!(
" {:<width$} {}{}",
query.name.green().bold(),
query.metadata.description.dimmed(),
ai_badge,
width = max_name_len,
);
}
println!("\n{} queries found.", queries.len());
Ok(())
}
fn create(name: &str, project: bool) -> Result<()> {
let path = query_file_path(name, project)?;
if path.exists() {
bail!(
"Query '{}' already exists at {}. Use `cosq queries edit {}` to modify it.",
name,
path.display(),
name
);
}
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let template = StoredQueryMetadata {
description: "TODO: describe what this query does".to_string(),
database: None,
container: None,
params: Vec::new(),
template: None,
template_file: None,
generated_by: None,
generated_from: None,
};
let yaml = serde_yaml::to_string(&template)?;
let contents =
format!("---\n{yaml}---\n-- Write your Cosmos DB SQL query below\nSELECT * FROM c\n");
std::fs::write(&path, &contents)?;
println!("{} Created {}", "OK".green().bold(), path.display());
open_in_editor(&path)?;
Ok(())
}
fn edit(name: &str) -> Result<()> {
let path = find_query_path(name)?;
let _ = StoredQuery::load(&path).map_err(|e| anyhow::anyhow!("Query file has errors: {e}"))?;
open_in_editor(&path)?;
Ok(())
}
fn delete(name: &str, yes: bool) -> Result<()> {
let path = find_query_path(name)?;
if !yes {
let confirm = dialoguer::Confirm::with_theme(&dialoguer::theme::ColorfulTheme::default())
.with_prompt(format!("Delete query '{name}' at {}?", path.display()))
.default(false)
.interact()
.context("confirmation cancelled")?;
if !confirm {
println!("Cancelled.");
return Ok(());
}
}
std::fs::remove_file(&path)?;
println!("{} Deleted query '{name}'.", "OK".green().bold());
Ok(())
}
fn show(name: &str) -> Result<()> {
let query =
find_stored_query(name).map_err(|e| anyhow::anyhow!("Query '{name}' not found: {e}"))?;
println!("{}", query.name.green().bold());
println!(" {} {}", "Description:".bold(), query.metadata.description);
if let Some(ref db) = query.metadata.database {
println!(" {} {}", "Database:".bold(), db);
}
if let Some(ref ctr) = query.metadata.container {
println!(" {} {}", "Container:".bold(), ctr);
}
let has_template = query.metadata.template.is_some() || query.metadata.template_file.is_some();
println!(
" {} {}",
"Template:".bold(),
if has_template { "yes" } else { "(none)" }
);
if let Some(ref generated) = query.metadata.generated_by {
println!(" {} {}", "Generated by:".bold(), generated);
}
if let Some(ref prompt) = query.metadata.generated_from {
println!(" {} \"{}\"", "Prompt:".bold(), prompt);
}
if !query.metadata.params.is_empty() {
println!("\n {}:", "Parameters".bold());
for param in &query.metadata.params {
let type_str = param.param_type.to_string();
let desc = param.description.as_deref().unwrap_or("");
let default_str = param
.default
.as_ref()
.map(|d| format!(" (default: {})", d))
.unwrap_or_default();
println!(
" {} <{}> {}{}",
format!("--{}", param.name).cyan(),
type_str,
desc,
default_str.dimmed(),
);
if let Some(ref choices) = param.choices {
let choices_str: Vec<String> = choices
.iter()
.map(|c| match c {
serde_json::Value::String(s) => s.clone(),
other => other.to_string(),
})
.collect();
println!(" {}: {}", "choices".dimmed(), choices_str.join(", "));
}
if let Some(min) = param.min {
print!(" {}: {min}", "min".dimmed());
}
if let Some(max) = param.max {
print!(" {}: {max}", "max".dimmed());
}
if param.min.is_some() || param.max.is_some() {
println!();
}
}
}
println!("\n {}:", "Query".bold());
for line in query.sql.lines() {
println!(" {}", line.dimmed());
}
if let Some(ref tmpl) = query.metadata.template {
println!("\n {}:", "Template".bold());
for line in tmpl.lines() {
println!(" {}", line.dimmed());
}
}
Ok(())
}
async fn generate(
description: Option<String>,
cli_db: Option<String>,
cli_container: Option<String>,
project: bool,
quiet: bool,
) -> Result<()> {
let mut config = Config::load()?;
let ai_config = config.ai.clone().ok_or_else(|| {
anyhow::anyhow!("AI is not configured. Run `cosq ai init` to set up an AI provider.")
})?;
let client = cosq_client::cosmos::CosmosClient::new(&config.account.endpoint).await?;
let (database, db_changed) =
super::common::resolve_database(&client, &mut config, cli_db, None).await?;
let (container, ctr_changed) =
super::common::resolve_container(&client, &mut config, &database, cli_container, None)
.await?;
if db_changed || ctr_changed {
config.save()?;
}
if !quiet {
eprintln!(
"{}",
format!("Sampling documents from {container}...").dimmed()
);
}
let sample_result = client
.query(&database, &container, "SELECT TOP 3 * FROM c")
.await
.context("failed to sample documents for schema context")?;
let sample_json = if sample_result.documents.is_empty() {
"(container is empty — no sample documents available)".to_string()
} else {
format_sample_documents(&sample_result.documents)
};
let description = if let Some(desc) = description {
desc
} else {
eprintln!();
dialoguer::Input::with_theme(&dialoguer::theme::ColorfulTheme::default())
.with_prompt("Describe the query you want to generate")
.interact_text()
.context("input cancelled")?
};
let system_prompt = build_system_prompt(&database, &container, &sample_json);
let user_prompt = format!("Generate a .cosq stored query for: {description}");
if !quiet {
eprintln!(
"{}",
format!("Generating via {}...", ai_config.provider.display_name()).dimmed()
);
}
let mut conversation_prompt = user_prompt;
let mut query = None;
let max_rounds = 3;
for round in 0..max_rounds {
let response =
cosq_client::ai::generate_text(&ai_config, &system_prompt, &conversation_prompt)
.await
.context("failed to generate query")?;
let content = strip_markdown_fences(&response);
match StoredQuery::parse("generated", &content) {
Ok(parsed) => {
query = Some(parsed);
break;
}
Err(_) if round < max_rounds - 1 => {
eprintln!("\n{}", "The AI needs clarification:".bold());
for line in content.lines() {
if !line.trim().is_empty() {
eprintln!(" {}", line);
}
}
eprintln!();
let answer: String =
dialoguer::Input::with_theme(&dialoguer::theme::ColorfulTheme::default())
.with_prompt("Your answer")
.interact_text()
.context("input cancelled")?;
conversation_prompt = format!(
"Original request: {description}\n\n\
You asked:\n{content}\n\n\
User answered: {answer}\n\n\
Now generate the .cosq file."
);
if !quiet {
eprintln!(
"{}",
format!("Generating via {}...", ai_config.provider.display_name()).dimmed()
);
}
}
Err(_) => {
bail!(
"AI did not produce a valid query after {max_rounds} attempts. \
Try rephrasing your description or run with more detail."
);
}
}
}
let mut query = query.ok_or_else(|| {
anyhow::anyhow!("AI did not produce a valid query. Try rephrasing your description.")
})?;
query.metadata.database = Some(database);
query.metadata.container = Some(container);
let provider_info = match ai_config.effective_model() {
Some(model) => format!("{} ({})", ai_config.provider.display_name(), model),
None => ai_config.provider.display_name().to_string(),
};
query.metadata.generated_by = Some(provider_info);
query.metadata.generated_from = Some(description.clone());
let suggested_name = generate_filename(&description);
eprintln!("\n{}", "Generated query:".bold());
eprintln!(" {} {}", "Name:".dimmed(), suggested_name.green());
eprintln!(
" {} {}",
"Description:".dimmed(),
query.metadata.description
);
if !query.metadata.params.is_empty() {
eprintln!(" {}:", "Parameters".dimmed());
for p in &query.metadata.params {
let default_str = p
.default
.as_ref()
.map(|d| format!(" (default: {d})"))
.unwrap_or_default();
eprintln!(
" --{} <{}>{}",
p.name,
p.param_type,
default_str.dimmed()
);
}
}
eprintln!("\n {}:", "SQL".dimmed());
for line in query.sql.lines() {
eprintln!(" {}", line.cyan());
}
if let Some(ref tmpl) = query.metadata.template {
eprintln!("\n {}:", "Template".dimmed());
for line in tmpl.lines() {
eprintln!(" {}", line);
}
}
let name: String = dialoguer::Input::with_theme(&dialoguer::theme::ColorfulTheme::default())
.with_prompt("Query name")
.default(suggested_name)
.interact_text()
.context("input cancelled")?;
query.name = name.clone();
let path = query_file_path(&name, project)?;
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)?;
}
let contents = query.to_file_contents()?;
std::fs::write(&path, &contents)?;
println!("{} Saved to {}", "OK".green().bold(), path.display());
let edit = dialoguer::Confirm::with_theme(&dialoguer::theme::ColorfulTheme::default())
.with_prompt("Open in editor to review?")
.default(true)
.interact()
.context("confirmation cancelled")?;
if edit {
open_in_editor(&path)?;
}
Ok(())
}
fn build_system_prompt(database: &str, container: &str, sample_json: &str) -> String {
format!(
r#"You are a Cosmos DB SQL query generator. You create .cosq stored query files.
TARGET:
Database: "{database}"
Container: "{container}"
SAMPLE DOCUMENTS from this container:
{sample_json}
FORMAT — .cosq files use YAML front matter between --- delimiters, followed by the SQL query.
SQL RULES:
- ONLY reference fields that exist in the sample documents above
- Use 'c' as the container alias (e.g., SELECT * FROM c)
- Use Cosmos DB SQL syntax (TOP not LIMIT, no OFFSET, use DateTimeAdd/GetCurrentDateTime for dates)
- Extract variable parts as @param parameters
- Parameters: define in params section with name, type (string/number/bool), description, and optional default/choices/min/max
- Set database and container in the YAML metadata
OUTPUT TEMPLATE RULES — ALWAYS include a MiniJinja template in the .cosq file:
- Templates use {{{{ variable }}}} syntax and {{% for doc in documents %}} loops
- Templates have access to 'documents' (array of results) and all parameter values
Choose the template style based on the query:
- MULTIPLE results (lists, searches, TOP N where N > 1): use a readable table layout
Example:
{{% for doc in documents %}}
{{{{ doc.displayName | truncate(25) }}}} {{{{ doc.email }}}} {{{{ doc.status }}}}
{{% endfor %}}
Total: {{{{ documents | length }}}}
- SINGLE result (TOP 1, lookup by ID, "latest", "last"): use a property-value layout
Example:
{{% for doc in documents %}}
Name: {{{{ doc.name }}}}
Email: {{{{ doc.email }}}}
Status: {{{{ doc.status }}}}
{{% endfor %}}
- If the user mentions "JSON" output: do NOT include a template field at all
- If the user mentions "CSV" output: create a template with a header row and comma-separated values
- For complex nested documents: format nested objects and arrays readably (indented sub-properties, bullet lists for arrays)
- When in doubt, pick the format that makes the data most readable
CONVERSATION RULES:
- Be CONFIDENT. If you can make a reasonable assumption, make it and generate the query.
- Only ask clarifying questions if the description is genuinely ambiguous (e.g., which field to filter on, or the user mentions something not in the schema)
- When asking questions, ask 1-3 short questions. Do NOT generate a .cosq file in the same response.
- When generating, respond with ONLY the .cosq file content — no explanation, no markdown fences."#
)
}
fn strip_markdown_fences(response: &str) -> String {
let trimmed = response.trim();
let stripped = trimmed
.strip_prefix("```yaml")
.or_else(|| trimmed.strip_prefix("```cosq"))
.or_else(|| trimmed.strip_prefix("```"))
.unwrap_or(trimmed);
let stripped = stripped.strip_suffix("```").unwrap_or(stripped);
stripped.trim().to_string()
}
fn format_sample_documents(docs: &[serde_json::Value]) -> String {
let truncated: Vec<serde_json::Value> = docs.iter().map(truncate_for_prompt).collect();
if let Ok(json) = serde_json::to_string_pretty(&truncated) {
if json.len() <= 4000 {
return json;
}
}
for n in (1..truncated.len()).rev() {
if let Ok(json) = serde_json::to_string_pretty(&truncated[..n]) {
if json.len() <= 4000 {
return format!("{json}\n(showing {n} of {} sampled documents)", docs.len());
}
}
}
serde_json::to_string_pretty(&truncated[..1])
.unwrap_or_else(|_| "(documents too large to display)".to_string())
}
fn truncate_for_prompt(value: &serde_json::Value) -> serde_json::Value {
use serde_json::Value;
match value {
Value::String(s) if s.len() > 100 => Value::String(format!("{}...", &s[..80])),
Value::Array(arr) if arr.len() > 3 => {
let mut truncated: Vec<Value> = arr[..3].iter().map(truncate_for_prompt).collect();
truncated.push(Value::String(format!("... ({} items total)", arr.len())));
Value::Array(truncated)
}
Value::Array(arr) => Value::Array(arr.iter().map(truncate_for_prompt).collect()),
Value::Object(map) => Value::Object(
map.iter()
.map(|(k, v)| (k.clone(), truncate_for_prompt(v)))
.collect(),
),
other => other.clone(),
}
}
fn generate_filename(description: &str) -> String {
let words: Vec<&str> = description
.split_whitespace()
.filter(|w| {
!matches!(
w.to_lowercase().as_str(),
"find"
| "get"
| "list"
| "show"
| "select"
| "all"
| "the"
| "a"
| "an"
| "from"
| "in"
| "of"
| "by"
| "with"
| "who"
| "that"
| "are"
| "is"
| "and"
| "or"
| "for"
| "to"
)
})
.take(4)
.collect();
let name = words.join("-").to_lowercase();
name.chars()
.filter(|c| c.is_alphanumeric() || *c == '-')
.collect()
}
fn open_in_editor(path: &std::path::Path) -> Result<()> {
let editor = std::env::var("VISUAL")
.or_else(|_| std::env::var("EDITOR"))
.unwrap_or_else(|_| {
if cfg!(target_os = "macos") {
"open".to_string()
} else if cfg!(target_os = "windows") {
"notepad".to_string()
} else {
"xdg-open".to_string()
}
});
eprintln!("{} Opening in {editor}...", ">>".dimmed());
std::process::Command::new(&editor)
.arg(path)
.status()
.with_context(|| format!("failed to open editor: {editor}"))?;
Ok(())
}
fn find_query_path(name: &str) -> Result<std::path::PathBuf> {
let filename = if name.ends_with(".cosq") {
name.to_string()
} else {
format!("{name}.cosq")
};
if let Some(project_dir) = cosq_core::stored_query::project_queries_dir() {
let path = project_dir.join(&filename);
if path.exists() {
return Ok(path);
}
}
let user_dir = user_queries_dir().map_err(|e| anyhow::anyhow!("{e}"))?;
let path = user_dir.join(&filename);
if path.exists() {
return Ok(path);
}
bail!("Query '{name}' not found. Run `cosq queries list` to see available queries.")
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_generate_filename() {
assert_eq!(
generate_filename("Find all users who signed up in the last 30 days"),
"users-signed-up-last"
);
assert_eq!(generate_filename("orders over $100"), "orders-over-100");
assert_eq!(
generate_filename("active subscriptions"),
"active-subscriptions"
);
}
#[test]
fn test_generate_filename_empty_input() {
let name = generate_filename("users");
assert_eq!(name, "users");
}
#[test]
fn test_strip_markdown_fences_yaml() {
let input = "```yaml\n---\ndescription: test\n---\nSELECT * FROM c\n```";
assert_eq!(
strip_markdown_fences(input),
"---\ndescription: test\n---\nSELECT * FROM c"
);
}
#[test]
fn test_strip_markdown_fences_plain() {
let input = "```\nsome content\n```";
assert_eq!(strip_markdown_fences(input), "some content");
}
#[test]
fn test_strip_markdown_fences_no_fences() {
let input = "---\ndescription: test\n---\nSELECT * FROM c";
assert_eq!(strip_markdown_fences(input), input);
}
#[test]
fn test_truncate_for_prompt_short_values() {
use serde_json::json;
let value = json!({"name": "short", "count": 42});
let result = truncate_for_prompt(&value);
assert_eq!(result, value);
}
#[test]
fn test_truncate_for_prompt_long_string() {
use serde_json::json;
let long = "x".repeat(200);
let value = json!({"content": long});
let result = truncate_for_prompt(&value);
let content = result["content"].as_str().unwrap();
assert!(content.len() < 200);
assert!(content.ends_with("..."));
}
#[test]
fn test_truncate_for_prompt_large_array() {
use serde_json::json;
let value = json!({"items": [1, 2, 3, 4, 5, 6, 7]});
let result = truncate_for_prompt(&value);
let items = result["items"].as_array().unwrap();
assert_eq!(items.len(), 4); }
#[test]
fn test_format_sample_documents() {
use serde_json::json;
let docs = vec![json!({"id": "1", "name": "test"})];
let formatted = format_sample_documents(&docs);
assert!(formatted.contains("\"id\": \"1\""));
assert!(formatted.contains("\"name\": \"test\""));
}
}