Fix Tracing bug and LSP Crash (#2030)

Fix an issue where BAML_LOG wouldn't be printed to the console in async
contexts in python due to the context stack not being propagated
correctly.

Log an error when we run into issues with the tracing spans

Fix an issue where LSP would crash when sending diagnostics for a baml
file not in a baml_src.

Author: aaronvg

engine/baml-lib/baml/tests/validation_files/prompt_fiddle_example.baml

@@ -199,7 +199,7 @@ client<llm> Claude {
 client<llm> Gemini {
   provider google-ai
   options {
-    model "gemini-1.5-pro-001"
+    model "gemini-1.5-pro"
     api_key env.GOOGLE_API_KEY
   }
 }

engine/baml-runtime/src/lib.rs

@@ -591,7 +591,7 @@ impl BamlRuntime {
                 .finish_call(call, ctx, None)
             {
                 Ok(id) => {}
-                Err(e) => log::debug!("Error during logging: {}", e),
+                Err(e) => baml_log::error!("Error during logging: {}", e),
             }
             #[cfg(target_arch = "wasm32")]
             match self
@@ -601,7 +601,7 @@ impl BamlRuntime {
                 .await
             {
                 Ok(id) => {}
-                Err(e) => log::debug!("Error during logging: {}", e),
+                Err(e) => log::error!("Error during logging: {}", e),
             }
         }
 
@@ -689,6 +689,7 @@ impl BamlRuntime {
 
         log::trace!("Calling function: {}", function_name);
         log::debug!("collectors: {:#?}", &collectors);
+
         let call = self
             .tracer_wrapper
             .get_or_create_tracer(&env_vars)
@@ -854,7 +855,7 @@ impl BamlRuntime {
             .finish_baml_call(call, ctx, &response)
         {
             Ok(id) => {}
-            Err(e) => log::debug!("Error during logging: {}", e),
+            Err(e) => baml_log::error!("Error during logging: {}", e),
         }
         #[cfg(target_arch = "wasm32")]
         match self
@@ -864,7 +865,7 @@ impl BamlRuntime {
             .await
         {
             Ok(id) => {}
-            Err(e) => log::debug!("Error during logging: {}", e),
+            Err(e) => log::error!("Error during logging: {}", e),
         }
 
         (response, curr_call_id)

engine/baml-runtime/src/tracing/mod.rs

@@ -49,6 +49,7 @@ cfg_if! {
 #[derive(Debug, Clone)]
 pub struct TracingCall {
     pub call_id: Uuid,
+    pub function_name: String,
     pub new_call_id_stack: Vec<baml_ids::FunctionCallId>,
     params: BamlMap<String, BamlValue>,
     start_time: web_time::SystemTime,
@@ -406,16 +407,25 @@ impl BamlTracer {
         self.trace_stats.guard().start();
         let (call_id, call_stack, last_tags, global_tags) = ctx.enter(function_name);
 
-        log::trace!(" Entering call {:#?} in {:?}", call_id, function_name);
+        log::trace!(
+            "\n{}------------------- Entering {:?}, ctx chain {:#?}",
+            "        ".repeat(ctx.context_depth()),
+            function_name,
+            ctx
+        );
+
         let call = TracingCall {
             call_id,
+            function_name: function_name.to_string(),
             new_call_id_stack: call_stack.clone(),
             params: params.clone(),
             start_time: web_time::SystemTime::now(),
             // Note these tags are the ones currently on the stack. While the function runs we may register
             // more tags with set_tags(). Those are picked up via a diff event (SetTags)
             tags: last_tags.clone(),
         };
+        // println!("---- {} ctx {:#?}", function_name, ctx);
+        // baml_log::info!("---- {} ctx {:#?}", function_name, ctx);
 
         // This must happen before the first event is sent.
         if let Some(collectors) = collectors {
@@ -506,8 +516,9 @@ impl BamlTracer {
             );
         };
         log::trace!(
-            "Finishing call: {:#?} {}\nevent chain {:?}",
-            call,
+            "\n{}------------------- Finishing call: {:#?} {}\nevent chain {:#?}",
+            "        ".repeat(ctx.context_depth()),
+            call.function_name,
             call_id,
             event_chain
         );
@@ -626,8 +637,8 @@ impl BamlTracer {
         };
 
         log::trace!(
-            "Finishing baml call: {:#?} {}\nevent chain {:?}",
-            call,
+            "Finishing baml call: {:#?} {}\nevent chain {:#?}",
+            call.function_name,
             call_id,
             event_chain
         );

engine/baml-runtime/src/tracingv2/publisher/publisher.rs

@@ -469,10 +469,10 @@ impl TracePublisher {
                 Ok(())
             }
             Err(e) => {
-                baml_log::error!("Failed to upload batch of {} events: {}", batch.len(), e);
+                log::info!("Failed to upload batch of {} events: {}", batch.len(), e);
                 // If batch size is at or below minimum, give up
                 if batch.len() <= min_batch_size {
-                    baml_log::error!(
+                    log::info!(
                         "Failed to upload single/minimum batch of {} events: {}",
                         batch.len(),
                         e
@@ -507,15 +507,15 @@ impl TracePublisher {
                         Ok(())
                     }
                     (Err(e1), Ok(())) => {
-                        baml_log::error!("First half failed: {}", e1);
+                        log::info!("First half failed: {}", e1);
                         Err(e1)
                     }
                     (Ok(()), Err(e2)) => {
-                        baml_log::error!("Second half failed: {}", e2);
+                        log::info!("Second half failed: {}", e2);
                         Err(e2)
                     }
                     (Err(e1), Err(e2)) => {
-                        baml_log::debug!("Both halves failed - first: {}, second: {}", e1, e2);
+                        log::debug!("Both halves failed - first: {}, second: {}", e1, e2);
                         Err(e1) // Return the first error
                     }
                 }
@@ -574,7 +574,7 @@ impl TracePublisher {
         {
             Ok(response) => response,
             Err(e) => {
-                baml_log::debug!("Failed to upload trace events: {}", e);
+                log::debug!("Failed to upload trace events: {}", e);
                 return Err(e.into());
             }
         };

engine/baml-runtime/src/tracingv2/storage/storage.rs

@@ -264,7 +264,6 @@ fn build_function_log(
     // Build each LLMCall or LLMStreamCall
     let mut calls = Vec::new();
     for (_rid, call_acc) in calls_map {
-        println!("### _rid: {:?}", _rid);
         let (client, provider) = parse_llm_client_and_provider(call_acc.llm_request.as_ref());
         let start_t = call_acc.timestamp_first_seen.unwrap_or(start_ms);
         let end_t = call_acc.timestamp_last_seen.unwrap_or(start_t);

engine/language_client_codegen/src/python/templates/async_client.py.j2

@@ -4,6 +4,7 @@ from typing_extensions import Literal
 import baml_py
 
 from . import _baml
+from ._baml import BamlCallOptions
 from .types import Checked, Check
 from .parser import LlmResponseParser, LlmStreamParser
 from .async_request import AsyncHttpRequest, AsyncHttpStreamRequest
@@ -104,7 +105,7 @@ class BamlAsyncClient:
           "{{name}}": {{name}},
           {%- endfor %}
         },
-        self.__ctx_manager.get(),
+        self.__ctx_manager.clone_context(),
         tb,
         __cr__,
         collectors,
@@ -167,4 +168,4 @@ class BamlStreamClient:
 
 b = BamlAsyncClient(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME, DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX)
 
-__all__ = ["b"]
+__all__ = ["b", "BamlCallOptions"]

engine/language_client_codegen/src/python/templates/sync_client.py.j2

@@ -4,6 +4,7 @@ from typing_extensions import Literal
 import baml_py
 
 from . import _baml
+from ._baml import BamlCallOptions
 from .types import Checked, Check
 from .parser import LlmResponseParser, LlmStreamParser
 from .sync_request import HttpRequest, HttpStreamRequest
@@ -166,4 +167,4 @@ class BamlStreamClient:
 
 b = BamlSyncClient(DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_RUNTIME, DO_NOT_USE_DIRECTLY_UNLESS_YOU_KNOW_WHAT_YOURE_DOING_CTX)
 
-__all__ = ["b"]
+__all__ = ["b", "BamlCallOptions"]

engine/language_client_python/python_src/baml_py/ctx_manager.py

@@ -10,6 +10,7 @@
 from .baml_py import BamlLogEvent, RuntimeContextManager, BamlRuntime, BamlSpan
 import atexit
 import threading
+from typing import Dict
 
 F = typing.TypeVar("F", bound=typing.Callable[..., typing.Any])
 
@@ -26,11 +27,28 @@ def current_thread_id() -> int:
     return current_thread.ident or 0
 
 
+prev_ctx_manager: typing.Optional["CtxManager"] = None
+
+
 class CtxManager:
+    def __new__(cls, *args, **kwargs):
+        if prev_ctx_manager is not None:
+            return prev_ctx_manager
+        return super().__new__(cls)
+
     def __init__(self, rt: BamlRuntime):
+        global prev_ctx_manager
+        if prev_ctx_manager is not None:
+            self.rt = prev_ctx_manager.rt
+            self.ctx = prev_ctx_manager.ctx
+            return
+
+        prev_ctx_manager = self
+
         self.rt = rt
+
         self.ctx = contextvars.ContextVar[typing.Dict[int, RuntimeContextManager]](
-            "baml_ctx", default={current_thread_id(): rt.create_context_manager()}
+            "baml_ctx", default={}
         )
         atexit.register(self.rt.flush)
 
@@ -71,20 +89,35 @@ def get(self) -> RuntimeContextManager:
         return self.__ctx()
 
     def start_trace_sync(
-        self, name: str, args: typing.Dict[str, typing.Any], env_vars: typing.Dict[str, str]
+        self,
+        name: str,
+        args: typing.Dict[str, typing.Any],
+        env_vars: typing.Dict[str, str],
     ) -> BamlSpan:
+        # Clone the current context before creating the span
         mng = self.__ctx()
         return BamlSpan.new(self.rt, name, args, mng, env_vars)
 
     def start_trace_async(
-        self, name: str, args: typing.Dict[str, typing.Any], env_vars: typing.Dict[str, str]
+        self,
+        name: str,
+        args: typing.Dict[str, typing.Any],
+        env_vars: typing.Dict[str, str],
     ) -> BamlSpan:
         mng = self.__ctx()
         cln = mng.deep_clone()
         self.ctx.set({current_thread_id(): cln})
         return BamlSpan.new(self.rt, name, args, cln, env_vars)
 
-    def end_trace(self, span: BamlSpan, response: typing.Any, env_vars: typing.Dict[str, str]) -> None:
+    def clone_context(self) -> RuntimeContextManager:
+        mng = self.__ctx()
+        cln = mng.deep_clone()
+        self.ctx.set({current_thread_id(): cln})
+        return cln
+
+    def end_trace(
+        self, span: BamlSpan, response: typing.Any, env_vars: typing.Dict[str, str]
+    ) -> None:
         span.finish(response, self.__ctx(), env_vars)
 
     def flush(self) -> None:

engine/language_server/src/server/api/diagnostics.rs

@@ -87,6 +87,9 @@ pub fn publish_session_lsp_diagnostics(
 ) -> Result<()> {
     // let keys = session.index().documents.keys();
     let path = file_url.to_file_path().unwrap_or(PathBuf::new());
+    if !file_url.to_string().contains("baml_src") {
+        return Ok(());
+    }
     let project = session
         .get_or_create_project(&path)
         .expect("We just ensured the session is valid");

integ-tests/baml_src/clients.baml

@@ -139,11 +139,11 @@ client<llm> GPT35AzureFailed {
     api_key env.AZURE_OPENAI_API_KEY
   }
 }
-
+  
 client<llm> Gemini {
   provider google-ai
   options {
-    model gemini-1.5-pro-001
+    model gemini-1.5-pro
     api_key env.GOOGLE_API_KEY
     safetySettings {
       category HARM_CATEGORY_HATE_SPEECH