From c07f68d0032f491d9c6b4d048e5e9d093623f2d8 Mon Sep 17 00:00:00 2001
From: Din <dinmukhamed.mailibay@gmail.com>
Date: Wed, 31 Jul 2024 14:48:06 -0700
Subject: [PATCH] fix parsing description from enum and class fields

---
 Cargo.lock                   | 18 +++++++++---------
 Cargo.toml                   |  2 +-
 baml-lib/baml/pyproject.toml |  2 +-
 baml-lib/baml/src/lib.rs     | 23 ++++++++++++++++++-----
 4 files changed, 29 insertions(+), 16 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index d8d22eb..fce6a46 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -74,7 +74,7 @@ checksum = "0c4b4d0bd25bd0b74681c0ad21497610ce1b7c91b1022cd21c80c6fbdd9476b0"
 
 [[package]]
 name = "baml-types"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "indexmap 2.2.6",
  "minijinja",
@@ -391,7 +391,7 @@ checksum = "b248f5224d1d606005e02c97f5aa4e88eeb230488bcc03bc9ca4d7991399f2b5"
 
 [[package]]
 name = "internal-baml-core"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "anyhow",
  "baml-types",
@@ -421,7 +421,7 @@ dependencies = [
 
 [[package]]
 name = "internal-baml-diagnostics"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "anyhow",
  "pest",
@@ -430,7 +430,7 @@ dependencies = [
 
 [[package]]
 name = "internal-baml-jinja"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "anyhow",
  "baml-types",
@@ -446,7 +446,7 @@ dependencies = [
 
 [[package]]
 name = "internal-baml-parser-database"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "baml-types",
  "colored",
@@ -465,7 +465,7 @@ dependencies = [
 
 [[package]]
 name = "internal-baml-prompt-parser"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "internal-baml-diagnostics",
  "internal-baml-schema-ast",
@@ -477,7 +477,7 @@ dependencies = [
 
 [[package]]
 name = "internal-baml-schema-ast"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "baml-types",
  "internal-baml-diagnostics",
@@ -501,7 +501,7 @@ checksum = "49f1f14873335454500d59611f1cf4a4b0f786f9ac11f4312a78e4cf2566695b"
 
 [[package]]
 name = "jsonish"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "anyhow",
  "baml-types",
@@ -532,7 +532,7 @@ checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c"
 
 [[package]]
 name = "lmnr-baml"
-version = "0.40.0"
+version = "0.40.1"
 dependencies = [
  "anyhow",
  "baml-types",
diff --git a/Cargo.toml b/Cargo.toml
index a6d2a7f..73e1f19 100644
--- 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
index 1d0b843..84a7c5a 100644
--- 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
index 7fd7e40..8b38f29 100644
--- 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 {