Source file convert_prompt.ml

open Melange_json.Primitives

type image_source =
  | Base64_image of {
      media_type : string;
      data : string;
    }
  | Url_image of { url : string }

type document_source =
  | Base64_document of {
      media_type : string;
      data : string;
    }

type anthropic_tool_result_content =
  | Tool_text of string
  | Tool_image of { source : image_source }

type anthropic_content =
  | A_text of {
      text : string;
      cache_control : Cache_control.t option;
    }
  | A_image of {
      source : image_source;
      cache_control : Cache_control.t option;
    }
  | A_document of {
      source : document_source;
      cache_control : Cache_control.t option;
    }
  | A_tool_use of {
      id : string;
      name : string;
      input : Yojson.Basic.t;
    }
  | A_tool_result of {
      tool_use_id : string;
      content : anthropic_tool_result_content list;
      is_error : bool;
      cache_control : Cache_control.t option;
    }
  | A_thinking of {
      thinking : string;
      signature : string;
    }
  | A_redacted_thinking of { data : string }

type anthropic_message = {
  role : [ `User | `Assistant ];
  content : anthropic_content list;
}

(* Extract system messages with per-message provider_options and return the rest.
   Each entry is (content, provider_options); the order matches the input. *)
let extract_system messages =
  let system_parts, rest =
    List.partition_map
      (fun (msg : Ai_provider.Prompt.message) ->
        match msg with
        | System { content; provider_options } -> Left (content, provider_options)
        | User _ | Assistant _ | Tool _ -> Right msg)
      messages
  in
  system_parts, rest

(* Get cache control from provider options. Exposed for tests. *)
let get_cc po = Cache_control_options.get_cache_control po

(* Convert file_data to image source *)
let file_data_to_image_source ~media_type (data : Ai_provider.Prompt.file_data) =
  match data with
  | Bytes b -> Base64_image { media_type; data = Base64.encode_string (Bytes.to_string b) }
  | Base64 s -> Base64_image { media_type; data = s }
  | Url u -> Url_image { url = u }

(* Convert a user part to anthropic content. The validator is consulted for
   every cache_control so we honour Anthropic's 4-breakpoint limit. *)
let convert_user_part ~validator (part : Ai_provider.Prompt.user_part) : anthropic_content =
  let cc po = Cache_control_validator.take validator (get_cc po) in
  match part with
  | Text { text; provider_options } -> A_text { text; cache_control = cc provider_options }
  | File { data; media_type; provider_options; _ } ->
    if String.starts_with ~prefix:"image/" media_type then
      A_image { source = file_data_to_image_source ~media_type data; cache_control = cc provider_options }
    else
      A_document
        {
          source =
            (match data with
            | Bytes b -> Base64_document { media_type; data = Base64.encode_string (Bytes.to_string b) }
            | Base64 s -> Base64_document { media_type; data = s }
            | Url u -> invalid_arg (Printf.sprintf "Anthropic documents must be base64-encoded, got URL: %s" u));
          cache_control = cc provider_options;
        }

(* Convert an assistant part to anthropic content *)
let convert_assistant_part ~validator (part : Ai_provider.Prompt.assistant_part) : anthropic_content =
  let cc po = Cache_control_validator.take validator (get_cc po) in
  let reasoning_metadata_field name provider_options =
    match Ai_provider.Provider_options.provider_metadata provider_options with
    | Some (`Assoc fields) ->
      (match List.assoc_opt "anthropic" fields with
      | Some (`Assoc fields) ->
        (match List.assoc_opt name fields with
        | Some (`String value) when String.length value > 0 -> Some value
        | Some _ | None -> None)
      | Some _ | None -> None)
    | Some _ | None -> None
  in
  match part with
  | Text { text; provider_options } -> A_text { text; cache_control = cc provider_options }
  | File { data; media_type; provider_options; _ } ->
    A_image { source = file_data_to_image_source ~media_type data; cache_control = cc provider_options }
  | Reasoning { text; provider_options } ->
    (match
       reasoning_metadata_field "signature" provider_options, reasoning_metadata_field "redactedData" provider_options
     with
    | Some signature, _ -> A_thinking { thinking = text; signature }
    | None, Some data -> A_redacted_thinking { data }
    | None, None -> invalid_arg "Anthropic reasoning block is missing its signature or redacted data")
  | Tool_call { id; name; args; provider_options = _ } -> A_tool_use { id; name; input = args }

(* Convert a tool result to anthropic content *)
let convert_tool_result ~validator (tr : Ai_provider.Prompt.tool_result) : anthropic_content =
  let content =
    List.map
      (fun (c : Ai_provider.Prompt.tool_result_content) ->
        match c with
        | Result_text s -> Tool_text s
        | Result_image { data; media_type } -> Tool_image { source = Base64_image { media_type; data } })
      tr.content
  in
  (* If no explicit content, use the result as text *)
  let content =
    match content with
    | [] ->
      (match tr.result with
      | `String s -> [ Tool_text s ]
      | json -> [ Tool_text (Yojson.Basic.to_string json) ])
    | _ -> content
  in
  A_tool_result
    {
      tool_use_id = tr.tool_call_id;
      content;
      is_error = tr.is_error;
      cache_control = Cache_control_validator.take validator (get_cc tr.provider_options);
    }

(* Convert a single SDK message to role + content parts *)
let convert_single_message ~validator (msg : Ai_provider.Prompt.message) :
  ([ `User | `Assistant ] * anthropic_content list) option =
  match msg with
  | System _ -> None (* already extracted *)
  | User { content } -> Some (`User, List.map (convert_user_part ~validator) content)
  | Assistant { content } -> Some (`Assistant, List.map (convert_assistant_part ~validator) content)
  | Tool { content } -> Some (`User, List.map (convert_tool_result ~validator) content)

(* Group messages to ensure alternating user/assistant roles.
   Consecutive messages with the same role are merged. *)
let role_equal (a : [ `User | `Assistant ]) (b : [ `User | `Assistant ]) =
  match a, b with
  | `User, `User | `Assistant, `Assistant -> true
  | `User, `Assistant | `Assistant, `User -> false

let group_messages (msgs : ([ `User | `Assistant ] * anthropic_content list) list) : anthropic_message list =
  let rec go acc = function
    | [] -> List.rev acc
    | (role, content) :: rest ->
    match acc with
    | { role = prev_role; content = prev_content } :: acc_rest when role_equal prev_role role ->
      go ({ role; content = prev_content @ content } :: acc_rest) rest
    | _ -> go ({ role; content } :: acc) rest
  in
  go [] msgs

let convert_messages ?validator messages =
  let validator =
    match validator with
    | Some v -> v
    | None -> Cache_control_validator.create ()
  in
  let role_content_pairs = List.filter_map (convert_single_message ~validator) messages in
  group_messages role_content_pairs

(* JSON serialization — typed records for each content shape *)

type cc = Cache_control.t

let cc_to_json (cc : cc) = Cache_control.to_json cc

type image_source_base64_json = {
  type_ : string; [@json.key "type"]
  media_type : string;
  data : string;
}
[@@deriving to_json]

type image_source_url_json = {
  type_ : string; [@json.key "type"]
  url : string;
}
[@@deriving to_json]

let image_source_to_json = function
  | Base64_image { media_type; data } -> image_source_base64_json_to_json { type_ = "base64"; media_type; data }
  | Url_image { url } -> image_source_url_json_to_json { type_ = "url"; url }

type text_content_json = {
  type_ : string; [@json.key "type"]
  text : string;
  cache_control : cc option; [@json.option] [@json.drop_default]
}
[@@deriving to_json]

(* Build the wire JSON for the system field. Always emit the array-of-blocks
   form, mirroring upstream @ai-sdk/anthropic's [convertToAnthropicPrompt]
   which unconditionally maps each system message to its own text block. The
   array form keeps the wire shape consistent regardless of cache_control,
   and avoids divergence between the cached and uncached paths. The validator
   is consulted to enforce the 4-breakpoint limit. *)
let system_to_json ?validator parts =
  let validator =
    match validator with
    | Some v -> v
    | None -> Cache_control_validator.create ()
  in
  match parts with
  | [] -> None
  | _ :: _ ->
    let blocks =
      List.map
        (fun (text, po) ->
          let cache_control = Cache_control_validator.take validator (get_cc po) in
          text_content_json_to_json { type_ = "text"; text; cache_control })
        parts
    in
    Some (`List blocks)

type source_content_json = {
  type_ : string; [@json.key "type"]
  source : Melange_json.t;
  cache_control : cc option; [@json.option] [@json.drop_default]
}
[@@deriving to_json]

type tool_use_json = {
  type_ : string; [@json.key "type"]
  id : string;
  name : string;
  input : Melange_json.t;
}
[@@deriving to_json]

type tool_result_json = {
  type_ : string; [@json.key "type"]
  tool_use_id : string;
  content : Melange_json.t list;
  is_error : bool;
  cache_control : cc option; [@json.option] [@json.drop_default]
}
[@@deriving to_json]

type thinking_json = {
  type_ : string; [@json.key "type"]
  thinking : string;
  signature : string;
}
[@@deriving to_json]

type redacted_thinking_json = {
  type_ : string; [@json.key "type"]
  data : string;
}
[@@deriving to_json]

let tool_result_content_to_json = function
  | Tool_text s -> text_content_json_to_json { type_ = "text"; text = s; cache_control = None }
  | Tool_image { source } ->
    source_content_json_to_json { type_ = "image"; source = image_source_to_json source; cache_control = None }

let anthropic_content_to_json = function
  | A_text { text; cache_control } -> text_content_json_to_json { type_ = "text"; text; cache_control }
  | A_image { source; cache_control } ->
    source_content_json_to_json { type_ = "image"; source = image_source_to_json source; cache_control }
  | A_document { source; cache_control } ->
    let (Base64_document { media_type; data }) = source in
    let source_json = image_source_base64_json_to_json { type_ = "base64"; media_type; data } in
    source_content_json_to_json { type_ = "document"; source = source_json; cache_control }
  | A_tool_use { id; name; input } -> tool_use_json_to_json { type_ = "tool_use"; id; name; input }
  | A_tool_result { tool_use_id; content; is_error; cache_control } ->
    let content_json = List.map tool_result_content_to_json content in
    tool_result_json_to_json { type_ = "tool_result"; tool_use_id; content = content_json; is_error; cache_control }
  | A_thinking { thinking; signature } -> thinking_json_to_json { type_ = "thinking"; thinking; signature }
  | A_redacted_thinking { data } -> redacted_thinking_json_to_json { type_ = "redacted_thinking"; data }

type message_json = {
  role : string;
  content : Melange_json.t list;
}
[@@deriving to_json]

let anthropic_message_to_json ({ role; content } : anthropic_message) =
  let role_str =
    match role with
    | `User -> "user"
    | `Assistant -> "assistant"
  in
  message_json_to_json { role = role_str; content = List.map anthropic_content_to_json content }