fix parsing description from enum and class fields

lmnr-ai · Jul 31, 2024 · c07f68d · c07f68d
1 parent bddc773
commit c07f68d
Show file tree

Hide file tree

Showing 4 changed files with 29 additions and 16 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -34,7 +34,7 @@ internal-baml-core = { path = "baml-lib/baml-core" }
 internal-baml-jinja = { path = "baml-lib/jinja" }
 
 [workspace.package]
-version = "0.40.0"
+version = "0.40.1"
 authors = ["Boundary <contact@boundaryml.com>"]
 
 description = "BAML Toolchain"

diff --git a/baml-lib/baml/pyproject.toml b/baml-lib/baml/pyproject.toml
@@ -1,6 +1,6 @@
 [project]
 name = "lmnr_baml"
-version = "0.40.0"
+version = "0.40.1"
 description = "LMNR BAML for Python"
 readme = "README.md"
 

diff --git a/baml-lib/baml/src/lib.rs b/baml-lib/baml/src/lib.rs
@@ -3,7 +3,7 @@
 
 use baml_types::{BamlValue, FieldType};
 use either::Either;
-use internal_baml_core::ast::{WithDocumentation, WithName};
+use internal_baml_core::ast::{WithAttributes, WithName};
 pub use internal_baml_core::{
     self,
     internal_baml_diagnostics::{self, Diagnostics, SourceFile},
@@ -146,8 +146,15 @@ impl BamlContext {
                     .iter_values()
                     .map(|(_id, v)| {
                         let name = internal_baml_jinja::Name::new(v.name().to_string());
-                        let doc = v.documentation().map(|d| d.to_string());
-                        (name, doc)
+                        let description = v
+                            .attributes()
+                            .iter()
+                            .find(|a| a.name() == "description")
+                            .and_then(|a| a.arguments.iter().next())
+                            .and_then(|(_id, val)| val.value.as_string_value())
+                            .map(|ast_string_val| ast_string_val.0.to_string());
+                        // let doc = v.documentation().map(|d| d.to_string());
+                        (name, description)
                     })
                     .collect::<Vec<_>>();
                 internal_baml_jinja::Enum {
@@ -167,8 +174,14 @@ impl BamlContext {
                     .map(|(_id, f)| {
                         let name = internal_baml_jinja::Name::new(f.name().to_string());
                         let t = validated_schema.db.to_raw_field_type(&f.field_type);
-                        let doc = f.documentation().map(|d| d.to_string());
-                        (name, t, doc)
+                        let description = f
+                            .attributes()
+                            .iter()
+                            .find(|a| a.name() == "description")
+                            .and_then(|a| a.arguments.iter().next())
+                            .and_then(|(_id, val)| val.value.as_string_value())
+                            .map(|ast_string_val| ast_string_val.0.to_string());
+                        (name, t, description)
                     })
                     .collect::<Vec<_>>();
                 internal_baml_jinja::Class {