use get method and update tracing message

Signed-off-by: Sahil Yeole <[email protected]>
beelchester · Aug 16, 2024 · 650c06a · 650c06a
1 parent 859daf2
commit 650c06a
Showing 1 changed file with 2 additions and 6 deletions.
diff --git a/src/cli/llm/infer_arg_name.rs b/src/cli/llm/infer_arg_name.rs
@@ -83,11 +83,7 @@ impl InferArgName {
 
         let mut new_name_mappings: HashMap<String, String> = HashMap::new();
 
-        let query_type = config
-            .types
-            .iter()
-            .find(|(type_name, _)| *type_name == "Query")
-            .map(|(_, type_)| type_);
+        let query_type = config.types.get("Query");
 
         if let Some(type_) = query_type {
             let mut args_to_be_processed = HashMap::new();
@@ -124,7 +120,7 @@ impl InferArgName {
                                 break;
                             }
                             tracing::info!(
-                                "Suggestions for {}: [{}] - {}/{}",
+                                "Suggestions for argument {}: [{}] - {}/{}",
                                 arg.0,
                                 name,
                                 i + 1,