feat(2690): make llm models configurable (#2716)

Co-authored-by: Tushar Mathur <[email protected]>
Co-authored-by: Sandipsinh Dilipsinh Rathod <[email protected]>
Co-authored-by: renovate[bot] <29139614+renovate[bot]@users.noreply.github.com>
Co-authored-by: Mehul Mathur <[email protected]>

src/cli/generator/config.rs

@@ -25,8 +25,18 @@ pub struct Config<Status = UnResolved> {
     #[serde(skip_serializing_if = "Option::is_none")]
     pub preset: Option<PresetConfig>,
     pub schema: Schema,
-    #[serde(skip_serializing_if = "TemplateString::is_empty")]
-    pub secret: TemplateString,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub llm: Option<LLMConfig>,
+}
+
+#[derive(Deserialize, Serialize, Debug, Default, PartialEq, Clone)]
+#[serde(rename_all = "camelCase")]
+#[serde(deny_unknown_fields)]
+pub struct LLMConfig {
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub model: Option<String>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub secret: Option<TemplateString>,
 }
 
 #[derive(Clone, Deserialize, Serialize, Debug, Default)]
@@ -273,13 +283,17 @@ impl Config {
             .collect::<anyhow::Result<Vec<Input<Resolved>>>>()?;
 
         let output = self.output.resolve(parent_dir)?;
+        let llm = self.llm.map(|llm| {
+            let secret = llm.secret.map(|s| s.resolve(&reader_context));
+            LLMConfig { model: llm.model, secret }
+        });
 
         Ok(Config {
             inputs,
             output,
             schema: self.schema,
             preset: self.preset,
-            secret: self.secret.resolve(&reader_context),
+            llm,
         })
     }
 }
@@ -419,7 +433,7 @@ mod tests {
     fn test_raise_error_unknown_field_at_root_level() {
         let json = r#"{"input": "value"}"#;
         let expected_error =
-            "unknown field `input`, expected one of `inputs`, `output`, `preset`, `schema`, `secret` at line 1 column 8";
+            "unknown field `input`, expected one of `inputs`, `output`, `preset`, `schema`, `llm` at line 1 column 8";
         assert_deserialization_error(json, expected_error);
     }
 
@@ -492,7 +506,7 @@ mod tests {
     }
 
     #[test]
-    fn test_secret() {
+    fn test_llm_config() {
         let mut env_vars = HashMap::new();
         let token = "eyJhbGciOiJIUzI1NiIsInR5";
         env_vars.insert("TAILCALL_SECRET".to_owned(), token.to_owned());
@@ -506,12 +520,17 @@ mod tests {
             headers: Default::default(),
         };
 
-        let config =
-            Config::default().secret(TemplateString::parse("{{.env.TAILCALL_SECRET}}").unwrap());
+        let config = Config::default().llm(Some(LLMConfig {
+            model: Some("gpt-3.5-turbo".to_string()),
+            secret: Some(TemplateString::parse("{{.env.TAILCALL_SECRET}}").unwrap()),
+        }));
         let resolved_config = config.into_resolved("", reader_ctx).unwrap();
 
-        let actual = resolved_config.secret;
-        let expected = TemplateString::from("eyJhbGciOiJIUzI1NiIsInR5");
+        let actual = resolved_config.llm;
+        let expected = Some(LLMConfig {
+            model: Some("gpt-3.5-turbo".to_string()),
+            secret: Some(TemplateString::from(token)),
+        });
 
         assert_eq!(actual, expected);
     }

src/cli/generator/generator.rs

@@ -6,7 +6,7 @@ use hyper::HeaderMap;
 use inquire::Confirm;
 use pathdiff::diff_paths;
 
-use super::config::{Config, Resolved, Source};
+use super::config::{Config, LLMConfig, Resolved, Source};
 use super::source::ConfigSource;
 use crate::cli::llm::InferTypeName;
 use crate::core::config::transformer::{Preset, RenameTypes};
@@ -164,7 +164,7 @@ impl Generator {
         let query_type = config.schema.query.clone();
         let mutation_type_name = config.schema.mutation.clone();
 
-        let secret = config.secret.clone();
+        let llm = config.llm.clone();
         let preset = config.preset.clone().unwrap_or_default();
         let preset: Preset = preset.validate_into().to_result()?;
         let input_samples = self.resolve_io(config).await?;
@@ -180,19 +180,15 @@ impl Generator {
         let mut config = config_gen.mutation(mutation_type_name).generate(true)?;
 
         if infer_type_names {
-            let key = if !secret.is_empty() {
-                Some(secret.to_string())
-            } else {
-                None
-            };
-
-            let mut llm_gen = InferTypeName::new(key);
-            let suggested_names = llm_gen.generate(config.config()).await?;
-            let cfg = RenameTypes::new(suggested_names.iter())
-                .transform(config.config().to_owned())
-                .to_result()?;
-
-            config = ConfigModule::from(cfg);
+            if let Some(LLMConfig { model: Some(model), secret }) = llm {
+                let mut llm_gen = InferTypeName::new(model, secret.map(|s| s.to_string()));
+                let suggested_names = llm_gen.generate(config.config()).await?;
+                let cfg = RenameTypes::new(suggested_names.iter())
+                    .transform(config.config().to_owned())
+                    .to_result()?;
+
+                config = ConfigModule::from(cfg);
+            }
         }
 
         self.write(&config, &path).await?;

src/cli/llm/infer_type_name.rs

@@ -4,14 +4,12 @@ use genai::chat::{ChatMessage, ChatRequest, ChatResponse};
 use serde::{Deserialize, Serialize};
 use serde_json::json;
 
-use super::model::groq;
 use super::{Error, Result, Wizard};
 use crate::core::config::Config;
 use crate::core::Mustache;
 
-#[derive(Default)]
 pub struct InferTypeName {
-    secret: Option<String>,
+    wizard: Wizard<Question, Answer>,
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -74,14 +72,11 @@ impl TryInto<ChatRequest> for Question {
 }
 
 impl InferTypeName {
-    pub fn new(secret: Option<String>) -> InferTypeName {
-        Self { secret }
+    pub fn new(model: String, secret: Option<String>) -> InferTypeName {
+        Self { wizard: Wizard::new(model, secret) }
     }
-    pub async fn generate(&mut self, config: &Config) -> Result<HashMap<String, String>> {
-        let secret = self.secret.as_ref().map(|s| s.to_owned());
-
-        let wizard: Wizard<Question, Answer> = Wizard::new(groq::LLAMA38192, secret);
 
+    pub async fn generate(&mut self, config: &Config) -> Result<HashMap<String, String>> {
         let mut new_name_mappings: HashMap<String, String> = HashMap::new();
 
         // removed root type from types.
@@ -104,7 +99,7 @@ impl InferTypeName {
 
             let mut delay = 3;
             loop {
-                let answer = wizard.ask(question.clone()).await;
+                let answer = self.wizard.ask(question.clone()).await;
                 match answer {
                     Ok(answer) => {
                         let name = &answer.suggestions.join(", ");

src/cli/llm/mod.rs

@@ -3,7 +3,6 @@ pub mod infer_type_name;
 pub use error::Error;
 use error::Result;
 pub use infer_type_name::InferTypeName;
-mod model;
 mod wizard;
 
 pub use wizard::Wizard;

src/cli/llm/wizard.rs

@@ -5,18 +5,17 @@ use genai::resolver::AuthResolver;
 use genai::Client;
 
 use super::Result;
-use crate::cli::llm::model::Model;
 
 #[derive(Setters, Clone)]
 pub struct Wizard<Q, A> {
     client: Client,
-    model: Model,
+    model: String,
     _q: std::marker::PhantomData<Q>,
     _a: std::marker::PhantomData<A>,
 }
 
 impl<Q, A> Wizard<Q, A> {
-    pub fn new(model: Model, secret: Option<String>) -> Self {
+    pub fn new(model: String, secret: Option<String>) -> Self {
         let mut config = genai::adapter::AdapterConfig::default();
         if let Some(key) = secret {
             config = config.with_auth_resolver(AuthResolver::from_key_value(key));