1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
type ('complete, 'partial) t = {
name : string;
response_format : Ai_provider.Mode.json_schema option;
parse_complete : string -> ('complete, string) result;
parse_partial : string -> 'partial option;
}
let partial_json_parse text =
match Partial_json.parse text with
| Some (json, _) -> Some json
| None -> None
let validate_schema ~schema json =
match Jsonschema.create_validator_from_json ~schema () with
| Error err -> Error (Format.asprintf "%a" Jsonschema.pp_compile_error err)
| Ok validator -> Jsonschema.validate validator json |> Result.map_error Jsonschema.Validation_error.to_string
let mode_of_output = function
| Some o ->
(match o.response_format with
| Some schema -> Ai_provider.Mode.Object_json (Some schema)
| None -> Ai_provider.Mode.Regular)
| None -> Ai_provider.Mode.Regular
let parse_output output steps =
match output with
| Some o ->
(match o.response_format with
| Some _ ->
let final_text = Generate_text_result.join_text steps in
(match o.parse_complete final_text with
| Ok json -> Some json
| Error _ -> None)
| None -> None)
| None -> None
let text =
{
name = "text";
response_format = None;
parse_complete = (fun s -> Ok s);
parse_partial =
(fun s ->
match String.length s with
| 0 -> None
| _ -> Some s);
}
let object_ ~name ~schema () =
let response_format = Some { Ai_provider.Mode.name; schema } in
let parse_complete text =
try
let json = Yojson.Basic.from_string text in
match validate_schema ~schema json with
| Ok () -> Ok json
| Error msg -> Error (Printf.sprintf "Schema validation failed: %s" msg)
with Yojson.Json_error msg -> Error (Printf.sprintf "Invalid JSON: %s" msg)
in
{ name; response_format; parse_complete; parse_partial = partial_json_parse }
let array ~name ~element_schema () =
let element_schema_clean =
match element_schema with
| `Assoc pairs -> `Assoc (List.filter (fun (k, _) -> not (String.equal k "$schema")) pairs)
| json -> json
in
let schema =
`Assoc
[
"$schema", `String "http://json-schema.org/draft-07/schema#";
"type", `String "object";
"properties", `Assoc [ "elements", `Assoc [ "type", `String "array"; "items", element_schema_clean ] ];
"required", `List [ `String "elements" ];
"additionalProperties", `Bool false;
]
in
let response_format = Some { Ai_provider.Mode.name; schema } in
let pairs =
match List.assoc_opt "elements" pairs with
| Some (`List elts) -> Some elts
| _ -> None
in
let validate_element elt =
match validate_schema ~schema:element_schema elt with
| Ok () -> true
| Error _ -> false
in
let parse_complete text =
try
match Yojson.Basic.from_string text with
| `Assoc pairs ->
(match extract_elements pairs with
| Some elts ->
let invalid = List.find_opt (fun elt -> not (validate_element elt)) elts in
(match invalid with
| Some elt -> Error (Printf.sprintf "Element validation failed: %s" (Yojson.Basic.to_string elt))
| None -> Ok (`List elts))
| None -> Error "missing or invalid 'elements' array")
| _ -> Error "expected JSON object with 'elements' array"
with Yojson.Json_error msg -> Error (Printf.sprintf "Invalid JSON: %s" msg)
in
let parse_partial text =
match Partial_json.parse text with
| None -> None
| Some (json, status) ->
match json with
| `Assoc pairs ->
(match extract_elements pairs with
| None -> None
| Some elts ->
let candidates =
match status with
| Partial_json.Repaired ->
(match List.rev elts with
| _ :: rest -> List.rev rest
| [] -> [])
| Partial_json.Successful -> elts
in
let valid = List.filter validate_element candidates in
Some (`List valid))
| _ -> None
in
{ name; response_format; parse_complete; parse_partial }
let choice ~name options =
let schema =
`Assoc
[
"$schema", `String "http://json-schema.org/draft-07/schema#";
"type", `String "object";
( "properties",
`Assoc
[ "result", `Assoc [ "type", `String "string"; "enum", `List (List.map (fun s -> `String s) options) ] ] );
"required", `List [ `String "result" ];
"additionalProperties", `Bool false;
]
in
let response_format = Some { Ai_provider.Mode.name; schema } in
let parse_complete text =
try
match Yojson.Basic.from_string text with
| `Assoc pairs as json ->
(match validate_schema ~schema json with
| Ok () ->
(match List.assoc_opt "result" pairs with
| Some value -> Ok value
| None -> Error "missing 'result' field in choice response")
| Error msg -> Error (Printf.sprintf "Schema validation failed: %s" msg))
| _ -> Error "expected JSON object with 'result' field"
with Yojson.Json_error msg -> Error (Printf.sprintf "Invalid JSON: %s" msg)
in
let parse_partial text =
match Partial_json.parse text with
| None -> None
| Some (json, status) ->
match json with
| `Assoc pairs ->
(match List.assoc_opt "result" pairs with
| Some (`String partial_result) ->
let potential_matches = List.filter (fun opt -> String.starts_with ~prefix:partial_result opt) options in
(match status with
| Partial_json.Successful ->
if List.exists (String.equal partial_result) options then Some (`String partial_result) else None
| Partial_json.Repaired ->
match potential_matches with
| [ single_match ] -> Some (`String single_match)
| _ -> None)
| _ -> None)
| _ -> None
in
{ name; response_format; parse_complete; parse_partial }
let enum = choice