···
94
+
(* Jsont codec for parsing incoming user messages from CLI *)
95
+
let incoming_jsont : t Jsont.t =
97
+
Jsont.Object.map ~kind:"UserMessage" (fun json_content ->
98
+
let content = decode_content json_content in
99
+
{ content; unknown = Unknown.empty }
101
+
|> Jsont.Object.mem "content" Jsont.json ~enc:(fun t -> encode_content (content t))
102
+
|> Jsont.Object.finish
104
+
Jsont.Object.map ~kind:"UserEnvelope" Fun.id
105
+
|> Jsont.Object.mem "message" message_jsont ~enc:Fun.id
106
+
|> Jsont.Object.finish
96
-
| Jsont.Object (fields, _) ->
97
-
let message = List.assoc (Jsont.Json.name "message") fields in
98
-
let content = match message with
99
-
| Jsont.Object (msg_fields, _) ->
100
-
(match List.assoc (Jsont.Json.name "content") msg_fields with
101
-
| Jsont.String (s, _) -> String s
102
-
| Jsont.Array (items, _) ->
103
-
Blocks (List.map Content_block.of_json items)
104
-
| _ -> raise (Invalid_argument "User.of_json: invalid content"))
105
-
| _ -> raise (Invalid_argument "User.of_json: invalid message")
107
-
{ content; unknown = Unknown.empty }
108
-
| _ -> raise (Invalid_argument "User.of_json: expected object")
109
+
match Jsont.Json.decode incoming_jsont json with
111
+
| Error msg -> raise (Invalid_argument ("User.of_json: " ^ msg))
···
(Jsont.Json.name "message", Jsont.Object (msg_fields, Jsont.Meta.none));
246
+
(* Jsont codec for parsing incoming assistant messages from CLI *)
247
+
let incoming_jsont : t Jsont.t =
248
+
Jsont.Object.map ~kind:"AssistantEnvelope" Fun.id
249
+
|> Jsont.Object.mem "message" jsont ~enc:Fun.id
250
+
|> Jsont.Object.finish
245
-
| Jsont.Object (fields, _) ->
246
-
let message = List.assoc (Jsont.Json.name "message") fields in
247
-
let content, model, error = match message with
248
-
| Jsont.Object (msg_fields, _) ->
250
-
match List.assoc (Jsont.Json.name "content") msg_fields with
251
-
| Jsont.Array (items, _) -> List.map Content_block.of_json items
252
-
| _ -> raise (Invalid_argument "Assistant.of_json: invalid content")
254
-
let model = match List.assoc (Jsont.Json.name "model") msg_fields with
255
-
| Jsont.String (s, _) -> s
256
-
| _ -> raise (Invalid_argument "Assistant.of_json: invalid model")
259
-
match List.assoc_opt (Jsont.Json.name "error") msg_fields with
260
-
| Some (Jsont.String (err_str, _)) -> Some (error_of_string err_str)
261
-
| Some _ -> raise (Invalid_argument "Assistant.of_json: invalid error")
264
-
content, model, error
265
-
| _ -> raise (Invalid_argument "Assistant.of_json: invalid message")
267
-
{ content; model; error; unknown = Unknown.empty }
268
-
| _ -> raise (Invalid_argument "Assistant.of_json: expected object")
253
+
match Jsont.Json.decode incoming_jsont json with
255
+
| Error msg -> raise (Invalid_argument ("Assistant.of_json: " ^ msg))
let text_count = List.length (get_text_blocks t) in
···
303
-
module Data = struct
304
-
(* Opaque JSON type with typed accessors *)
305
-
type t = Jsont.json
290
+
(** Typed data for system init messages *)
291
+
module Init = struct
292
+
module Unknown = struct
293
+
type t = Jsont.json
294
+
let empty = Jsont.Object ([], Jsont.Meta.none)
295
+
let _jsont = Jsont.json
299
+
session_id : string option;
300
+
model : string option;
301
+
cwd : string option;
302
+
unknown : Unknown.t;
305
+
let make session_id model cwd unknown = { session_id; model; cwd; unknown }
307
-
let jsont = Jsont.json
307
+
let create ?session_id ?model ?cwd () =
308
+
{ session_id; model; cwd; unknown = Unknown.empty }
310
+
let session_id t = t.session_id
311
+
let model t = t.model
313
+
let unknown t = t.unknown
309
-
let empty = Jsont.Object ([], Jsont.Meta.none)
315
+
let jsont : t Jsont.t =
316
+
Jsont.Object.map ~kind:"SystemInit" make
317
+
|> Jsont.Object.opt_mem "session_id" Jsont.string ~enc:session_id
318
+
|> Jsont.Object.opt_mem "model" Jsont.string ~enc:model
319
+
|> Jsont.Object.opt_mem "cwd" Jsont.string ~enc:cwd
320
+
|> Jsont.Object.keep_unknown Jsont.json_mems ~enc:unknown
321
+
|> Jsont.Object.finish
311
-
let of_assoc (assoc : (string * Jsont.json) list) : t =
313
-
List.map (fun (k, v) -> (Jsont.Json.name k, v)) assoc,
324
+
(** Typed data for system error messages *)
325
+
module Error = struct
326
+
module Unknown = struct
327
+
type t = Jsont.json
328
+
let empty = Jsont.Object ([], Jsont.Meta.none)
329
+
let _jsont = Jsont.json
334
+
unknown : Unknown.t;
337
+
let make error unknown = { error; unknown }
339
+
let create ~error = { error; unknown = Unknown.empty }
341
+
let error t = t.error
342
+
let unknown t = t.unknown
344
+
let jsont : t Jsont.t =
345
+
Jsont.Object.map ~kind:"SystemError" make
346
+
|> Jsont.Object.mem "error" Jsont.string ~enc:error
347
+
|> Jsont.Object.keep_unknown Jsont.json_mems ~enc:unknown
348
+
|> Jsont.Object.finish
317
-
let get_field t key =
319
-
| Jsont.Object (members, _) ->
320
-
List.find_map (fun ((name, _), value) ->
321
-
if name = key then Some value else None
351
+
(** Sum type for system message data *)
352
+
module Data = struct
356
+
| Other of Jsont.json (** Unknown subtypes preserve raw JSON *)
325
-
let get_string t key =
326
-
match get_field t key with
327
-
| Some (Jsont.String (s, _)) -> Some s
358
+
let init ?session_id ?model ?cwd () = Init (Init.create ?session_id ?model ?cwd ())
359
+
let error ~error = Error (Error.create ~error)
360
+
let other json = Other json
330
-
let get_int t key =
331
-
match get_field t key with
332
-
| Some (Jsont.Number (f, _)) ->
333
-
let i = int_of_float f in
334
-
if float_of_int i = f then Some i else None
362
+
let session_id = function
363
+
| Init i -> Init.session_id i
337
-
let get_bool t key =
338
-
match get_field t key with
339
-
| Some (Jsont.Bool (b, _)) -> Some b
366
+
let model = function
367
+
| Init i -> Init.model i
342
-
let get_float t key =
343
-
match get_field t key with
344
-
| Some (Jsont.Number (f, _)) -> Some f
371
+
| Init i -> Init.cwd i
347
-
let get_list t key =
348
-
match get_field t key with
349
-
| Some (Jsont.Array (items, _)) -> Some items
374
+
let error_msg = function
375
+
| Error e -> Some (Error.error e)
378
+
let to_json = function
380
+
(match Jsont.Json.encode Init.jsont i with
382
+
| Error msg -> failwith ("Init.to_json: " ^ msg))
384
+
(match Jsont.Json.encode Error.jsont e with
386
+
| Error msg -> failwith ("Error.to_json: " ^ msg))
387
+
| Other json -> json
355
-
let of_json json = json
389
+
let of_json ~subtype json =
392
+
(match Jsont.Json.decode Init.jsont json with
394
+
| Error _ -> Other json)
396
+
(match Jsont.Json.decode Error.jsont json with
398
+
| Error _ -> Other json)
···
let create ~subtype ~data = { subtype; data; unknown = Unknown.empty }
372
-
let make subtype data unknown = { subtype; data; unknown }
let subtype t = t.subtype
let unknown t = t.unknown
420
+
(** Create a system init message *)
421
+
let init ?session_id ?model ?cwd () =
422
+
{ subtype = "init";
423
+
data = Data.init ?session_id ?model ?cwd ();
424
+
unknown = Unknown.empty }
426
+
(** Create a system error message *)
428
+
{ subtype = "error";
429
+
data = Data.error ~error;
430
+
unknown = Unknown.empty }
(* Custom jsont that handles both formats:
- Old format: {"type":"system","subtype":"init","data":{...}}
- New format: {"type":"system","subtype":"init","cwd":"...","session_id":"...",...}
When data field is not present, we use the entire object as data *)
382
-
let make_with_optional_data subtype opt_data unknown_json =
383
-
let data = match opt_data with
385
-
| None -> unknown_json (* Use the full unknown object as data *)
438
+
(* First decode just the subtype *)
439
+
let subtype_codec = Jsont.Object.map ~kind:"SystemSubtype" Fun.id
440
+
|> Jsont.Object.mem "subtype" Jsont.string ~enc:Fun.id
441
+
|> Jsont.Object.finish
387
-
make subtype data Unknown.empty
443
+
match Jsont.Json.decode subtype_codec json with
444
+
| Error msg -> failwith ("System.jsont: " ^ msg)
446
+
(* Try to get data field, otherwise use full object *)
447
+
let data_codec = Jsont.Object.map ~kind:"SystemDataField" Fun.id
448
+
|> Jsont.Object.opt_mem "data" Jsont.json ~enc:Fun.id
449
+
|> Jsont.Object.finish
451
+
let data_json = match Jsont.Json.decode data_codec json with
455
+
let data = Data.of_json ~subtype data_json in
456
+
{ subtype; data; unknown = Unknown.empty }
389
-
Jsont.Object.map ~kind:"System" make_with_optional_data
390
-
|> Jsont.Object.mem "subtype" Jsont.string ~enc:subtype
391
-
|> Jsont.Object.opt_mem "data" Data.jsont ~enc:(fun t -> Some (data t))
392
-
|> Jsont.Object.keep_unknown Jsont.json_mems ~enc:(fun _ -> Unknown.empty)
393
-
|> Jsont.Object.finish
459
+
Jsont.Json.object' [
460
+
Jsont.Json.mem (Jsont.Json.name "type") (Jsont.Json.string "system");
461
+
Jsont.Json.mem (Jsont.Json.name "subtype") (Jsont.Json.string t.subtype);
462
+
Jsont.Json.mem (Jsont.Json.name "data") (Data.to_json t.data);
465
+
Jsont.map ~kind:"System" ~dec ~enc Jsont.json
397
-
(Jsont.Json.name "type", Jsont.String ("system", Jsont.Meta.none));
398
-
(Jsont.Json.name "subtype", Jsont.String (t.subtype, Jsont.Meta.none));
399
-
(Jsont.Json.name "data", Data.to_json t.data);
400
-
], Jsont.Meta.none)
468
+
match Jsont.Json.encode jsont t with
470
+
| Error msg -> failwith ("System.to_json: " ^ msg)
404
-
| Jsont.Object (fields, _) ->
405
-
let subtype = match List.assoc (Jsont.Json.name "subtype") fields with
406
-
| Jsont.String (s, _) -> s
407
-
| _ -> raise (Invalid_argument "System.of_json: invalid subtype")
409
-
let data = Data.of_json (
410
-
try List.assoc (Jsont.Json.name "data") fields
411
-
with Not_found -> Jsont.Object (fields, Jsont.Meta.none)
413
-
{ subtype; data; unknown = Unknown.empty }
414
-
| _ -> raise (Invalid_argument "System.of_json: expected object")
473
+
match Jsont.Json.decode jsont json with
475
+
| Error msg -> raise (Invalid_argument ("System.of_json: " ^ msg))
417
-
match t.subtype with
419
-
let session_id = Data.get_string t.data "session_id" in
420
-
let model = Data.get_string t.data "model" in
421
-
let cwd = Data.get_string t.data "cwd" in
Fmt.pf fmt "@[<2>System.init@ { session_id = %a;@ model = %a;@ cwd = %a }@]"
423
-
Fmt.(option string) session_id
424
-
Fmt.(option string) model
425
-
Fmt.(option string) cwd
427
-
let error = Data.get_string t.data "error" in
428
-
Fmt.pf fmt "@[<2>System.error@ { error = %a }@]"
429
-
Fmt.(option string) error
481
+
Fmt.(option string) (Init.session_id i)
482
+
Fmt.(option string) (Init.model i)
483
+
Fmt.(option string) (Init.cwd i)
485
+
Fmt.pf fmt "@[<2>System.error@ { error = %s }@]" (Error.error e)
Fmt.pf fmt "@[<2>System.%s@ { ... }@]" t.subtype
436
-
(* Opaque JSON type with typed accessors *)
437
-
type t = Jsont.json
492
+
module Unknown = struct
493
+
type t = Jsont.json
494
+
let empty = Jsont.Object ([], Jsont.Meta.none)
495
+
let is_empty = function Jsont.Object ([], _) -> true | _ -> false
496
+
let jsont = Jsont.json
439
-
let jsont = Jsont.json
500
+
input_tokens : int option;
501
+
output_tokens : int option;
502
+
total_tokens : int option;
503
+
cache_creation_input_tokens : int option;
504
+
cache_read_input_tokens : int option;
505
+
unknown : Unknown.t;
508
+
let make input_tokens output_tokens total_tokens
509
+
cache_creation_input_tokens cache_read_input_tokens unknown =
510
+
{ input_tokens; output_tokens; total_tokens;
511
+
cache_creation_input_tokens; cache_read_input_tokens; unknown }
let create ?input_tokens ?output_tokens ?total_tokens
?cache_creation_input_tokens ?cache_read_input_tokens () =
444
-
let fields = match input_tokens with
445
-
| Some n -> (Jsont.Json.name "input_tokens", Jsont.Number (float_of_int n, Jsont.Meta.none)) :: fields
446
-
| None -> fields in
447
-
let fields = match output_tokens with
448
-
| Some n -> (Jsont.Json.name "output_tokens", Jsont.Number (float_of_int n, Jsont.Meta.none)) :: fields
449
-
| None -> fields in
450
-
let fields = match total_tokens with
451
-
| Some n -> (Jsont.Json.name "total_tokens", Jsont.Number (float_of_int n, Jsont.Meta.none)) :: fields
452
-
| None -> fields in
453
-
let fields = match cache_creation_input_tokens with
454
-
| Some n -> (Jsont.Json.name "cache_creation_input_tokens", Jsont.Number (float_of_int n, Jsont.Meta.none)) :: fields
455
-
| None -> fields in
456
-
let fields = match cache_read_input_tokens with
457
-
| Some n -> (Jsont.Json.name "cache_read_input_tokens", Jsont.Number (float_of_int n, Jsont.Meta.none)) :: fields
458
-
| None -> fields in
459
-
Jsont.Object (fields, Jsont.Meta.none)
461
-
let get_field t key =
463
-
| Jsont.Object (members, _) ->
464
-
List.find_map (fun ((name, _), value) ->
465
-
if name = key then Some value else None
469
-
let get_int t key =
470
-
match get_field t key with
471
-
| Some (Jsont.Number (f, _)) ->
472
-
let i = int_of_float f in
473
-
if float_of_int i = f then Some i else None
515
+
{ input_tokens; output_tokens; total_tokens;
516
+
cache_creation_input_tokens; cache_read_input_tokens;
517
+
unknown = Unknown.empty }
476
-
let input_tokens t = get_int t "input_tokens"
519
+
let input_tokens t = t.input_tokens
520
+
let output_tokens t = t.output_tokens
521
+
let total_tokens t = t.total_tokens
522
+
let cache_creation_input_tokens t = t.cache_creation_input_tokens
523
+
let cache_read_input_tokens t = t.cache_read_input_tokens
524
+
let unknown t = t.unknown
478
-
let output_tokens t = get_int t "output_tokens"
480
-
let total_tokens t = get_int t "total_tokens"
482
-
let cache_creation_input_tokens t = get_int t "cache_creation_input_tokens"
484
-
let cache_read_input_tokens t = get_int t "cache_read_input_tokens"
526
+
let jsont : t Jsont.t =
527
+
Jsont.Object.map ~kind:"Usage" make
528
+
|> Jsont.Object.opt_mem "input_tokens" Jsont.int ~enc:input_tokens
529
+
|> Jsont.Object.opt_mem "output_tokens" Jsont.int ~enc:output_tokens
530
+
|> Jsont.Object.opt_mem "total_tokens" Jsont.int ~enc:total_tokens
531
+
|> Jsont.Object.opt_mem "cache_creation_input_tokens" Jsont.int ~enc:cache_creation_input_tokens
532
+
|> Jsont.Object.opt_mem "cache_read_input_tokens" Jsont.int ~enc:cache_read_input_tokens
533
+
|> Jsont.Object.keep_unknown Jsont.json_mems ~enc:unknown
534
+
|> Jsont.Object.finish
let effective_input_tokens t =
487
-
match input_tokens t with
537
+
match t.input_tokens with
490
-
let cached = Option.value (cache_read_input_tokens t) ~default:0 in
540
+
let cached = Option.value t.cache_read_input_tokens ~default:0 in
let total_cost_estimate t ~input_price ~output_price =
494
-
match input_tokens t, output_tokens t with
544
+
match t.input_tokens, t.output_tokens with
| Some input, Some output ->
let input_cost = float_of_int input *. input_price /. 1_000_000. in
let output_cost = float_of_int output *. output_price /. 1_000_000. in
···
Fmt.pf fmt "@[<2>Usage@ { input = %a;@ output = %a;@ total = %a;@ \
cache_creation = %a;@ cache_read = %a }@]"
504
-
Fmt.(option int) (input_tokens t)
505
-
Fmt.(option int) (output_tokens t)
506
-
Fmt.(option int) (total_tokens t)
507
-
Fmt.(option int) (cache_creation_input_tokens t)
508
-
Fmt.(option int) (cache_read_input_tokens t)
554
+
Fmt.(option int) t.input_tokens
555
+
Fmt.(option int) t.output_tokens
556
+
Fmt.(option int) t.total_tokens
557
+
Fmt.(option int) t.cache_creation_input_tokens
558
+
Fmt.(option int) t.cache_read_input_tokens
511
-
let of_json json = json
561
+
match Jsont.Json.encode jsont t with
563
+
| Error msg -> failwith ("Usage.to_json: " ^ msg)
566
+
match Jsont.Json.decode jsont json with
568
+
| Error msg -> raise (Invalid_argument ("Usage.of_json: " ^ msg))
···
Jsont.Object (fields, Jsont.Meta.none)
602
-
| Jsont.Object (fields, _) ->
603
-
let subtype = match List.assoc (Jsont.Json.name "subtype") fields with
604
-
| Jsont.String (s, _) -> s
605
-
| _ -> raise (Invalid_argument "Result.of_json: invalid subtype")
607
-
let duration_ms = match List.assoc (Jsont.Json.name "duration_ms") fields with
608
-
| Jsont.Number (f, _) -> int_of_float f
609
-
| _ -> raise (Invalid_argument "Result.of_json: invalid duration_ms")
611
-
let duration_api_ms = match List.assoc (Jsont.Json.name "duration_api_ms") fields with
612
-
| Jsont.Number (f, _) -> int_of_float f
613
-
| _ -> raise (Invalid_argument "Result.of_json: invalid duration_api_ms")
615
-
let is_error = match List.assoc (Jsont.Json.name "is_error") fields with
616
-
| Jsont.Bool (b, _) -> b
617
-
| _ -> raise (Invalid_argument "Result.of_json: invalid is_error")
619
-
let num_turns = match List.assoc (Jsont.Json.name "num_turns") fields with
620
-
| Jsont.Number (f, _) -> int_of_float f
621
-
| _ -> raise (Invalid_argument "Result.of_json: invalid num_turns")
623
-
let session_id = match List.assoc (Jsont.Json.name "session_id") fields with
624
-
| Jsont.String (s, _) -> s
625
-
| _ -> raise (Invalid_argument "Result.of_json: invalid session_id")
627
-
let total_cost_usd = match List.assoc_opt (Jsont.Json.name "total_cost_usd") fields with
628
-
| Some (Jsont.Number (f, _)) -> Some f
629
-
| Some _ -> raise (Invalid_argument "Result.of_json: invalid total_cost_usd")
632
-
let usage = Option.map Usage.of_json (List.assoc_opt (Jsont.Json.name "usage") fields) in
633
-
let result = match List.assoc_opt (Jsont.Json.name "result") fields with
634
-
| Some (Jsont.String (s, _)) -> Some s
635
-
| Some _ -> raise (Invalid_argument "Result.of_json: invalid result")
638
-
let structured_output = List.assoc_opt (Jsont.Json.name "structured_output") fields in
639
-
{ subtype; duration_ms; duration_api_ms; is_error; num_turns;
640
-
session_id; total_cost_usd; usage; result; structured_output; unknown = Unknown.empty }
641
-
| _ -> raise (Invalid_argument "Result.of_json: expected object")
658
+
match Jsont.Json.decode jsont json with
660
+
| Error msg -> raise (Invalid_argument ("Result.of_json: " ^ msg))
···
let system ~subtype ~data = System (System.create ~subtype ~data)
let system_init ~session_id =
686
-
let data = System.Data.of_assoc [(("session_id", Jsont.String (session_id, Jsont.Meta.none)))] in
687
-
System (System.create ~subtype:"init" ~data)
705
+
System (System.init ~session_id ())
let system_error ~error =
689
-
let data = System.Data.of_assoc [(("error", Jsont.String (error, Jsont.Meta.none)))] in
690
-
System (System.create ~subtype:"error" ~data)
707
+
System (System.error ~error)
let result ~subtype ~duration_ms ~duration_api_ms ~is_error ~num_turns
~session_id ?total_cost_usd ?usage ?result ?structured_output () =
···
| System t -> System.to_json t
| Result t -> Result.to_json t
705
-
| Jsont.Object (fields, _) -> (
706
-
match List.assoc_opt (Jsont.Json.name "type") fields with
707
-
| Some (Jsont.String ("user", _)) -> User (User.of_json json)
708
-
| Some (Jsont.String ("assistant", _)) -> Assistant (Assistant.of_json json)
709
-
| Some (Jsont.String ("system", _)) -> System (System.of_json json)
710
-
| Some (Jsont.String ("result", _)) -> Result (Result.of_json json)
711
-
| Some _ -> raise (Invalid_argument "Message.of_json: invalid type")
712
-
| None -> raise (Invalid_argument "Message.of_json: missing type field")
714
-
| _ -> raise (Invalid_argument "Message.of_json: expected object")
(* Jsont codec for the main Message variant type.
717
-
Uses a custom decoder to handle both old and new formats. *)
721
+
Uses case_mem for discriminated union based on "type" field. *)
719
-
Jsont.map ~kind:"Message" ~dec:of_json ~enc:to_json Jsont.json
723
+
let case_map kind obj dec = Jsont.Object.Case.map kind obj ~dec in
724
+
let case_user = case_map "user" User.incoming_jsont (fun v -> User v) in
725
+
let case_assistant = case_map "assistant" Assistant.incoming_jsont (fun v -> Assistant v) in
726
+
let case_system = case_map "system" System.jsont (fun v -> System v) in
727
+
let case_result = case_map "result" Result.jsont (fun v -> Result v) in
728
+
let enc_case = function
729
+
| User v -> Jsont.Object.Case.value case_user v
730
+
| Assistant v -> Jsont.Object.Case.value case_assistant v
731
+
| System v -> Jsont.Object.Case.value case_system v
732
+
| Result v -> Jsont.Object.Case.value case_result v
734
+
let cases = Jsont.Object.Case.[
736
+
make case_assistant;
740
+
Jsont.Object.map ~kind:"Message" Fun.id
741
+
|> Jsont.Object.case_mem "type" Jsont.string ~enc:Fun.id ~enc_case cases
742
+
~tag_to_string:Fun.id ~tag_compare:String.compare
743
+
|> Jsont.Object.finish
746
+
match Jsont.Json.decode jsont json with
748
+
| Error msg -> raise (Invalid_argument ("Message.of_json: " ^ msg))
| User t -> User.pp fmt t
···
let get_session_id = function
749
-
| System s when System.subtype s = "init" ->
750
-
System.Data.get_string (System.data s) "session_id"
778
+
| System s -> System.Data.session_id (System.data s)
| Result r -> Some (Result.session_id r)
anil.recoil.org