add a separator to maps

engine/baml-lib/ast/src/parser/datamodel.pest

@@ -102,7 +102,7 @@ arguments_list = { "(" ~ (NEWLINE?) ~ expression? ~ ("," ~ (NEWLINE?) ~ expressi
 // Expressions & Functions
 // ######################################
 map_key   = { identifier | quoted_string_literal }
-map_entry = { (comment_block | empty_lines)* ~ map_key ~ (expression | ENTRY_CATCH_ALL)? ~ trailing_comment? }
+map_entry = { (comment_block | empty_lines)* ~ map_key ~ COLON? ~ empty_lines* ~ (expression | ENTRY_CATCH_ALL)? ~ trailing_comment? }
 
 splitter          = _{ ("," ~ NEWLINE?) | NEWLINE }
 map_expression    =  { "{" ~ empty_lines? ~ (map_entry ~ (splitter ~ map_entry)*)? ~ (comment_block | empty_lines)* ~ "}" }

engine/baml-lib/ast/src/parser/parse_expression.rs

@@ -155,6 +155,9 @@ fn parse_map_entry(
         match current.as_rule() {
             Rule::map_key => key = Some(parse_map_key(current, diagnostics)),
             Rule::expression => value = parse_expression(current, diagnostics),
+            Rule::COLON => {
+                // colon is optional, so we just ignore it
+            }
             Rule::ENTRY_CATCH_ALL => {
                 diagnostics.push_error(
                     internal_baml_diagnostics::DatamodelError::new_validation_error(
@@ -416,6 +419,25 @@ mod tests {
         };
     }
 
+    #[test]
+    fn test_parse_map() {
+        let input = "{\"foo\": 1, \"bar\": 2}";
+        let root_path = "test_file.baml";
+        let source = SourceFile::new_static(root_path.into(), input);
+        let mut diagnostics = Diagnostics::new(root_path.into());
+        diagnostics.set_source(&source);
+
+        let pair = BAMLParser::parse(Rule::map_expression, input)
+            .unwrap()
+            .next()
+            .unwrap();
+        let expr = parse_map(pair, &mut diagnostics);
+        match expr {
+            Expression::Map(entries, _) => assert_eq!(entries.len(), 2),
+            _ => panic!("Expected Map, got {expr:?}"),
+        }
+    }
+
     #[test]
     fn test_parse_jinja_expression() {
         let input = "{{ 1 + 1 }}";
@@ -434,4 +456,45 @@ mod tests {
             _ => panic!("Expected JinjaExpression, got {expr:?}"),
         }
     }
+
+    #[test]
+    fn test_parse_map_without_colons() {
+        let input = "{foo 1, bar 2}";
+        let root_path = "test_file.baml";
+        let source = SourceFile::new_static(root_path.into(), input);
+        let mut diagnostics = Diagnostics::new(root_path.into());
+        diagnostics.set_source(&source);
+
+        let pair = BAMLParser::parse(Rule::map_expression, input)
+            .unwrap()
+            .next()
+            .unwrap();
+        let expr = parse_map(pair, &mut diagnostics);
+        match expr {
+            Expression::Map(entries, _) => assert_eq!(entries.len(), 2),
+            _ => panic!("Expected Map, got {expr:?}"),
+        }
+    }
+
+    #[test]
+    fn test_parse_map_multiline_without_colons() {
+        let input = r#"{
+    foo 1
+    bar 2
+}"#;
+        let root_path = "test_file.baml";
+        let source = SourceFile::new_static(root_path.into(), input);
+        let mut diagnostics = Diagnostics::new(root_path.into());
+        diagnostics.set_source(&source);
+
+        let pair = BAMLParser::parse(Rule::map_expression, input)
+            .unwrap()
+            .next()
+            .unwrap();
+        let expr = parse_map(pair, &mut diagnostics);
+        match expr {
+            Expression::Map(entries, _) => assert_eq!(entries.len(), 2),
+            _ => panic!("Expected Map, got {expr:?}"),
+        }
+    }
 }