Source file convert_stream.ml

open Melange_json.Primitives

(* Content block tracking state *)
type block_state =
  | Text_block
  | Tool_use_block of {
      id : string;
      name : string;
    }
  | Thinking_block

(* Typed records for SSE event JSON payloads *)

type content_block_info = {
  type_ : string; [@json.key "type"]
  id : string option; [@json.default None]
  name : string option; [@json.default None]
  data : string option; [@json.default None]
}
[@@json.allow_extra_fields] [@@deriving of_json]

type content_block_start_event = {
  index : int;
  content_block : content_block_info;
}
[@@json.allow_extra_fields] [@@deriving of_json]

type delta_info = {
  type_ : string; [@json.key "type"]
  text : string option; [@json.default None]
  partial_json : string option; [@json.default None]
  thinking : string option; [@json.default None]
  signature : string option; [@json.default None]
}
[@@json.allow_extra_fields] [@@deriving of_json]

type content_block_delta_event = {
  index : int;
  delta : delta_info;
}
[@@json.allow_extra_fields] [@@deriving of_json]

type content_block_stop_event = { index : int } [@@json.allow_extra_fields] [@@deriving of_json]

type message_delta_info = { stop_reason : string option [@json.default None] }
[@@json.allow_extra_fields] [@@deriving of_json]

type message_delta_event = {
  delta : message_delta_info;
  usage : Convert_usage.anthropic_usage option; [@json.default None]
}
[@@json.allow_extra_fields] [@@deriving of_json]

type message_start_message = { usage : Convert_usage.anthropic_usage option [@json.default None] }
[@@json.allow_extra_fields] [@@deriving of_json]

type message_start_event = { message : message_start_message } [@@json.allow_extra_fields] [@@deriving of_json]

type error_info = {
  type_ : string; [@json.key "type"] [@json.default "unknown"]
  message : string; [@json.default ""]
}
[@@json.allow_extra_fields] [@@deriving of_json]

type error_event = { error : error_info } [@@json.allow_extra_fields] [@@deriving of_json]

let transform events ~warnings =
  let blocks : (int, block_state) Hashtbl.t = Hashtbl.create 8 in
  let is_first = ref true in
  (* Anthropic emits cache_creation_input_tokens / cache_read_input_tokens on
     [message_start.message.usage]; [message_delta.usage] carries the final
     output count. Stash the start usage and fall back to it when the delta
     omits the cache fields, mirroring upstream @ai-sdk/anthropic. *)
  let start_usage : Convert_usage.anthropic_usage option ref = ref None in
  let stream, push = Lwt_stream.create () in
  Lwt.async (fun () ->
    let%lwt () =
      Lwt_stream.iter
        (fun (evt : Sse.event) ->
          try
            let json = Yojson.Basic.from_string evt.data in
            match evt.event_type with
            | "message_start" ->
              (try
                 let evt = message_start_event_of_json json in
                 start_usage := evt.message.usage
               with Melange_json.Of_json_error _ -> ());
              if !is_first then begin
                push (Some (Ai_provider.Stream_part.Stream_start { warnings }));
                is_first := false
              end
            | "content_block_start" ->
              let { index; content_block } = content_block_start_event_of_json json in
              (match content_block.type_ with
              | "text" -> Hashtbl.replace blocks index Text_block
              | "tool_use" ->
                (match content_block.id, content_block.name with
                | Some id, Some name -> Hashtbl.replace blocks index (Tool_use_block { id; name })
                | _ -> ())
              | "thinking" -> Hashtbl.replace blocks index Thinking_block
              | "redacted_thinking" ->
                Option.iter
                  (fun data ->
                    push
                      (Some
                         (Ai_provider.Stream_part.Reasoning
                            {
                              text = "";
                              signature = None;
                              provider_options = Convert_response.redacted_reasoning_provider_options data;
                            })))
                  content_block.data
              | _ -> ())
            | "content_block_delta" ->
              let { index; delta } = content_block_delta_event_of_json json in
              (match delta.type_ with
              | "text_delta" ->
                (match delta.text with
                | Some text -> push (Some (Ai_provider.Stream_part.Text { text }))
                | None -> ())
              | "input_json_delta" ->
                (match delta.partial_json with
                | Some partial ->
                  (match Hashtbl.find_opt blocks index with
                  | Some (Tool_use_block { id; name }) ->
                    push
                      (Some
                         (Ai_provider.Stream_part.Tool_call_delta
                            {
                              tool_call_type = "function";
                              tool_call_id = id;
                              tool_name = name;
                              args_text_delta = partial;
                            }))
                  | _ -> ())
                | None -> ())
              | "thinking_delta" ->
                (match delta.thinking with
                | Some text ->
                  push
                    (Some
                       (Ai_provider.Stream_part.Reasoning
                          { text; signature = None; provider_options = Ai_provider.Provider_options.empty }))
                | None -> ())
              | "signature_delta" ->
                (match delta.signature with
                | Some signature ->
                  push
                    (Some
                       (Ai_provider.Stream_part.Reasoning
                          {
                            text = "";
                            signature = Some signature;
                            provider_options = Convert_response.reasoning_provider_options (Some signature);
                          }))
                | None -> ())
              | _ -> ())
            | "content_block_stop" ->
              let { index } = content_block_stop_event_of_json json in
              (match Hashtbl.find_opt blocks index with
              | Some (Tool_use_block { id; _ }) ->
                push (Some (Ai_provider.Stream_part.Tool_call_finish { tool_call_id = id }))
              | _ -> ());
              Hashtbl.remove blocks index
            | "message_delta" ->
              let { delta; usage } = message_delta_event_of_json json in
              let has_cache_signal (u : Convert_usage.anthropic_usage) =
                Option.is_some u.cache_read_input_tokens
                || Option.is_some u.cache_creation_input_tokens
                || Option.is_some u.cache_creation
              in
              (* Delta carries the final output_tokens; start carries the cache
                 fields. When delta lacks cache signal, overlay it onto start
                 so the Finish chunk reports both accurately. *)
              let anthropic_usage =
                match usage, !start_usage with
                | None, _ -> !start_usage
                | Some u, _ when has_cache_signal u -> Some u
                | Some u, None -> Some u
                | Some u, Some s -> Some { s with output_tokens = u.output_tokens }
              in
              let usage_ai =
                match anthropic_usage with
                | Some u -> Convert_usage.to_usage u
                | None -> { Ai_provider.Usage.input_tokens = 0; output_tokens = 0; total_tokens = None }
              in
              (* Surface Anthropic cache token metrics to streaming consumers, matching
                 the non-streaming path. Upstream attaches providerMetadata to the
                 [finish] LanguageModelV4 stream part rather than emitting it as a
                 separate chunk. *)
              let provider_metadata =
                match anthropic_usage with
                | Some u when has_cache_signal u -> Some (Convert_usage.to_provider_metadata u)
                | _ -> None
              in
              push
                (Some
                   (Ai_provider.Stream_part.Finish
                      {
                        finish_reason = Convert_response.map_stop_reason delta.stop_reason;
                        usage = usage_ai;
                        provider_metadata;
                      }))
            | "message_stop" | "ping" -> ()
            | "error" ->
              let error_type, message =
                try
                  let { error = { type_; message } } = error_event_of_json json in
                  type_, message
                with Melange_json.Of_json_error _ -> "unknown", evt.data
              in
              push
                (Some
                   (Ai_provider.Stream_part.Error
                      {
                        error =
                          Ai_provider.Provider_error.make_api_error ~provider:"anthropic" ~status:0
                            ~body:(Printf.sprintf "%s: %s" error_type message)
                            ~is_retryable:false ();
                      }))
            | _ -> ()
          with (Yojson.Json_error _ | Melange_json.Of_json_error _) as exn ->
            push
              (Some
                 (Ai_provider.Stream_part.Error
                    {
                      error =
                        {
                          Ai_provider.Provider_error.provider = "anthropic";
                          kind = Deserialization_error { message = Printexc.to_string exn; raw = evt.data };
                          is_retryable = false;
                          retry_after_s = None;
                        };
                    })))
        events
    in
    push None;
    Lwt.return_unit);
  stream