My agentic slop goes here. Not intended for anyone else!
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