add test cases which show collector fails in a non-recoverable way whenever baml-runtime fails (bad type_builder, bad client_registry, bad args) (#1633)

<!-- ELLIPSIS_HIDDEN -->



> [!IMPORTANT]
> Add test cases to ensure collector fails non-recoverably on
baml-runtime errors, with async handling in `runtime_interface.rs`.
> 
>   - **Behavior**:
> - Refactor `call_function_impl` in `runtime_interface.rs` to handle
async execution and error handling.
> - Ensure function execution logs errors when BAML function does not
exist or argument validation fails.
>   - **Tests**:
> - Add `test_collector_failures_arg_type` to check for
`BamlInvalidArgumentError` on bad argument types.
> - Add `test_collector_failures_client_registry` to check for
`BamlError` when client registry is incorrect.
> - Add streaming variants `test_collector_failures_arg_type_streaming`
and `test_collector_failures_client_registry_streaming` for similar
checks.
>   - **Misc**:
>     - Format adjustments in `test_collector.py` for readability.
> 
> <sup>This description was created by </sup>[<img alt="Ellipsis"
src="https://img.shields.io/badge/Ellipsis-blue?color=175173">](https://www.ellipsis.dev?ref=BoundaryML%2Fbaml&utm_source=github&utm_medium=referral)<sup>
for 816aec8. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->

---------

Co-authored-by: Aaron Villalpando <aaron@boundaryml.com>

Author: hellovai

engine/baml-runtime/src/runtime/runtime_interface.rs

@@ -372,47 +372,8 @@ impl RuntimeInterface for InternalBamlRuntime {
         params: &BamlMap<String, BamlValue>,
         ctx: RuntimeContext,
     ) -> Result<crate::FunctionResult> {
-        let func = match self.get_function(&function_name, &ctx) {
-            Ok(func) => func,
-            Err(e) => {
-                return Ok(FunctionResult::new(
-                    OrchestrationScope::default(),
-                    LLMResponse::UserFailure(format!(
-                        "BAML function {function_name} does not exist in baml_src/ (did you typo it?): {:?}",
-                        e
-                    )),
-                    None,
-                ))
-            }
-        };
-        let baml_args = self.ir().check_function_params(
-            &func,
-            params,
-            ArgCoercer {
-                span_path: None,
-                allow_implicit_cast_to_string: false,
-            },
-        )?;
-        // let baml_args = match self.ir().check_function_params(
-        //     &func,
-        //     &params,
-        //     ArgCoercer {
-        //         span_path: None,
-        //         allow_implicit_cast_to_string: false,
-        //     },
-        // ) {
-        //     Ok(args) => args,
-        //     Err(e) => {
-        //         return Ok(FunctionResult::new(
-        //             OrchestrationScope::default(),
-        //             LLMResponse::UserFailure(format!(
-        //                 "Failed while validating args for {function_name}: {:?}",
-        //                 e
-        //             )),
-        //             None,
-        //         ))
-        //     }
-        // };
+        let local_span_id = ctx.span_id.clone();
+        let local_function_name = function_name.clone();
 
         if let Some(span_id) = ctx.span_id {
             let trace_event = TraceEvent {
@@ -443,15 +404,44 @@ impl RuntimeInterface for InternalBamlRuntime {
             );
         }
 
-        let renderer = PromptRenderer::from_function(&func, self.ir(), &ctx)?;
-        let orchestrator = self.orchestration_graph(renderer.client_spec(), &ctx)?;
+        let future = async {
+            let func = match self.get_function(&function_name, &ctx) {
+                Ok(func) => func,
+                Err(e) => {
+                return Ok(FunctionResult::new(
+                    OrchestrationScope::default(),
+                    LLMResponse::UserFailure(format!(
+                        "BAML function {function_name} does not exist in baml_src/ (did you typo it?): {:?}",
+                        e
+                    )),
+                    None,
+                ))
+                }
+            };
 
-        // Now actually execute the code.
-        let (history, _) =
-            orchestrate_call(orchestrator, self.ir(), &ctx, &renderer, &baml_args, |s| {
-                renderer.parse(self.ir(), &ctx, s, false)
-            })
-            .await;
+            let baml_args = self.ir().check_function_params(
+                &func,
+                params,
+                ArgCoercer {
+                    span_path: None,
+                    allow_implicit_cast_to_string: false,
+                },
+            )?;
+
+            let renderer = PromptRenderer::from_function(&func, self.ir(), &ctx)?;
+            let orchestrator = self.orchestration_graph(renderer.client_spec(), &ctx)?;
+
+            // Now actually execute the code.
+            let (history, _) =
+                orchestrate_call(orchestrator, self.ir(), &ctx, &renderer, &baml_args, |s| {
+                    renderer.parse(self.ir(), &ctx, s, false)
+                })
+                .await;
+
+            FunctionResult::new_chain(history)
+        };
+
+        let result = future.await;
 
         let end_time = web_time::SystemTime::now();
         if let Some(span_id) = ctx.span_id {
@@ -475,7 +465,7 @@ impl RuntimeInterface for InternalBamlRuntime {
             );
         }
 
-        FunctionResult::new_chain(history)
+        result
     }
 
     // Note that this only returns a FunctionResultStream object,

integ-tests/python/tests/test_collector.py

@@ -1,19 +1,22 @@
+from baml_py.errors import BamlInvalidArgumentError, BamlError, BamlClientError
 import pytest
 import dotenv
 from openai.types.chat import ChatCompletion
 
 from ..baml_client import b
 from ..baml_client.sync_client import b as b_sync
-from baml_py import Collector
+from baml_py import ClientRegistry, Collector
 import gc
 import sys
 import asyncio
 
 dotenv.load_dotenv()
 
+
 def function_span_count():
     return Collector.__function_span_count()  # type: ignore
 
+
 @pytest.fixture(autouse=True)
 def ensure_collector_is_empty():
     assert function_span_count() == 0
@@ -240,8 +243,12 @@ async def test_collector_async_multiple_calls_usage():
 
     # Capture usage after second call and verify it's the sum of both calls
     second_call_usage = function_logs[1].usage
-    total_input = (first_call_usage.input_tokens or 0) + (second_call_usage.input_tokens or 0)
-    total_output = (first_call_usage.output_tokens or 0) + (second_call_usage.output_tokens or 0)
+    total_input = (first_call_usage.input_tokens or 0) + (
+        second_call_usage.input_tokens or 0
+    )
+    total_output = (first_call_usage.output_tokens or 0) + (
+        second_call_usage.output_tokens or 0
+    )
     assert collector.usage.input_tokens == total_input
     assert collector.usage.output_tokens == total_output
 
@@ -286,11 +293,11 @@ async def test_collector_multiple_collectors():
 
     # Verify coll1 usage is now the sum of both calls
     usage_second_call_coll1 = logs1[1].usage
-    total_input = (
-        (usage_first_call_coll1.input_tokens or 0) + (usage_second_call_coll1.input_tokens or 0)
+    total_input = (usage_first_call_coll1.input_tokens or 0) + (
+        usage_second_call_coll1.input_tokens or 0
     )
-    total_output = (
-        (usage_first_call_coll1.output_tokens or 0) + (usage_second_call_coll1.output_tokens or 0)
+    total_output = (usage_first_call_coll1.output_tokens or 0) + (
+        usage_second_call_coll1.output_tokens or 0
     )
     assert coll1.usage.input_tokens == total_input
     assert coll1.usage.output_tokens == total_output
@@ -322,8 +329,12 @@ async def test_collector_mixed_async_sync_calls():
     # Verify the second call's usage
     usage_second_call = logs[1].usage
     assert logs[1].timing.start_time_utc_ms > logs[0].timing.start_time_utc_ms
-    total_input = (usage_first_call.input_tokens or 0) + (usage_second_call.input_tokens or 0)
-    total_output = (usage_first_call.output_tokens or 0) + (usage_second_call.output_tokens or 0)
+    total_input = (usage_first_call.input_tokens or 0) + (
+        usage_second_call.input_tokens or 0
+    )
+    total_output = (usage_first_call.output_tokens or 0) + (
+        usage_second_call.output_tokens or 0
+    )
     assert collector.usage.input_tokens == total_input
     assert collector.usage.output_tokens == total_output
 
@@ -361,3 +372,81 @@ async def test_collector_parallel_async_calls():
     # total_output = usage_call1.output_tokens + usage_call2.output_tokens
     # assert collector.usage.input_tokens == total_input
     # assert collector.usage.output_tokens == total_output
+
+
+@pytest.mark.asyncio
+async def test_collector_failures_arg_type():
+    collector = Collector(name="my-collector")
+    with pytest.raises(BamlInvalidArgumentError):
+        value: str = 124  # type: ignore (We want to test the error)
+        await b.TestOpenAIGPT4oMini(value, baml_options={"collector": collector})
+
+    assert len(collector.logs) == 1
+    last_log = collector.last
+    print("------------------------- last_log", last_log)
+    assert last_log is not None
+    assert last_log.function_name == "TestOpenAIGPT4oMini"
+
+
+@pytest.mark.asyncio
+async def test_collector_failures_client_registry():
+    collector = Collector(name="my-collector")
+    client_registry = ClientRegistry()
+    client_registry.set_primary("DoesNotExist")
+    with pytest.raises(BamlError):
+        await b.TestOpenAIGPT4oMini(
+            "hi there",
+            baml_options={"collector": collector, "client_registry": client_registry},
+        )
+    assert len(collector.logs) == 1
+    last_log = collector.last
+    assert last_log is not None
+    assert last_log.function_name == "TestOpenAIGPT4oMini"
+
+
+@pytest.mark.asyncio
+async def test_collector_failures_arg_type_streaming():
+    collector = Collector(name="my-collector")
+    with pytest.raises(BamlInvalidArgumentError):
+        value: str = 124  # type: ignore (We want to test the error)
+        async for _ in b.stream.TestOpenAIGPT4oMini(
+            value, baml_options={"collector": collector}
+        ):
+            pass
+
+    # Fails before the stream is even started
+    # We don't have a state for streams that were "registered" but not started
+    assert len(collector.logs) == 0
+
+
+@pytest.mark.asyncio
+async def test_collector_failures_client_registry_streaming():
+    collector = Collector(name="my-collector")
+    client_registry = ClientRegistry()
+    client_registry.add_llm_client(
+        "TestClient",
+        "openai",
+        {"model": "gpt-4o-mini", "base_url": "https://does-not-exist.com"},
+    )
+    client_registry.set_primary("TestClient")
+    with pytest.raises(BamlClientError):
+        try:
+            stream = b.stream.TestOpenAIGPT4oMini(
+                "hi there",
+                baml_options={
+                    "collector": collector,
+                    "client_registry": client_registry,
+                },
+            )
+            # TODO: baml doesnt yet throw if theres a connection error during the stream..
+            async for _ in stream:
+                pass
+            # So we try to call get final response to make sure it fails
+            await stream.get_final_response()
+        except Exception as e:
+            print(f"Error occurred: {e}")
+            raise
+    assert len(collector.logs) == 1
+    last_log = collector.last
+    assert last_log is not None
+    assert last_log.function_name == "TestOpenAIGPT4oMini"