1use anyhow::Result;
  2use async_trait::async_trait;
  3use serde::Deserialize;
  4use std::collections::BTreeMap;
  5use std::fs;
  6use std::{
  7    path::{Path, PathBuf},
  8    rc::Rc,
  9};
 10use util::serde::default_true;
 11
 12use crate::example::{Example, ExampleContext, ExampleMetadata, JudgeAssertion};
 13
 14mod add_arg_to_trait_method;
 15mod code_block_citations;
 16mod file_search;
 17
 18pub fn all(examples_dir: &Path) -> Vec<Rc<dyn Example>> {
 19    let mut threads: Vec<Rc<dyn Example>> = vec![
 20        Rc::new(file_search::FileSearchExample),
 21        Rc::new(add_arg_to_trait_method::AddArgToTraitMethod),
 22        Rc::new(code_block_citations::CodeBlockCitations),
 23    ];
 24
 25    for example_path in list_declarative_examples(examples_dir).unwrap() {
 26        threads.push(Rc::new(DeclarativeExample::load(&example_path).unwrap()));
 27    }
 28
 29    threads
 30}
 31
 32struct DeclarativeExample {
 33    metadata: ExampleMetadata,
 34    prompt: String,
 35    diff_assertions: Vec<JudgeAssertion>,
 36    thread_assertions: Vec<JudgeAssertion>,
 37}
 38
 39impl DeclarativeExample {
 40    pub fn load(example_path: &Path) -> Result<Self> {
 41        let name = Self::name_from_path(example_path);
 42        let base: ExampleToml = toml::from_str(&fs::read_to_string(&example_path)?)?;
 43
 44        let language_server = if base.require_lsp {
 45            Some(crate::example::LanguageServer {
 46                file_extension: base
 47                    .language_extension
 48                    .expect("Language extension is required when require_lsp = true"),
 49                allow_preexisting_diagnostics: base.allow_preexisting_diagnostics,
 50            })
 51        } else {
 52            None
 53        };
 54
 55        let metadata = ExampleMetadata {
 56            name,
 57            url: base.url,
 58            revision: base.revision,
 59            language_server,
 60            max_assertions: None,
 61        };
 62
 63        Ok(DeclarativeExample {
 64            metadata,
 65            prompt: base.prompt,
 66            thread_assertions: base
 67                .thread_assertions
 68                .into_iter()
 69                .map(|(id, description)| JudgeAssertion { id, description })
 70                .collect(),
 71            diff_assertions: base
 72                .diff_assertions
 73                .into_iter()
 74                .map(|(id, description)| JudgeAssertion { id, description })
 75                .collect(),
 76        })
 77    }
 78
 79    pub fn name_from_path(path: &Path) -> String {
 80        path.file_stem().unwrap().to_string_lossy().to_string()
 81    }
 82}
 83
 84#[derive(Clone, Debug, Deserialize)]
 85pub struct ExampleToml {
 86    pub url: String,
 87    pub revision: String,
 88    pub language_extension: Option<String>,
 89    pub insert_id: Option<String>,
 90    #[serde(default = "default_true")]
 91    pub require_lsp: bool,
 92    #[serde(default)]
 93    pub allow_preexisting_diagnostics: bool,
 94    pub prompt: String,
 95    #[serde(default)]
 96    pub diff_assertions: BTreeMap<String, String>,
 97    #[serde(default)]
 98    pub thread_assertions: BTreeMap<String, String>,
 99}
100
101#[async_trait(?Send)]
102impl Example for DeclarativeExample {
103    fn meta(&self) -> ExampleMetadata {
104        self.metadata.clone()
105    }
106
107    async fn conversation(&self, cx: &mut ExampleContext) -> Result<()> {
108        cx.push_user_message(&self.prompt);
109        let _ = cx.run_to_end().await;
110        Ok(())
111    }
112
113    fn diff_assertions(&self) -> Vec<JudgeAssertion> {
114        self.diff_assertions.clone()
115    }
116
117    fn thread_assertions(&self) -> Vec<JudgeAssertion> {
118        self.thread_assertions.clone()
119    }
120}
121
122fn list_declarative_examples(examples_dir: &Path) -> Result<Vec<PathBuf>> {
123    let path = std::fs::canonicalize(examples_dir).unwrap();
124    let entries = std::fs::read_dir(path).unwrap();
125    let mut result_paths = Vec::new();
126    for entry in entries {
127        let entry = entry?;
128        let path = entry.path();
129        if path.extension() == Some("toml".as_ref()) {
130            result_paths.push(path);
131        }
132    }
133    Ok(result_paths)
134}