Add HTTPBody to HTTPResponse (#1684)

<!-- ELLIPSIS_HIDDEN -->


> [!IMPORTANT]
> Replaces `serde_json::Value` with `HTTPBody` for HTTP response bodies
across multiple components, updating handling, logging, and tests.
> 
>   - **Behavior**:
> - Replaces `serde_json::Value` with `HTTPBody` for `body` in
`HTTPResponse` in `events.rs` and `request.rs`.
> - Updates `log_http_response()` and `execute_request()` in
`request.rs` to use `HTTPBody`.
> - Modifies `HTTPResponse` class in Python, Ruby, and TypeScript to use
`HTTPBody`.
>   - **Tests**:
> - Updates Python, Ruby, and TypeScript integration tests to handle
`HTTPBody` in request and response validation.
> - Ensures tests verify JSON parsing and content extraction from
`HTTPBody`.
> 
> <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 edc4312. It will automatically
update as commits are pushed.</sup>


<!-- ELLIPSIS_HIDDEN -->

Author: antoniosarosi

engine/baml-lib/baml-types/src/tracing/events.rs

@@ -197,7 +197,7 @@ pub struct HTTPResponse {
     pub request_id: HttpRequestId,
     pub status: u16,
     pub headers: serde_json::Value,
-    pub body: serde_json::Value,
+    pub body: HTTPBody,
 }
 
 #[derive(Debug, Serialize, Deserialize)]

engine/baml-runtime/src/internal/llm_client/primitive/request.rs

@@ -126,7 +126,7 @@ async fn log_http_response(
     http_request_id: HttpRequestId,
     status: u16,
     headers: serde_json::Value,
-    body: serde_json::Value,
+    body: HTTPBody,
 ) {
     if let Some(span_id) = runtime_context.span_id {
         BAML_TRACER.lock().unwrap().put(Arc::new(TraceEvent {
@@ -253,7 +253,7 @@ pub async fn execute_request(
                     .unwrap_or(reqwest::StatusCode::INTERNAL_SERVER_ERROR)
                     .as_u16(),
                 serde_json::Value::Null,
-                serde_json::Value::String(format!("No response. Error: {:?}", e)),
+                HTTPBody::new(format!("No response. Error: {:?}", e).into_bytes()),
             )
             .await;
 
@@ -295,7 +295,7 @@ pub async fn execute_request(
                     http_request_id.clone(),
                     0,
                     serde_json::Value::Null,
-                    serde_json::Value::String(format!("Could not read response body: {:?}", e)),
+                    HTTPBody::new(format!("Could not read response body: {:?}", e).into_bytes()),
                 )
                 .await;
                 return Err(LLMResponse::LLMFailure(LLMErrorResponse {
@@ -317,15 +317,17 @@ pub async fn execute_request(
             Ok(s) if !s.is_empty() => s.to_string(),
             _ => "<no response or invalid utf-8>".to_string(),
         };
+
         log_http_response(
             runtime_context,
             TraceLevel::Error,
             http_request_id.clone(),
             logged_res.status.as_u16(),
             json_headers(&logged_res.headers),
-            serde_json::Value::String(resp_body.clone()),
+            HTTPBody::new(resp_body.clone().into_bytes()),
         )
         .await;
+
         return Err(LLMResponse::LLMFailure(LLMErrorResponse {
             client: client.context().name.to_string(),
             model: None,
@@ -335,8 +337,7 @@ pub async fn execute_request(
             latency: instant_now.elapsed(),
             message: format!(
                 "Request failed with status code: {}, \n{}",
-                logged_res.status,
-                resp_body.clone()
+                logged_res.status, resp_body
             ),
             code: ErrorCode::from_status(logged_res.status),
         }));
@@ -352,7 +353,7 @@ pub async fn execute_request(
                     http_request_id.clone(),
                     0,
                     serde_json::Value::Null,
-                    serde_json::Value::String(format!("Could not read response body: {:?}", e)),
+                    HTTPBody::new(format!("Could not read response body: {:?}", e).into_bytes()),
                 )
                 .await;
                 return Err(LLMResponse::LLMFailure(LLMErrorResponse {
@@ -380,7 +381,7 @@ pub async fn execute_request(
             http_request_id.clone(),
             logged_response.status.as_u16(),
             json_headers(&logged_response.headers),
-            json_body(JsonBodyInput::String(resp_body)).unwrap_or_default(),
+            HTTPBody::new(resp_body.into_bytes()),
         )
         .await;
 

engine/language_client_python/python_src/baml_py/baml_py.pyi

@@ -174,6 +174,19 @@ class BamlRuntime:
         cr: Optional[ClientRegistry],
         is_stream: bool,
     ) -> HTTPRequest: ...
+    def parse_llm_response(
+        self,
+        function_name: str,
+        llm_response: str,
+        enum_module: Any,
+        cls_module: Any,
+        partial_cls_module: Any,
+        allow_partials: bool,
+        ctx: RuntimeContextManager,
+        tb: Optional[TypeBuilder],
+        cr: Optional[ClientRegistry],
+    ) -> Any: ...
+
 
 class LogEventMetadata:
     event_id: str
@@ -354,7 +367,7 @@ class HTTPResponse:
     @property
     def headers(self) -> Dict[str, Any]: ...
     @property
-    def body(self) -> Union[Dict[str, Any], str]: ...
+    def body(self) -> HTTPBody: ...
 
 class ClientRegistry:
     def __init__(self) -> None: ...

engine/language_client_python/src/types/response.rs

@@ -1,10 +1,10 @@
 use pyo3::{
     prelude::{pymethods, Py},
     types::{PyDict, PyDictMethods},
-    PyObject, PyResult, Python,
+    PyResult, Python,
 };
 
-use super::log_collector::serde_value_to_py;
+use super::request::HTTPBody;
 
 crate::lang_wrapper!(
     HTTPResponse,
@@ -20,7 +20,7 @@ impl HTTPResponse {
             "HTTPResponse(status={}, headers={}, body={})",
             self.inner.status,
             serde_json::to_string_pretty(&self.inner.headers).unwrap(),
-            serde_json::to_string_pretty(&self.inner.body).unwrap()
+            serde_json::to_string_pretty(&self.inner.body.as_serde_value()).unwrap()
         )
     }
 
@@ -40,9 +40,9 @@ impl HTTPResponse {
         Ok(dict.into())
     }
 
-    // note the body may be an error string, not a dict
     #[getter]
-    pub fn body(&self, py: Python<'_>) -> PyResult<PyObject> {
-        serde_value_to_py(py, &self.inner.body)
+    pub fn body(&self) -> HTTPBody {
+        // TODO: Avoid clone.
+        HTTPBody::from(self.inner.body.clone())
     }
 }

engine/language_client_ruby/ext/ruby_ffi/src/types/response.rs

@@ -1,6 +1,8 @@
 use crate::Result;
 use magnus::{class, method, Error, Module, RModule, Ruby};
 
+use super::request::HTTPBody;
+
 crate::lang_wrapper!(
     HTTPResponse,
     "Baml::Ffi::HTTPResponse",
@@ -14,7 +16,7 @@ impl HTTPResponse {
             "HTTPResponse(status={}, headers={}, body={})",
             self.inner.status,
             serde_json::to_string_pretty(&self.inner.headers).unwrap_or_default(),
-            serde_json::to_string_pretty(&self.inner.body).unwrap_or_default()
+            serde_json::to_string_pretty(&self.inner.body.as_serde_value()).unwrap_or_default()
         )
     }
 
@@ -28,9 +30,9 @@ impl HTTPResponse {
             .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{:?}", e)))
     }
 
-    pub fn body(ruby: &Ruby, rb_self: &Self) -> Result<magnus::Value> {
-        serde_magnus::serialize(&rb_self.inner.body)
-            .map_err(|e| Error::new(ruby.exception_runtime_error(), format!("{:?}", e)))
+    pub fn body(&self) -> HTTPBody {
+        // TODO: Avoid clone.
+        HTTPBody::from(self.inner.body.clone())
     }
 
     pub fn define_in_ruby(module: &RModule) -> Result<()> {

engine/language_client_typescript/native.d.ts

@@ -128,7 +128,7 @@ export declare class HttpResponse {
   toString(): string
   get status(): number
   get headers(): object
-  get body(): any
+  get body(): HttpBody
 }
 export type HTTPResponse = HttpResponse
 

engine/language_client_typescript/src/types/response.rs

@@ -1,6 +1,8 @@
 use napi::{bindgen_prelude::Env, JsObject};
 use napi_derive::napi;
 
+use super::request::HTTPBody;
+
 crate::lang_wrapper!(
     HTTPResponse,
     baml_types::tracing::events::HTTPResponse,
@@ -15,7 +17,7 @@ impl HTTPResponse {
             "HTTPResponse(status={}, headers={}, body={})",
             self.inner.status,
             serde_json::to_string_pretty(&self.inner.headers).unwrap(),
-            serde_json::to_string_pretty(&self.inner.body).unwrap()
+            serde_json::to_string_pretty(&self.inner.body.as_serde_value()).unwrap()
         )
     }
 
@@ -37,7 +39,8 @@ impl HTTPResponse {
     }
 
     #[napi(getter)]
-    pub fn body(&self) -> napi::Result<serde_json::Value> {
-        Ok(self.inner.body.clone())
+    pub fn body(&self) -> HTTPBody {
+        // TODO: Avoid clone.
+        HTTPBody::from(self.inner.body.clone())
     }
 }

integ-tests/python/tests/test_collector.py

@@ -82,14 +82,15 @@ async def test_collector_async_no_stream_success():
     # Verify http response
     response = call.http_response
     assert response is not None
+    response_body = response.body.json()
     assert response.status == 200
-    assert response.body is not None
-    assert isinstance(response.body, dict)
-    completion = ChatCompletion(**response.body)
-    assert "choices" in response.body
-    assert len(response.body["choices"]) > 0
-    assert "message" in response.body["choices"][0]
-    assert "content" in response.body["choices"][0]["message"]
+    assert response_body is not None
+    assert isinstance(response_body, dict)
+    completion = ChatCompletion(**response_body)
+    assert "choices" in response_body
+    assert len(response_body["choices"]) > 0
+    assert "message" in response_body["choices"][0]
+    assert "content" in response_body["choices"][0]["message"]
     assert completion.choices[0].message.content is not None
 
     # Verify call timing

integ-tests/ruby/test_collector.rb

@@ -17,7 +17,7 @@
     assert_equal 0, Baml::Collector.__function_span_count if Baml::Collector.respond_to?(:__function_span_count)
   end
 
-  after do 
+  after do
     # Force garbage collection and check collector is empty
     GC.start
     # GC.start(full_mark: true, immediate_sweep: true);
@@ -35,7 +35,7 @@
     b.TestOpenAIGPT4oMini(input: "hi there", baml_options: {collector: collector})
 
     puts "func called"
-    
+
 
     puts "#{Baml::Collector.__function_span_count}"
     puts "#{Baml::Collector.__print_storage}"
@@ -71,24 +71,26 @@
 
     # Verify request/response
     request = call.http_request
-    refute_nil request
-    assert_kind_of Hash, request.body
-    assert_includes request.body, "messages"
-    assert_includes request.body["messages"][0], "content"
-    refute_nil request.body["messages"][0]["content"]
-    assert_equal "gpt-4o-mini", request.body["model"]
+    body = request.body.json()
+    refute_nil body
+    assert_kind_of Hash, body
+    assert_includes body, "messages"
+    assert_includes body["messages"][0], "content"
+    refute_nil body["messages"][0]["content"]
+    assert_equal "gpt-4o-mini", body["model"]
 
     # Verify http response
     response = call.http_response
     refute_nil response
+    body = response.body.json()
     assert_equal 200, response.status
-    refute_nil response.body
-    assert_kind_of Hash, response.body
-    assert_includes response.body, "choices"
-    assert response.body["choices"].length > 0
-    assert_includes response.body["choices"][0], "message"
-    assert_includes response.body["choices"][0]["message"], "content"
-    refute_nil response.body["choices"][0]["message"]["content"]
+    refute_nil body
+    assert_kind_of Hash, body
+    assert_includes body, "choices"
+    assert body["choices"].length > 0
+    assert_includes body["choices"][0], "message"
+    assert_includes body["choices"][0]["message"], "content"
+    refute_nil body["choices"][0]["message"]["content"]
 
     puts "call.body.headers: #{call.http_response.headers}"
     # Verify response headers contain openai-version
@@ -109,7 +111,7 @@
     assert call_usage.input_tokens > 0
     refute_nil call_usage.output_tokens
     assert call_usage.output_tokens > 0
-    
+
     # Matches log usage
     assert_equal call_usage.input_tokens, log.usage.input_tokens
     assert_equal call_usage.output_tokens, log.usage.output_tokens
@@ -149,7 +151,7 @@
     assert_equal 0, function_logs.length
 
     stream = b.stream.TestOpenAIGPT4oMini(input: "hi there", baml_options: {collector: collector})
-    
+
     chunks = []
     stream.each do |chunk|
       puts "### chunk: #{chunk}"
@@ -158,7 +160,7 @@
 
     res = stream.get_final_response
     puts "### res: #{res}"
-    
+
     function_logs = collector.logs
     assert_equal 1, function_logs.length
 
@@ -329,7 +331,7 @@
     threads << Thread.new { b.TestOpenAIGPT4oMini(input: "call #1", baml_options: {collector: collector}) }
     threads << Thread.new { b.TestOpenAIGPT4oMini(input: "call #2", baml_options: {collector: collector}) }
     threads.each(&:join)
-    
+
     puts "------------------------- ended parallel calls"
 
     # Verify the collector has two function logs

integ-tests/typescript/tests/collector.test.ts

@@ -73,12 +73,13 @@ describe('Collector Tests', () => {
 
     // Verify http response
     const response = call.httpResponse;
+    const responseBody = response?.body.json();
     expect(response).not.toBeNull();
     expect(response?.status).toBe(200);
-    expect(response?.body).not.toBeNull();
-    expect(response?.body.choices).toBeDefined();
-    expect(response?.body.choices.length).toBeGreaterThan(0);
-    expect(response?.body.choices[0].message.content).not.toBeNull();
+    expect(responseBody).not.toBeNull();
+    expect(responseBody?.choices).toBeDefined();
+    expect(responseBody?.choices.length).toBeGreaterThan(0);
+    expect(responseBody?.choices[0].message.content).not.toBeNull();
 
     // Verify call timing
     const callTiming = call.timing;
@@ -157,7 +158,7 @@ describe('Collector Tests', () => {
     const request = call.httpRequest;
     expect(request).not.toBeNull();
     expect(typeof request?.body).toBe('object');
-    expect((request?.body.json() as any).messages).toBeDefined();
+    expect((request?.body.json()).messages).toBeDefined();
 
     // For streaming, httpResponse might be null since it's streaming
     const response = call.httpResponse;
@@ -292,13 +293,13 @@ describe('Collector Tests', () => {
   it('should handle sync calls correctly', async () => {
     const collector = new Collector("sync-collector");
     const result = b_sync.TestOpenAIGPT4oMini("sync call", { collector });
-    
+
     const logs = collector.logs;
     expect(logs.length).toBe(1);
     expect(logs[0].functionName).toBe("TestOpenAIGPT4oMini");
     expect(logs[0].logType).toBe("call");
     expect(logs[0].usage).not.toBeNull();
   });
-  
-  
+
+
 });