My agentic slop goes here. Not intended for anyone else!
at jsont 19 kB view raw
1open Ezjsonm 2module JU = Json_utils 3 4let src = Logs.Src.create "claude.message" ~doc:"Claude messages" 5module Log = (val Logs.src_log src : Logs.LOG) 6 7 8module User = struct 9 type content = 10 | String of string 11 | Blocks of Content_block.t list 12 13 type t = { content : content } 14 15 let create_string s = { content = String s } 16 let create_blocks blocks = { content = Blocks blocks } 17 18 let create_with_tool_result ~tool_use_id ~content ?is_error () = 19 let tool_result = Content_block.tool_result ~tool_use_id ~content ?is_error () in 20 { content = Blocks [tool_result] } 21 22 let create_mixed ~text ~tool_results = 23 let blocks = 24 let text_blocks = match text with 25 | Some t -> [Content_block.text t] 26 | None -> [] 27 in 28 let tool_blocks = List.map (fun (tool_use_id, content, is_error) -> 29 Content_block.tool_result ~tool_use_id ~content ?is_error () 30 ) tool_results in 31 text_blocks @ tool_blocks 32 in 33 { content = Blocks blocks } 34 35 let content t = t.content 36 37 let as_text t = match t.content with 38 | String s -> Some s 39 | Blocks _ -> None 40 41 let get_blocks t = match t.content with 42 | String s -> [Content_block.text s] 43 | Blocks blocks -> blocks 44 45 let to_json t = 46 let content_json = match t.content with 47 | String s -> `String s 48 | Blocks blocks -> 49 `A (List.map Content_block.to_json blocks) 50 in 51 `O [ 52 ("type", `String "user"); 53 ("message", `O [ 54 ("role", `String "user"); 55 ("content", content_json); 56 ]); 57 ] 58 59 let of_json = function 60 | `O fields -> 61 let message = List.assoc "message" fields in 62 let content = match message with 63 | `O msg_fields -> 64 (match List.assoc "content" msg_fields with 65 | `String s -> String s 66 | `A blocks -> Blocks (List.map Content_block.of_json blocks) 67 | _ -> raise (Invalid_argument "User.of_json: invalid content")) 68 | _ -> raise (Invalid_argument "User.of_json: invalid message") 69 in 70 { content } 71 | _ -> raise (Invalid_argument "User.of_json: expected object") 72 73 let pp fmt t = 74 match t.content with 75 | String s -> 76 if String.length s > 60 then 77 let truncated = String.sub s 0 57 in 78 Fmt.pf fmt "@[<2>User:@ %s...@]" truncated 79 else 80 Fmt.pf fmt "@[<2>User:@ %S@]" s 81 | Blocks blocks -> 82 let text_count = List.length (List.filter (function 83 | Content_block.Text _ -> true | _ -> false) blocks) in 84 let tool_result_count = List.length (List.filter (function 85 | Content_block.Tool_result _ -> true | _ -> false) blocks) in 86 match text_count, tool_result_count with 87 | 1, 0 -> 88 let text = List.find_map (function 89 | Content_block.Text t -> Some (Content_block.Text.text t) 90 | _ -> None) blocks in 91 Fmt.pf fmt "@[<2>User:@ %a@]" Fmt.(option string) text 92 | 0, 1 -> 93 Fmt.pf fmt "@[<2>User:@ [tool result]@]" 94 | 0, n when n > 1 -> 95 Fmt.pf fmt "@[<2>User:@ [%d tool results]@]" n 96 | _ -> 97 Fmt.pf fmt "@[<2>User:@ [%d blocks]@]" (List.length blocks) 98end 99 100module Assistant = struct 101 type error = [ 102 | `Authentication_failed 103 | `Billing_error 104 | `Rate_limit 105 | `Invalid_request 106 | `Server_error 107 | `Unknown 108 ] 109 110 let error_to_string = function 111 | `Authentication_failed -> "authentication_failed" 112 | `Billing_error -> "billing_error" 113 | `Rate_limit -> "rate_limit" 114 | `Invalid_request -> "invalid_request" 115 | `Server_error -> "server_error" 116 | `Unknown -> "unknown" 117 118 let error_of_string = function 119 | "authentication_failed" -> `Authentication_failed 120 | "billing_error" -> `Billing_error 121 | "rate_limit" -> `Rate_limit 122 | "invalid_request" -> `Invalid_request 123 | "server_error" -> `Server_error 124 | "unknown" | _ -> `Unknown 125 126 type t = { 127 content : Content_block.t list; 128 model : string; 129 error : error option; 130 } 131 132 let create ~content ~model ?error () = { content; model; error } 133 let content t = t.content 134 let model t = t.model 135 let error t = t.error 136 137 let get_text_blocks t = 138 List.filter_map (function 139 | Content_block.Text text -> Some (Content_block.Text.text text) 140 | _ -> None 141 ) t.content 142 143 let get_tool_uses t = 144 List.filter_map (function 145 | Content_block.Tool_use tool -> Some tool 146 | _ -> None 147 ) t.content 148 149 let get_thinking t = 150 List.filter_map (function 151 | Content_block.Thinking thinking -> Some thinking 152 | _ -> None 153 ) t.content 154 155 let has_tool_use t = 156 List.exists (function 157 | Content_block.Tool_use _ -> true 158 | _ -> false 159 ) t.content 160 161 let combined_text t = 162 String.concat "\n" (get_text_blocks t) 163 164 let to_json t = 165 let msg_fields = [ 166 ("content", `A (List.map Content_block.to_json t.content)); 167 ("model", `String t.model); 168 ] in 169 let msg_fields = match t.error with 170 | Some err -> ("error", `String (error_to_string err)) :: msg_fields 171 | None -> msg_fields 172 in 173 `O [ 174 ("type", `String "assistant"); 175 ("message", `O msg_fields); 176 ] 177 178 let of_json = function 179 | `O fields -> 180 let message = List.assoc "message" fields in 181 let content, model, error = match message with 182 | `O msg_fields -> 183 let content = 184 match List.assoc "content" msg_fields with 185 | `A blocks -> List.map Content_block.of_json blocks 186 | _ -> raise (Invalid_argument "Assistant.of_json: invalid content") 187 in 188 let model = JU.assoc_string "model" msg_fields in 189 let error = 190 match JU.assoc_string_opt "error" msg_fields with 191 | Some err_str -> Some (error_of_string err_str) 192 | None -> None 193 in 194 content, model, error 195 | _ -> raise (Invalid_argument "Assistant.of_json: invalid message") 196 in 197 { content; model; error } 198 | _ -> raise (Invalid_argument "Assistant.of_json: expected object") 199 200 let pp fmt t = 201 let text_count = List.length (get_text_blocks t) in 202 let tool_count = List.length (get_tool_uses t) in 203 let thinking_count = List.length (get_thinking t) in 204 match text_count, tool_count, thinking_count with 205 | 1, 0, 0 -> 206 (* Simple text response *) 207 Fmt.pf fmt "@[<2>Assistant@ [%s]:@ %S@]" 208 t.model (combined_text t) 209 | _, 0, 0 when text_count > 0 -> 210 (* Multiple text blocks *) 211 Fmt.pf fmt "@[<2>Assistant@ [%s]:@ %d text blocks@]" 212 t.model text_count 213 | 0, _, 0 when tool_count > 0 -> 214 (* Only tool uses *) 215 let tools = get_tool_uses t in 216 let tool_names = List.map Content_block.Tool_use.name tools in 217 Fmt.pf fmt "@[<2>Assistant@ [%s]:@ tools(%a)@]" 218 t.model Fmt.(list ~sep:comma string) tool_names 219 | _ -> 220 (* Mixed content *) 221 let parts = [] in 222 let parts = if text_count > 0 then 223 Printf.sprintf "%d text" text_count :: parts else parts in 224 let parts = if tool_count > 0 then 225 Printf.sprintf "%d tools" tool_count :: parts else parts in 226 let parts = if thinking_count > 0 then 227 Printf.sprintf "%d thinking" thinking_count :: parts else parts in 228 Fmt.pf fmt "@[<2>Assistant@ [%s]:@ %s@]" 229 t.model (String.concat ", " (List.rev parts)) 230end 231 232module System = struct 233 module Data = struct 234 (* Store both the raw JSON and provide typed accessors *) 235 type t = value (* The full JSON data *) 236 237 let empty = `O [] 238 239 let of_assoc assoc = `O assoc 240 241 let get_string t key = JU.get_field_string_opt t key 242 243 let get_int t key = JU.get_field_int_opt t key 244 245 let get_bool t key = JU.get_field_bool_opt t key 246 247 let get_float t key = JU.get_field_float_opt t key 248 249 let get_list t key = 250 match t with 251 | `O fields -> 252 (match List.assoc_opt key fields with 253 | Some (`A lst) -> Some lst 254 | _ -> None) 255 | _ -> None 256 257 let get_field t key = 258 match t with 259 | `O fields -> List.assoc_opt key fields 260 | _ -> None 261 262 let raw_json t = t 263 264 let to_json t = t 265 let of_json json = json 266 end 267 268 type t = { 269 subtype : string; 270 data : Data.t; 271 } 272 273 let create ~subtype ~data = { subtype; data } 274 let subtype t = t.subtype 275 let data t = t.data 276 277 let to_json t = 278 `O [ 279 ("type", `String "system"); 280 ("subtype", `String t.subtype); 281 ("data", Data.to_json t.data); 282 ] 283 284 let of_json = function 285 | `O fields -> 286 let subtype = JU.assoc_string "subtype" fields in 287 let data = Data.of_json (try List.assoc "data" fields with Not_found -> `O fields) in 288 { subtype; data } 289 | _ -> raise (Invalid_argument "System.of_json: expected object") 290 291 let pp fmt t = 292 match t.subtype with 293 | "init" -> 294 let session_id = Data.get_string t.data "session_id" in 295 let model = Data.get_string t.data "model" in 296 let cwd = Data.get_string t.data "cwd" in 297 Fmt.pf fmt "@[<2>System.init@ { session_id = %a;@ model = %a;@ cwd = %a }@]" 298 Fmt.(option string) session_id 299 Fmt.(option string) model 300 Fmt.(option string) cwd 301 | "error" -> 302 let error = Data.get_string t.data "error" in 303 Fmt.pf fmt "@[<2>System.error@ { error = %a }@]" 304 Fmt.(option string) error 305 | _ -> 306 Fmt.pf fmt "@[<2>System.%s@ { ... }@]" t.subtype 307end 308 309module Result = struct 310 module Usage = struct 311 type t = value 312 313 let create ?input_tokens ?output_tokens ?total_tokens 314 ?cache_creation_input_tokens ?cache_read_input_tokens () = 315 let fields = [] in 316 let fields = match input_tokens with 317 | Some n -> ("input_tokens", `Float (float_of_int n)) :: fields 318 | None -> fields in 319 let fields = match output_tokens with 320 | Some n -> ("output_tokens", `Float (float_of_int n)) :: fields 321 | None -> fields in 322 let fields = match total_tokens with 323 | Some n -> ("total_tokens", `Float (float_of_int n)) :: fields 324 | None -> fields in 325 let fields = match cache_creation_input_tokens with 326 | Some n -> ("cache_creation_input_tokens", `Float (float_of_int n)) :: fields 327 | None -> fields in 328 let fields = match cache_read_input_tokens with 329 | Some n -> ("cache_read_input_tokens", `Float (float_of_int n)) :: fields 330 | None -> fields in 331 `O fields 332 333 let input_tokens t = JU.get_field_int_opt t "input_tokens" 334 335 let output_tokens t = JU.get_field_int_opt t "output_tokens" 336 337 let total_tokens t = JU.get_field_int_opt t "total_tokens" 338 339 let cache_creation_input_tokens t = JU.get_field_int_opt t "cache_creation_input_tokens" 340 341 let cache_read_input_tokens t = JU.get_field_int_opt t "cache_read_input_tokens" 342 343 let effective_input_tokens t = 344 match input_tokens t with 345 | None -> 0 346 | Some input -> 347 let cached = Option.value (cache_read_input_tokens t) ~default:0 in 348 max 0 (input - cached) 349 350 let total_cost_estimate t ~input_price ~output_price = 351 match input_tokens t, output_tokens t with 352 | Some input, Some output -> 353 let input_cost = float_of_int input *. input_price /. 1_000_000. in 354 let output_cost = float_of_int output *. output_price /. 1_000_000. in 355 Some (input_cost +. output_cost) 356 | _ -> None 357 358 let pp fmt t = 359 Fmt.pf fmt "@[<2>Usage@ { input = %a;@ output = %a;@ total = %a;@ \ 360 cache_creation = %a;@ cache_read = %a }@]" 361 Fmt.(option int) (input_tokens t) 362 Fmt.(option int) (output_tokens t) 363 Fmt.(option int) (total_tokens t) 364 Fmt.(option int) (cache_creation_input_tokens t) 365 Fmt.(option int) (cache_read_input_tokens t) 366 367 let to_json t = t 368 let of_json json = json 369 end 370 371 type t = { 372 subtype : string; 373 duration_ms : int; 374 duration_api_ms : int; 375 is_error : bool; 376 num_turns : int; 377 session_id : string; 378 total_cost_usd : float option; 379 usage : Usage.t option; 380 result : string option; 381 structured_output : value option; 382 } 383 384 let create ~subtype ~duration_ms ~duration_api_ms ~is_error ~num_turns 385 ~session_id ?total_cost_usd ?usage ?result ?structured_output () = 386 { subtype; duration_ms; duration_api_ms; is_error; num_turns; 387 session_id; total_cost_usd; usage; result; structured_output } 388 389 let subtype t = t.subtype 390 let duration_ms t = t.duration_ms 391 let duration_api_ms t = t.duration_api_ms 392 let is_error t = t.is_error 393 let num_turns t = t.num_turns 394 let session_id t = t.session_id 395 let total_cost_usd t = t.total_cost_usd 396 let usage t = t.usage 397 let result t = t.result 398 let structured_output t = t.structured_output 399 400 let to_json t = 401 let fields = [ 402 ("type", `String "result"); 403 ("subtype", `String t.subtype); 404 ("duration_ms", `Float (float_of_int t.duration_ms)); 405 ("duration_api_ms", `Float (float_of_int t.duration_api_ms)); 406 ("is_error", `Bool t.is_error); 407 ("num_turns", `Float (float_of_int t.num_turns)); 408 ("session_id", `String t.session_id); 409 ] in 410 let fields = match t.total_cost_usd with 411 | Some cost -> ("total_cost_usd", `Float cost) :: fields 412 | None -> fields 413 in 414 let fields = match t.usage with 415 | Some usage -> ("usage", Usage.to_json usage) :: fields 416 | None -> fields 417 in 418 let fields = match t.result with 419 | Some result -> ("result", `String result) :: fields 420 | None -> fields 421 in 422 let fields = match t.structured_output with 423 | Some output -> ("structured_output", output) :: fields 424 | None -> fields 425 in 426 `O fields 427 428 let of_json = function 429 | `O fields -> 430 let subtype = JU.assoc_string "subtype" fields in 431 let duration_ms = int_of_float (JU.assoc_float "duration_ms" fields) in 432 let duration_api_ms = int_of_float (JU.assoc_float "duration_api_ms" fields) in 433 let is_error = JU.assoc_bool "is_error" fields in 434 let num_turns = int_of_float (JU.assoc_float "num_turns" fields) in 435 let session_id = JU.assoc_string "session_id" fields in 436 let total_cost_usd = JU.assoc_float_opt "total_cost_usd" fields in 437 let usage = Option.map Usage.of_json (List.assoc_opt "usage" fields) in 438 let result = JU.assoc_string_opt "result" fields in 439 let structured_output = List.assoc_opt "structured_output" fields in 440 { subtype; duration_ms; duration_api_ms; is_error; num_turns; 441 session_id; total_cost_usd; usage; result; structured_output } 442 | _ -> raise (Invalid_argument "Result.of_json: expected object") 443 444 let pp fmt t = 445 if t.is_error then 446 Fmt.pf fmt "@[<2>Result.error@ { session = %S;@ result = %a }@]" 447 t.session_id 448 Fmt.(option string) t.result 449 else 450 let tokens_info = match t.usage with 451 | Some u -> 452 let input = Usage.input_tokens u in 453 let output = Usage.output_tokens u in 454 let cached = Usage.cache_read_input_tokens u in 455 (match input, output, cached with 456 | Some i, Some o, Some c when c > 0 -> 457 Printf.sprintf " (tokens: %d+%d, cached: %d)" i o c 458 | Some i, Some o, _ -> 459 Printf.sprintf " (tokens: %d+%d)" i o 460 | _ -> "") 461 | None -> "" 462 in 463 Fmt.pf fmt "@[<2>Result.%s@ { duration = %dms;@ cost = $%.4f%s }@]" 464 t.subtype 465 t.duration_ms 466 (Option.value t.total_cost_usd ~default:0.0) 467 tokens_info 468end 469 470type t = 471 | User of User.t 472 | Assistant of Assistant.t 473 | System of System.t 474 | Result of Result.t 475 476let user_string s = User (User.create_string s) 477let user_blocks blocks = User (User.create_blocks blocks) 478let user_with_tool_result ~tool_use_id ~content ?is_error () = 479 User (User.create_with_tool_result ~tool_use_id ~content ?is_error ()) 480 481let assistant ~content ~model ?error () = Assistant (Assistant.create ~content ~model ?error ()) 482let assistant_text ~text ~model ?error () = 483 Assistant (Assistant.create ~content:[Content_block.text text] ~model ?error ()) 484 485let system ~subtype ~data = System (System.create ~subtype ~data) 486let system_init ~session_id = 487 let data = System.Data.of_assoc [("session_id", `String session_id)] in 488 System (System.create ~subtype:"init" ~data) 489let system_error ~error = 490 let data = System.Data.of_assoc [("error", `String error)] in 491 System (System.create ~subtype:"error" ~data) 492 493let result ~subtype ~duration_ms ~duration_api_ms ~is_error ~num_turns 494 ~session_id ?total_cost_usd ?usage ?result ?structured_output () = 495 Result (Result.create ~subtype ~duration_ms ~duration_api_ms ~is_error 496 ~num_turns ~session_id ?total_cost_usd ?usage ?result ?structured_output ()) 497 498let to_json = function 499 | User t -> User.to_json t 500 | Assistant t -> Assistant.to_json t 501 | System t -> System.to_json t 502 | Result t -> Result.to_json t 503 504let of_json json = 505 match json with 506 | `O fields -> ( 507 match List.assoc_opt "type" fields with 508 | Some (`String "user") -> User (User.of_json json) 509 | Some (`String "assistant") -> Assistant (Assistant.of_json json) 510 | Some (`String "system") -> System (System.of_json json) 511 | Some (`String "result") -> Result (Result.of_json json) 512 | _ -> raise (Invalid_argument "Message.of_json: unknown type") 513 ) 514 | _ -> raise (Invalid_argument "Message.of_json: expected object") 515 516let pp fmt = function 517 | User t -> User.pp fmt t 518 | Assistant t -> Assistant.pp fmt t 519 | System t -> System.pp fmt t 520 | Result t -> Result.pp fmt t 521 522let is_user = function User _ -> true | _ -> false 523let is_assistant = function Assistant _ -> true | _ -> false 524let is_system = function System _ -> true | _ -> false 525let is_result = function Result _ -> true | _ -> false 526 527let is_error = function 528 | Result r -> Result.is_error r 529 | System s -> System.subtype s = "error" 530 | _ -> false 531 532let extract_text = function 533 | User u -> User.as_text u 534 | Assistant a -> 535 let text = Assistant.combined_text a in 536 if text = "" then None else Some text 537 | _ -> None 538 539let extract_tool_uses = function 540 | Assistant a -> Assistant.get_tool_uses a 541 | _ -> [] 542 543let get_session_id = function 544 | System s when System.subtype s = "init" -> 545 System.Data.get_string (System.data s) "session_id" 546 | Result r -> Some (Result.session_id r) 547 | _ -> None 548 549let log_received t = 550 Log.info (fun m -> m "← %a" pp t) 551 552let log_sending t = 553 Log.info (fun m -> m "→ %a" pp t) 554 555let log_error msg t = 556 Log.err (fun m -> m "%s: %a" msg pp t) 557