Support thinking models from anthropic. (#1555)

<!-- ELLIPSIS_HIDDEN -->



> [!IMPORTANT]
> Add support for Anthropic thinking models by introducing
`TestThinking` function and `CustomStory` model, updating response
handling and integration tests.
> 
>   - **Behavior**:
> - Add `TestThinking` function to handle thinking events, returning
`CustomStory`.
> - Update `parse_anthropic_response` in `response_handler.rs` to handle
multiple content types.
>   - **Models**:
> - Add `CustomStory` model to `types.ts`, `types.py`, and `types.rb`.
> - Update `AnthropicMessageContent` enum in `types.rs` to include new
content types.
>   - **Integration Tests**:
>     - Add `TestThinking` to `clients.baml` and `anthropic.baml`.
>     - Add tests for `TestThinking` in `test_functions.py`.
>   - **Misc**:
> - Remove `Drop` implementation from `FunctionResultStream` in
`stream.rs`.
>     - Update `openapi.yaml` to include `TestThinking` endpoint.
> 
> <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 ae0e94f. It will automatically
update as commits are pushed.</sup>

<!-- ELLIPSIS_HIDDEN -->

Author: hellovai

engine/baml-runtime/src/internal/llm_client/primitive/anthropic/response_handler.rs

@@ -1,8 +1,11 @@
 use anyhow::Result;
 use baml_types::BamlMap;
 
-use super::types::{AnthropicMessageResponse, MessageChunk, StopReason};
-use crate::internal::llm_client::{primitive::request::RequestBuilder, traits::WithClient, ErrorCode, LLMCompleteResponse, LLMCompleteResponseMetadata, LLMErrorResponse, LLMResponse};
+use super::types::{AnthropicMessageContent, AnthropicMessageResponse, MessageChunk, StopReason};
+use crate::internal::llm_client::{
+    primitive::request::RequestBuilder, traits::WithClient, ErrorCode, LLMCompleteResponse,
+    LLMCompleteResponseMetadata, LLMErrorResponse, LLMResponse,
+};
 use anyhow::Context;
 use serde::Deserialize;
 use serde_json::Value;
@@ -24,11 +27,13 @@ pub fn parse_anthropic_response<C: WithClient + RequestBuilder>(
     instant_now: web_time::Instant,
     model_name: Option<String>,
 ) -> LLMResponse {
-    let response = match AnthropicMessageResponse::deserialize(&response_body).context(format!(
-        "Failed to parse into a response accepted by {}: {}",
-        std::any::type_name::<AnthropicMessageResponse>(),
-        response_body
-    )).map_err(|e| LLMErrorResponse {
+    let response = match AnthropicMessageResponse::deserialize(&response_body)
+        .context(format!(
+            "Failed to parse into a response accepted by {}: {}",
+            std::any::type_name::<AnthropicMessageResponse>(),
+            response_body
+        ))
+        .map_err(|e| LLMErrorResponse {
             client: client.context().name.to_string(),
             model: model_name.clone(),
             prompt: to_prompt(prompt),
@@ -37,33 +42,41 @@ pub fn parse_anthropic_response<C: WithClient + RequestBuilder>(
             latency: instant_now.elapsed(),
             message: format!("{:?}", e),
             code: ErrorCode::Other(2),
-        })
-    {
+        }) {
         Ok(response) => response,
         Err(e) => return LLMResponse::LLMFailure(e),
     };
 
+    println!("response: {:?}", response.content);
+
+    let content = response
+        .content
+        .iter()
+        .filter_map(|v| match v {
+            AnthropicMessageContent::Text { text } => Some(text),
+            _ => None,
+        })
+        .next();
 
-    if response.content.len() != 1 {
+    let content = if let Some(content) = content {
+        content
+    } else {
         return LLMResponse::LLMFailure(LLMErrorResponse {
             client: client.context().name.to_string(),
             model: model_name.clone(),
             prompt: to_prompt(prompt),
             start_time: system_now,
             request_options: client.request_options().clone(),
             latency: instant_now.elapsed(),
-            message: format!(
-                "Expected exactly one content block, got {}",
-                response.content.len()
-            ),
-            code: ErrorCode::Other(200),
+            message: "Anthropic response contains no text".to_string(),
+            code: ErrorCode::Other(2),
         });
-    }
+    };
 
     LLMResponse::Success(LLMCompleteResponse {
         client: client.context().name.to_string(),
         prompt: to_prompt(prompt),
-        content: response.content[0].text.clone(),
+        content: content.to_string(),
         start_time: system_now,
         latency: instant_now.elapsed(),
         request_options: client.request_options().clone(),
@@ -84,7 +97,6 @@ pub fn parse_anthropic_response<C: WithClient + RequestBuilder>(
     })
 }
 
-
 pub fn scan_anthropic_response_stream(
     client_name: &str,
     request_options: &BamlMap<String, serde_json::Value>,
@@ -98,14 +110,16 @@ pub fn scan_anthropic_response_stream(
     let inner = match accumulated {
         Ok(accumulated) => accumulated,
         // We'll just keep the first error and return it
-        Err(e) => return Ok(())
+        Err(e) => return Ok(()),
     };
 
-    let event = match MessageChunk::deserialize(&event_body).context(format!(
-        "Failed to parse into a response accepted by {}: {}",
-        std::any::type_name::<MessageChunk>(),
-        event_body
-    )).map_err(|e| LLMErrorResponse {
+    let event = match MessageChunk::deserialize(&event_body)
+        .context(format!(
+            "Failed to parse into a response accepted by {}: {}",
+            std::any::type_name::<MessageChunk>(),
+            event_body
+        ))
+        .map_err(|e| LLMErrorResponse {
             client: client_name.to_string(),
             model: model_name.clone(),
             prompt: prompt.clone(),
@@ -114,8 +128,7 @@ pub fn scan_anthropic_response_stream(
             latency: instant_now.elapsed(),
             message: format!("{:?}", e),
             code: ErrorCode::Other(2),
-        })
-    {
+        }) {
         Ok(response) => response,
         Err(e) => return Err(e),
     };
@@ -128,15 +141,18 @@ pub fn scan_anthropic_response_stream(
                 body.stop_reason,
                 Some(StopReason::StopSequence) | Some(StopReason::EndTurn)
             );
-            inner.finish_reason =
-                body.stop_reason.as_ref().map(ToString::to_string);
+            inner.finish_reason = body.stop_reason.as_ref().map(ToString::to_string);
             inner.prompt_tokens = Some(body.usage.input_tokens);
             inner.output_tokens = Some(body.usage.output_tokens);
-            inner.total_tokens =
-                Some(body.usage.input_tokens + body.usage.output_tokens);
+            inner.total_tokens = Some(body.usage.input_tokens + body.usage.output_tokens);
         }
         MessageChunk::ContentBlockDelta(event) => {
-            inner.content += &event.delta.text;
+            match event.delta {
+                super::types::ContentBlockDelta::TextDelta { text } => {
+                    inner.content += &text;
+                }
+                _ => (),
+            }
         }
         MessageChunk::ContentBlockStart(_) => (),
         MessageChunk::ContentBlockStop(_) => (),
@@ -154,25 +170,22 @@ pub fn scan_anthropic_response_stream(
                 .as_ref()
                 .map(|r| serde_json::to_string(r).unwrap_or("".into()));
             inner.output_tokens = Some(body.usage.output_tokens);
-            inner.total_tokens = Some(
-                inner.prompt_tokens.unwrap_or(0) + body.usage.output_tokens,
-            );
+            inner.total_tokens = Some(inner.prompt_tokens.unwrap_or(0) + body.usage.output_tokens);
         }
         MessageChunk::MessageStop => (),
         MessageChunk::Error(err) => {
-            return Err(
-                LLMErrorResponse {
-                    client: client_name.to_string(),
-                    model: model_name.clone(),
-                    prompt: prompt.clone(),
-                    request_options: request_options.clone(),
-                    start_time: system_now.clone(),
-                    latency: instant_now.elapsed(),
-                    message: err.message,
-                    code: ErrorCode::Other(2),
-                }
-            );
+            return Err(LLMErrorResponse {
+                client: client_name.to_string(),
+                model: model_name.clone(),
+                prompt: prompt.clone(),
+                request_options: request_options.clone(),
+                start_time: system_now.clone(),
+                latency: instant_now.elapsed(),
+                message: err.message,
+                code: ErrorCode::Other(2),
+            });
         }
+        MessageChunk::Other => (),
     };
 
     inner.latency = instant_now.elapsed();

engine/baml-runtime/src/internal/llm_client/primitive/anthropic/types.rs

@@ -2,9 +2,30 @@ use serde::{Deserialize, Serialize};
 
 // https://docs.anthropic.com/claude/reference/messages_post
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
-pub struct AnthropicMessageContent {
-    pub r#type: String,
-    pub text: String,
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum AnthropicMessageContent {
+    // type: text
+    Text {
+        text: String,
+    },
+    // // type: tool_use
+    ToolUse {
+        id: Option<String>,
+        input: serde_json::Value,
+        name: String,
+    },
+    // // type: thinking
+    // Thinking {
+    //     signature: Option<String>,
+    //     thinking: String,
+    // },
+    // type: redacted_thinking
+    RedactedThinking {
+        data: String,
+    },
+    // fallback for unknown types
+    #[serde(other)]
+    Other,
 }
 
 #[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
@@ -98,6 +119,9 @@ pub enum MessageChunk {
     /// Message stop chunk.
     MessageStop,
     Error(AnthropicErrorInner),
+    /// Fallback for unknown types
+    #[serde(other)]
+    Other,
 }
 
 /// The message start chunk.
@@ -122,7 +146,16 @@ pub struct ContentBlockDeltaChunk {
     /// The index.
     pub index: u32,
     /// The text delta content block.
-    pub delta: TextDeltaContentBlock,
+    pub delta: ContentBlockDelta,
+}
+
+#[derive(Debug, Deserialize, Clone, PartialEq, Serialize)]
+#[serde(tag = "type", rename_all = "snake_case")]
+pub enum ContentBlockDelta {
+    TextDelta { text: String },
+    SignatureDelta { signature: String },
+    ThinkingDelta { thinking: String },
+    Other,
 }
 
 /// The content block stop chunk.
@@ -178,8 +211,7 @@ mod tests {
 
         let chunk = MessageChunk::ContentBlockDelta(ContentBlockDeltaChunk {
             index: 0,
-            delta: TextDeltaContentBlock {
-                _type: ContentType::TextDelta,
+            delta: ContentBlockDelta::TextDelta {
                 text: "Hello".to_string(),
             },
         });

engine/baml-runtime/src/types/stream.rs

@@ -37,12 +37,6 @@ pub struct FunctionResultStream {
     pub(crate) collectors: Vec<Arc<Collector>>,
 }
 
-impl Drop for FunctionResultStream {
-    fn drop(&mut self) {
-        log::info!("Dropping FunctionResultStream: {}", self.function_name);
-    }
-}
-
 #[cfg(target_arch = "wasm32")]
 // JsFuture is !Send, so when building for WASM, we have to drop that requirement from StreamCallback
 static_assertions::assert_impl_all!(FunctionResultStream: Send);

integ-tests/baml_src/clients.baml

@@ -265,6 +265,20 @@ client<llm> Sonnet {
   }
 }
 
+
+client<llm> SonnetThinking {
+  provider anthropic
+  options {
+    model "claude-3-7-sonnet-20250219"
+    api_key env.ANTHROPIC_API_KEY
+    max_tokens 2048
+    thinking {
+      type "enabled"
+      budget_tokens 1024
+    }
+  }
+}
+
 client<llm> Claude {
   provider anthropic
   options {

integ-tests/baml_src/test-files/providers/anthropic.baml

@@ -25,4 +25,22 @@ function TestCaching(input: string, not_cached: string) -> string {
     {{ _.role('user') }}
     {{ not_cached }}
   "#
+}
+
+class CustomStory {
+  title string
+  characters string[]
+  content string
+}
+
+function TestThinking(input: string) -> CustomStory {
+  client SonnetThinking
+  prompt #"
+    {{ _.role('system') }}
+    Generate the following story
+    {{ ctx.output_format }}
+
+    {{ _.role('user') }}
+    {{ input }}
+  "#
 }

integ-tests/openapi/baml_client/openapi.yaml

@@ -2146,6 +2146,19 @@ paths:
                 title: TestSingleFallbackClientResponse
                 type: string
       operationId: TestSingleFallbackClient
+  /call/TestThinking:
+    post:
+      requestBody:
+        $ref: '#/components/requestBodies/TestThinking'
+      responses:
+        '200':
+          description: Successful operation
+          content:
+            application/json:
+              schema:
+                title: TestThinkingResponse
+                $ref: '#/components/schemas/CustomStory'
+      operationId: TestThinking
   /call/TestUniverseQuestion:
     post:
       requestBody:
@@ -4797,6 +4810,22 @@ components:
                 $ref: '#/components/schemas/BamlOptions'
             required: []
             additionalProperties: false
+    TestThinking:
+      required: true
+      content:
+        application/json:
+          schema:
+            title: TestThinkingRequest
+            type: object
+            properties:
+              input:
+                type: string
+              __baml_options__:
+                nullable: true
+                $ref: '#/components/schemas/BamlOptions'
+            required:
+            - input
+            additionalProperties: false
     TestUniverseQuestion:
       required: true
       content:
@@ -5239,6 +5268,22 @@ components:
       required:
       - primary
       additionalProperties: false
+    CustomStory:
+      type: object
+      properties:
+        characters:
+          type: array
+          items:
+            type: string
+        content:
+          type: string
+        title:
+          type: string
+      required:
+      - characters
+      - content
+      - title
+      additionalProperties: false
     CustomTaskResult:
       type: object
       properties: