Model Context Protocol in OCaml
anil.recoil.org
1open Mcp
2open Mcp_sdk
3open Mcp_server
4
5(* Helper for extracting string value from JSON *)
6let get_string_param json name =
7 match json with
8 | `Assoc fields ->
9 (match List.assoc_opt name fields with
10 | Some (`String value) -> value
11 | _ -> raise (Failure (Printf.sprintf "Missing or invalid parameter: %s" name)))
12 | _ -> raise (Failure "Expected JSON object")
13
14(* Create a server *)
15let server = create_server
16 ~name:"OCaml MCP Completion Example"
17 ~version:"0.1.0"
18 ~protocol_version:"2024-11-05"
19 ()
20
21(* Define startup and shutdown hooks *)
22let startup () =
23 (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
24 Printf.fprintf stderr "CompletionServer is starting up!\n";
25 flush stderr;
26 Log.info "CompletionServer is starting up!"
27
28let shutdown () =
29 Printf.fprintf stderr "CompletionServer is shutting down. Goodbye!\n";
30 flush stderr;
31 Log.info "CompletionServer is shutting down. Goodbye!"
32
33(* Register the hooks *)
34let () =
35 set_startup_hook server startup;
36 set_shutdown_hook server shutdown
37
38(* Database of programming languages and their features *)
39let languages = [
40 ("ocaml", ["functional"; "static typing"; "pattern matching"; "modules"; "type inference"]);
41 ("python", ["dynamic typing"; "interpreted"; "object-oriented"; "high-level"; "scripting"]);
42 ("rust", ["memory safety"; "performance"; "static typing"; "ownership"; "zero-cost abstractions"]);
43 ("javascript", ["dynamic typing"; "interpreted"; "prototypes"; "single-threaded"; "event-driven"]);
44 ("go", ["concurrency"; "garbage collection"; "simplicity"; "static typing"; "compiled"]);
45]
46
47(* Helper function to create a completion response *)
48let create_completion values ?(has_more=false) ?(total=None) () =
49 Completion.Result.{
50 completion = {
51 values;
52 has_more = Some has_more;
53 total;
54 };
55 meta = None;
56 }
57
58(* Define and register a tool that handles completions *)
59let _ = add_tool server
60 ~name:"complete"
61 ~description:"Handles completion requests for programming languages and features"
62 ~schema_properties:[
63 ("argument_name", "string", "The name of the argument to complete");
64 ("argument_value", "string", "The partial value to complete");
65 ]
66 ~schema_required:["argument_name"; "argument_value"]
67 (fun args ->
68 try
69 let argument_name = get_string_param args "argument_name" in
70 let argument_value = get_string_param args "argument_value" in
71
72 Log.info (Printf.sprintf "Completion request for %s = %s" argument_name argument_value);
73
74 (* Handle different completion requests *)
75 let result =
76 match argument_name with
77 | "language" ->
78 (* Complete programming language names *)
79 let matches =
80 List.filter
81 (fun (lang, _) ->
82 let lang_lower = String.lowercase_ascii lang in
83 let arg_lower = String.lowercase_ascii argument_value in
84 String.length lang_lower >= String.length arg_lower &&
85 String.sub lang_lower 0 (String.length arg_lower) = arg_lower)
86 languages
87 in
88 let values = List.map fst matches in
89 create_completion values ~has_more:false ~total:(Some (List.length values)) ()
90
91 | "feature" ->
92 (* Complete programming language features *)
93 let all_features =
94 List.flatten (List.map snd languages) |>
95 List.sort_uniq String.compare
96 in
97 let matches =
98 List.filter
99 (fun feature ->
100 let feature_lower = String.lowercase_ascii feature in
101 let arg_lower = String.lowercase_ascii argument_value in
102 String.length feature_lower >= String.length arg_lower &&
103 String.sub feature_lower 0 (String.length arg_lower) = arg_lower)
104 all_features
105 in
106 create_completion matches ~has_more:false ~total:(Some (List.length matches)) ()
107
108 | _ ->
109 (* Default completions for unknown arguments *)
110 create_completion ["unknown argument"] ()
111 in
112
113 (* Convert to JSON and return *)
114 TextContent.yojson_of_t TextContent.{
115 text = Yojson.Safe.to_string (Completion.Result.to_result result);
116 annotations = None
117 }
118 with
119 | Failure msg ->
120 Log.error (Printf.sprintf "Error handling completion request: %s" msg);
121 TextContent.yojson_of_t TextContent.{
122 text = Printf.sprintf "Error: %s" msg;
123 annotations = None
124 }
125 )
126
127(* Define and register a prompt that provides programming language info *)
128let _ = add_prompt server
129 ~name:"language-info-prompt"
130 ~description:"A prompt that provides information about programming languages"
131 ~arguments:[
132 ("language", Some "Name of the programming language", true);
133 ]
134 (fun args ->
135 let language =
136 try List.assoc "language" args
137 with Not_found -> "ocaml" (* Default to OCaml *)
138 in
139
140 let features =
141 try
142 let features = List.assoc (String.lowercase_ascii language) languages in
143 String.concat ", " features
144 with Not_found -> "unknown language"
145 in
146
147 [
148 Prompt.{
149 role = `User;
150 content = make_text_content (Printf.sprintf "Tell me about the %s programming language" language)
151 };
152 Prompt.{
153 role = `Assistant;
154 content = make_text_content (Printf.sprintf "%s is a programming language with the following features: %s" language features)
155 }
156 ]
157 )
158
159(* Main function *)
160let () =
161 (* Parse command line arguments *)
162 let transport_type = ref Stdio in
163 let args = [
164 ("--http", Arg.Unit (fun () -> transport_type := Http),
165 "Start server with HTTP transport (default is stdio)");
166 ] in
167 let usage_msg = "Usage: completion_example [--http]" in
168 Arg.parse args (fun _ -> ()) usage_msg;
169
170 (* Use stderr for direct printing to avoid interfering with JSON-RPC protocol *)
171 Printf.fprintf stderr "Starting CompletionServer...\n";
172 flush stderr;
173 Log.info "Starting CompletionServer...";
174
175 (* Set custom capabilities to indicate support for completions *)
176 let capabilities = `Assoc [
177 ("completions", `Assoc []); (* Indicate support for completions *)
178 ("prompts", `Assoc [
179 ("listChanged", `Bool true)
180 ]);
181 ("resources", `Assoc [
182 ("listChanged", `Bool true);
183 ("subscribe", `Bool true)
184 ]);
185 ("tools", `Assoc [
186 ("listChanged", `Bool true)
187 ])
188 ] in
189 set_capabilities server capabilities;
190
191 (* Create and start MCP server with the selected transport *)
192 let mcp_server = create ~server ~transport:!transport_type () in
193 start mcp_server