This works!

BoundaryML · Mar 5, 2024 · d3cfb56 · d3cfb56
1 parent ffea2ee
commit d3cfb56
Show file tree

Hide file tree

Showing 15 changed files with 219 additions and 50 deletions.
diff --git a/clients/ts/src/client_manager/llm_base_provider.ts b/clients/ts/src/client_manager/llm_base_provider.ts
@@ -53,6 +53,10 @@ abstract class LLMBaseProvider extends BaseProvider {
       throw new Error(`Unknown arguments: ${Object.keys(rest).join(', ')}`);
     }
 
+    if (typeof provider !== 'string') {
+      throw new Error(`provider must be a string: ${provider}`);
+    }
+
     super();
 
     this.provider = provider;

diff --git a/engine/Cargo.lock b/engine/Cargo.lock
diff --git a/engine/baml-client-lib/Cargo.toml b/engine/baml-client-lib/Cargo.toml
@@ -28,6 +28,7 @@ lazy_static = "1.4.0"
 neon = { version = "1.0.0-alpha.4", features = ["napi-6", "futures"] }
 tracing-subscriber = "0.3.18"
 hostname = "0.3.1"
+colored = "2.1.0"
 # tokio_unstable, feature = "tracing"
 # tokio_unstable = { version = "*", features = ["tracing"] }
 # [dependencies.neon]

diff --git a/engine/baml-client-lib/index.js b/engine/baml-client-lib/index.js
@@ -66,7 +66,12 @@ const tracerAsync = (cb, name, args, asKwargs, returnType) => {
 };
 
 function logLLMEvent(event) {
-  logLLMEventNode({
+  const span = getSpan();
+  if (span === undefined) {
+    console.warn('BAML: Attempting to call an LLM event without an active span. Ignoring.');
+    return;
+  }
+  logLLMEventNode(span, {
     name: event.name, 
     meta: JSON.stringify(event.data)
   });

diff --git a/engine/baml-client-lib/index.node b/engine/baml-client-lib/index.node
diff --git a/engine/baml-client-lib/src/api_wrapper/core_types.rs b/engine/baml-client-lib/src/api_wrapper/core_types.rs
@@ -1,5 +1,7 @@
-use std::collections::HashMap;
+use colored::*;
+use std::{collections::HashMap, str::FromStr};
 
+use neon::meta;
 use serde::{Deserialize, Serialize};
 use serde_json::Value;
 
@@ -201,12 +203,147 @@ impl Default for Template {
     }
 }
 
+impl LLMEventInputPrompt {
+    fn pretty_print(&self) -> String {
+        match &self.template {
+            Template::Single(s) => s.clone(),
+            Template::Multiple(chats) => chats
+                .iter()
+                .map(|chat| {
+                    format!(
+                        "{} {}\n{}",
+                        "Role:".yellow(),
+                        chat.role.as_str().yellow(),
+                        chat.content
+                    )
+                })
+                .collect::<Vec<String>>()
+                .join("\n"),
+        }
+    }
+}
+
 impl LogSchema {
     pub fn pretty_string(&self) -> Option<String> {
         match self.event_type {
             EventType::FuncLlm => {
-                // Returns a colored string representation of the LogSchema
-                todo!();
+                let log = self;
+                let (llm_prompt, llm_raw_output) = match log.metadata.as_ref().and_then(|meta| {
+                    // TODO: Swap out template vars
+                    let input = match &meta.input.prompt.template {
+                        Template::Single(o) => o.clone(),
+                        Template::Multiple(chats) => chats
+                            .iter()
+                            .map(|c| {
+                                format!(
+                                    "{}:\n{}",
+                                    c.role.as_str().yellow().bold(),
+                                    c.content.white()
+                                )
+                            })
+                            .collect::<Vec<_>>()
+                            .join("\n"),
+                    };
+
+                    let mut colored_input = input.clone();
+                    meta.input.prompt.template_args.iter().for_each(|(k, v)| {
+                        let replacement = format!("{}", v.blue()); // Colorize the replacement text in magenta
+                        colored_input = colored_input.replace(k, &replacement);
+                    });
+
+                    let raw_output = match &meta.output {
+                        Some(output) => Some(output.raw_text.clone()),
+                        None => None,
+                    };
+
+                    Some((colored_input, raw_output))
+                }) {
+                    Some((llm_prompt, llm_raw_output)) => (Some(llm_prompt), llm_raw_output),
+                    None => (None, None),
+                };
+
+                let err = log.error.as_ref().and_then(|error| match &error.traceback {
+                    Some(traceback) => Some(format!("{}\n{}", error.message, traceback)),
+                    None => Some(error.message.clone()),
+                });
+
+                let parsed_output = match log.io.output.as_ref().and_then(|output| {
+                    let r#type = match &output.r#type.name {
+                        TypeSchemaName::Single => {
+                            let fields = output
+                                .r#type
+                                .fields
+                                .iter()
+                                .map(|(k, v)| format!("{}: {}", k, v))
+                                .collect::<Vec<_>>()
+                                .join(", ");
+                            format!("{}", fields)
+                        }
+                        TypeSchemaName::Multi => {
+                            let fields = output
+                                .r#type
+                                .fields
+                                .iter()
+                                .map(|(k, v)| format!("{}: {}", k, v))
+                                .collect::<Vec<_>>()
+                                .join(", ");
+                            format!("{}", fields)
+                        }
+                    };
+                    let output = match &output.value {
+                        ValueType::String(s) => serde_json::Value::from_str(s),
+                        ValueType::List(l) => l
+                            .iter()
+                            .map(|v| serde_json::Value::from_str(v))
+                            .collect::<Result<Vec<_>, _>>()
+                            .map(|v| serde_json::Value::Array(v)),
+                    }
+                    .map(|v| {
+                        serde_json::to_string_pretty(&v)
+                            .unwrap_or_else(|_| format!("Failed to serialize output: {:?}", v))
+                    })
+                    .ok();
+                    Some((output, Some(r#type)))
+                }) {
+                    Some((Some(output), Some(r#type))) => Some((output, r#type)),
+                    _ => None,
+                };
+
+                let res = match (llm_prompt, llm_raw_output, err, parsed_output) {
+                    (Some(llm_prompt), Some(llm_raw_output), Some(err), _) => vec![
+                        format!("\n{}", "---- Prompt ---------".dimmed()),
+                        format!("{}", llm_prompt),
+                        format!("\n{}", "---- Raw Response ---".dimmed()),
+                        format!("{}", llm_raw_output.white()),
+                        format!("{}", "----- Error -----".dimmed()),
+                        format!("{}", err.red()),
+                    ],
+                    (Some(llm_prompt), None, Some(err), _) => vec![
+                        format!("\n{}", "---- Prompt ---------".dimmed()),
+                        format!("{}", llm_prompt),
+                        format!("{}", "----- Error -----".dimmed()),
+                        format!("{}", err.red()),
+                    ],
+                    (Some(llm_prompt), Some(llm_raw_output), None, Some((output, output_type))) => {
+                        vec![
+                            format!("\n{}", "------- Prompt ------".yellow()),
+                            format!("{}", llm_prompt),
+                            format!("\n{}", "---- Raw Response ---".dimmed()),
+                            format!("{}", llm_raw_output.dimmed()),
+                            format!(
+                                "\n{}{}{}",
+                                "----- Parsed Response (".green(),
+                                output_type.green(),
+                                ") -----".green()
+                            ),
+                            format!("{}", output.green()),
+                        ]
+                    }
+                    _ => vec![],
+                }
+                .join("\n");
+
+                Some(res)
             }
             _ => None,
         }

diff --git a/engine/baml-client-lib/src/api_wrapper/mod.rs b/engine/baml-client-lib/src/api_wrapper/mod.rs
@@ -1,4 +1,5 @@
 use anyhow::Result;
+use colored::*;
 pub(super) mod api_interface;
 pub(super) mod core_types;
 use serde_json::Value;
@@ -22,7 +23,7 @@ impl APIWrapper {
 
         if log_level {
             match payload.pretty_string() {
-                Some(s) => println!("{}", s),
+                Some(s) => println!("{s}"),
                 None => println!("Failed to pretty print log schema {:?}", payload),
             }
         }

diff --git a/engine/baml-client-lib/src/lib.rs b/engine/baml-client-lib/src/lib.rs
@@ -519,13 +519,8 @@ fn set_tags(mut cx: FunctionContext) -> JsResult<JsUndefined> {
 }
 
 fn set_llm_event(mut cx: FunctionContext) -> JsResult<JsUndefined> {
-    let event = cx.argument::<JsObject>(0)?;
-    /*
-        Event is {
-            name: string,
-            meta: string // JSON string
-        }
-    */
+    let span = cx.argument::<JsBox<BamlSpanOwner>>(0)?;
+    let event = cx.argument::<JsObject>(1)?;
 
     let name: Handle<JsString> = event
         .get_value(&mut cx, "name")?
@@ -538,6 +533,7 @@ fn set_llm_event(mut cx: FunctionContext) -> JsResult<JsUndefined> {
 
     let meta = meta.value(&mut cx);
 
+    let _guard = span.enter();
     match otel::log_event(name.as_str(), meta.as_str()) {
         Ok(_) => Ok(cx.undefined()),
         Err(e) => cx.throw_error(e.to_string()),

diff --git a/engine/baml-client-lib/src/otel/span_events.rs b/engine/baml-client-lib/src/otel/span_events.rs
@@ -17,7 +17,7 @@ use tracing::Subscriber;
 use tracing_subscriber::layer::Layer;
 
 use crate::{
-    api_wrapper::core_types::{EventChain, LogSchemaContext},
+    api_wrapper::core_types::{EventChain, EventType, LogSchemaContext},
     baml_event,
 };
 
@@ -201,7 +201,14 @@ where
                 if let Err(e) = parse_event(event, &span) {
                     println!("Error parsing event: {:?}", e);
                 }
+            } else {
+                println!(
+                    "No inner span found for event {:?}",
+                    event.metadata().name()
+                );
             }
+        } else {
+            println!("No span id found for event: {:?}", event.metadata().name());
         }
     }
 

diff --git a/engine/baml-client-lib/src/otel/span_events/llm_prompt_template.rs b/engine/baml-client-lib/src/otel/span_events/llm_prompt_template.rs
@@ -59,13 +59,16 @@ where
         event: LlmPromptTemplate,
         _span: &tracing_subscriber::registry::SpanRef<'a, S>,
     ) {
-        let meta = self.get_meta_data_mut(true).unwrap();
-        if let Some(input) = &mut meta.input {
-            input.prompt = LLMEventInputPrompt {
-                template: event.template,
-                template_args: event.template_args,
-                ..Default::default()
+        if let Some(meta) = self.get_meta_data_mut(false) {
+            if let Some(input) = &mut meta.input {
+                input.prompt = LLMEventInputPrompt {
+                    template: event.template,
+                    template_args: event.template_args,
+                    ..Default::default()
+                }
             }
+        } else {
+            println!("No metadata found for llm event");
         }
     }
 }
diff --git a/engine/baml-client-lib/src/otel/span_events/llm_request_args.rs b/engine/baml-client-lib/src/otel/span_events/llm_request_args.rs
@@ -48,9 +48,12 @@ where
         event: LlmRequestArgs,
         _span: &tracing_subscriber::registry::SpanRef<'a, S>,
     ) {
-        let meta = self.get_meta_data_mut(true).unwrap();
-        if let Some(input) = &mut meta.input {
-            input.invocation_params = event.invocation_params;
+        if let Some(meta) = self.get_meta_data_mut(false) {
+            if let Some(input) = &mut meta.input {
+                input.invocation_params = event.invocation_params;
+            }
+        } else {
+            println!("No metadata found for llm event");
         }
     }
 }
diff --git a/engine/baml-client-lib/src/otel/span_events/llm_request_end.rs b/engine/baml-client-lib/src/otel/span_events/llm_request_end.rs
@@ -70,23 +70,26 @@ where
         event: LlmRequestEnd,
         _span: &tracing_subscriber::registry::SpanRef<'a, S>,
     ) {
-        let meta = self.get_meta_data_mut(true).unwrap();
-        meta.model_name = event.model_name;
-        match (event.generated, event.metadata) {
-            (None, _) => {}
-            (Some(generated), Some(metadata)) => {
-                meta.output = Some(LLMOutputModel {
-                    raw_text: generated,
-                    metadata,
-                    ..Default::default()
-                });
-            }
-            (Some(generated), None) => {
-                meta.output = Some(LLMOutputModel {
-                    raw_text: generated,
-                    ..Default::default()
-                });
+        if let Some(meta) = self.get_meta_data_mut(false) {
+            meta.model_name = event.model_name;
+            match (event.generated, event.metadata) {
+                (None, _) => {}
+                (Some(generated), Some(metadata)) => {
+                    meta.output = Some(LLMOutputModel {
+                        raw_text: generated,
+                        metadata,
+                        ..Default::default()
+                    });
+                }
+                (Some(generated), None) => {
+                    meta.output = Some(LLMOutputModel {
+                        raw_text: generated,
+                        ..Default::default()
+                    });
+                }
             }
+        } else {
+            println!("No metadata found for llm event");
         }
     }
 }
diff --git a/engine/baml-client-lib/src/otel/span_events/llm_request_error.rs b/engine/baml-client-lib/src/otel/span_events/llm_request_error.rs
@@ -71,6 +71,8 @@ where
                 traceback: event.traceback,
                 ..Default::default()
             });
+        } else {
+            println!("No metadata found for llm event");
         }
     }
 }