uncloseai.com/languages/ocaml/uncloseai.ml

(* UncloseAI - OCaml client for OpenAI-compatible APIs with streaming support *)

open Lwt
open Cohttp
open Cohttp_lwt_unix

type model_info = {
  id : string;
  endpoint : string;
  max_tokens : int;
}

type chat_message = {
  role : string;
  content : string;
}

type t = {
  models : model_info list;
  tts_endpoints : string list;
  api_key : string option;
  timeout : int;
  debug : bool;
}

let discover_env_endpoints prefix =
  let rec discover i acc =
    if i >= 10000 then List.rev acc
    else
      match Sys.getenv_opt (Printf.sprintf "%s_%d" prefix i) with
      | None -> List.rev acc
      | Some endpoint -> discover (i + 1) (endpoint :: acc)
  in
  discover 1 []

let discover_models client endpoints =
  let models = ref [] in
  let discover_from_endpoint endpoint =
    if client.debug then
      Printf.printf "[DEBUG] Discovering from: %s\n%!" endpoint;
    Lwt.catch
      (fun () ->
        let url = Printf.sprintf "%s/models" endpoint in
        Client.get (Uri.of_string url) >>= fun (_resp, body) ->
        Cohttp_lwt.Body.to_string body >|= fun body_str ->
        let json = Yojson.Basic.from_string body_str in
        let open Yojson.Basic.Util in
        let model_list = json |> member "data" |> to_list in
        List.iter (fun model ->
          let model_id = model |> member "id" |> to_string in
          (* Skip permission entries *)
          if not (String.starts_with ~prefix:"modelperm-" model_id ||
                  String.starts_with ~prefix:"chatcmpl-" model_id) then begin
            let max_tokens =
              try model |> member "max_model_len" |> to_int
              with _ -> 8192
            in
            models := { id = model_id; endpoint; max_tokens } :: !models;
            if client.debug then
              Printf.printf "[DEBUG] Discovered: %s\n%!" model_id
          end
        ) model_list)
      (fun _exn ->
        if client.debug then
          Printf.printf "[DEBUG] Error discovering from %s\n%!" endpoint;
        Lwt.return_unit)
  in
  Lwt_main.run (Lwt_list.iter_s discover_from_endpoint endpoints);
  List.rev !models

let create ?(model_endpoints=[]) ?(tts_endpoints=[]) ?(api_key=None) ?(timeout=30000) ?(debug=false) () =
  let model_ends = if List.length model_endpoints > 0 then model_endpoints
                   else discover_env_endpoints "MODEL_ENDPOINT" in
  let tts_ends = if List.length tts_endpoints > 0 then tts_endpoints
                 else discover_env_endpoints "TTS_ENDPOINT" in

  if debug then
    Printf.printf "[DEBUG] Initialized with %d endpoint(s)\n%!" (List.length model_ends);

  let client = {
    models = [];
    tts_endpoints = tts_ends;
    api_key;
    timeout;
    debug;
  } in

  let models = discover_models client model_ends in
  { client with models }

let list_models client = client.models

let resolve_model client model_id =
  if List.length client.models = 0 then
    failwith "No models available"
  else if model_id = "" then
    List.hd client.models
  else
    try
      List.find (fun m -> m.id = model_id) client.models
    with Not_found ->
      failwith (Printf.sprintf "Model '%s' not found" model_id)

let chat client messages ?(model="") ?(max_tokens=100) ?(temperature=0.7) () =
  let model_info = resolve_model client model in

  let messages_json = `List (List.map (fun msg ->
    `Assoc [("role", `String msg.role); ("content", `String msg.content)]
  ) messages) in

  let payload = `Assoc [
    ("model", `String model_info.id);
    ("messages", messages_json);
    ("max_tokens", `Int max_tokens);
    ("temperature", `Float temperature);
    ("stream", `Bool false)
  ] in

  let body = Yojson.Basic.to_string payload |> Cohttp_lwt.Body.of_string in
  let headers = Header.init ()
    |> fun h -> Header.add h "Content-Type" "application/json" in
  let headers = match client.api_key with
    | Some key -> Header.add headers "Authorization" (Printf.sprintf "Bearer %s" key)
    | None -> headers
  in

  let url = Printf.sprintf "%s/chat/completions" model_info.endpoint in
  Client.post ~headers ~body (Uri.of_string url) >>= fun (_resp, body) ->
  Cohttp_lwt.Body.to_string body >|= fun body_str ->
  Yojson.Basic.from_string body_str

let chat_stream client messages ?(model="") ?(max_tokens=500) ?(temperature=0.7) callback =
  let model_info = resolve_model client model in

  let messages_json = `List (List.map (fun msg ->
    `Assoc [("role", `String msg.role); ("content", `String msg.content)]
  ) messages) in

  let payload = `Assoc [
    ("model", `String model_info.id);
    ("messages", messages_json);
    ("max_tokens", `Int max_tokens);
    ("temperature", `Float temperature);
    ("stream", `Bool true)
  ] in

  let body = Yojson.Basic.to_string payload |> Cohttp_lwt.Body.of_string in
  let headers = Header.init ()
    |> fun h -> Header.add h "Content-Type" "application/json"
    |> fun h -> Header.add h "Accept" "text/event-stream" in
  let headers = match client.api_key with
    | Some key -> Header.add headers "Authorization" (Printf.sprintf "Bearer %s" key)
    | None -> headers
  in

  let url = Printf.sprintf "%s/chat/completions" model_info.endpoint in
  Lwt.catch
    (fun () ->
      Client.post ~headers ~body (Uri.of_string url) >>= fun (_resp, body) ->
      let stream = Cohttp_lwt.Body.to_stream body in
      let buffer = ref "" in

      Lwt_stream.iter_s (fun chunk ->
        buffer := !buffer ^ chunk;
        let lines = String.split_on_char '\n' !buffer in
        let rec process_lines = function
          | [] -> Lwt.return_unit
          | [last] ->
              buffer := last;
              Lwt.return_unit
          | line :: rest ->
              let trimmed = String.trim line in
              if String.starts_with ~prefix:"data: " trimmed then begin
                let data = String.sub trimmed 6 (String.length trimmed - 6) in
                let data = String.trim data in
                if data = "[DONE]" then
                  Lwt.return_unit
                else begin
                  try
                    let chunk = Yojson.Basic.from_string data in
                    let open Yojson.Basic.Util in
                    let choices = chunk |> member "choices" |> to_list in
                    if List.length choices > 0 then begin
                      let delta = List.hd choices |> member "delta" in
                      try
                        let content = delta |> member "content" |> to_string in
                        if String.length content > 0 then
                          callback content
                      with _ -> ()
                    end;
                    process_lines rest
                  with _ ->
                    if client.debug then
                      Printf.printf "[DEBUG] Parse error\n%!";
                    process_lines rest
                end
              end else
                process_lines rest
        in
        process_lines lines
      ) stream
    )
    (fun _exn ->
      if client.debug then
        Printf.printf "[DEBUG] Stream error\n%!";
      Lwt.return_unit)

let tts client text ?(voice="alloy") ?(model="tts-1") ?(response_format="mp3") () =
  if List.length client.tts_endpoints = 0 then
    failwith "No TTS endpoints available"
  else
    let endpoint = List.hd client.tts_endpoints in

    let payload = `Assoc [
      ("model", `String model);
      ("voice", `String voice);
      ("input", `String text);
      ("response_format", `String response_format)
    ] in

    let body = Yojson.Basic.to_string payload |> Cohttp_lwt.Body.of_string in
    let headers = Header.init ()
      |> fun h -> Header.add h "Content-Type" "application/json" in
    let headers = match client.api_key with
      | Some key -> Header.add headers "Authorization" (Printf.sprintf "Bearer %s" key)
      | None -> headers
    in

    let url = Printf.sprintf "%s/audio/speech" endpoint in
    Client.post ~headers ~body (Uri.of_string url) >>= fun (_resp, body) ->
    Cohttp_lwt.Body.to_string body

(* Demo when run as main module *)
let () =
  Printf.printf "=== UncloseAI OCaml Client (with Streaming) ===\n\n%!";

  let client = create ~debug:true () in

  let models = list_models client in
  if List.length models = 0 then begin
    Printf.printf "ERROR: No models discovered. Set environment variables:\n";
    Printf.printf "  MODEL_ENDPOINT_1, MODEL_ENDPOINT_2, etc.\n%!";
    exit 1
  end;

  Printf.printf "\nDiscovered %d model(s):\n%!" (List.length models);
  List.iter (fun m ->
    Printf.printf "  - %s (max_tokens: %d)\n%!" m.id m.max_tokens
  ) models;
  Printf.printf "\n%!";

  (* Non-streaming chat *)
  Printf.printf "=== Non-Streaming Chat ===\n%!";
  Lwt_main.run (
    Lwt.catch
      (fun () ->
        chat client [
          {role="system"; content="You are a helpful AI assistant."};
          {role="user"; content="Explain quantum computing in one sentence."}
        ] () >>= fun response ->
        let open Yojson.Basic.Util in
        let content = response |> member "choices" |> to_list |> List.hd
                      |> member "message" |> member "content" |> to_string in
        Printf.printf "Response: %s\n\n%!" content;
        Lwt.return_unit)
      (fun exn ->
        Printf.printf "Error: %s\n\n%!" (Printexc.to_string exn);
        Lwt.return_unit)
  );

  (* Streaming chat *)
  Printf.printf "=== Streaming Chat ===\n%!";
  let model_id = if List.length models > 1 then (List.nth models 1).id else "" in
  let model_name = if model_id = "" then (List.hd models).id else model_id in
  Printf.printf "Model: %s\n%!" model_name;
  Printf.printf "Response: %!";

  Lwt_main.run (
    Lwt.catch
      (fun () ->
        chat_stream client [
          {role="system"; content="You are a coding assistant."};
          {role="user"; content="Write an OCaml function to check if a number is prime"}
        ] ~model:model_id ~max_tokens:200 (fun content ->
          Printf.printf "%s%!" content
        ) >>= fun () ->
        Printf.printf "\n\n%!";
        Lwt.return_unit)
      (fun exn ->
        Printf.printf "\nError: %s\n\n%!" (Printexc.to_string exn);
        Lwt.return_unit)
  );

  (* TTS *)
  if List.length client.tts_endpoints > 0 then begin
    Printf.printf "=== TTS Speech Generation ===\n%!";
    Lwt_main.run (
      Lwt.catch
        (fun () ->
          tts client "Hello from UncloseAI OCaml client! This demonstrates streaming support." () >>= fun audio_data ->
          let oc = open_out_bin "speech.mp3" in
          output_string oc audio_data;
          close_out oc;
          Printf.printf "[OK] Speech file created: speech.mp3 (%d bytes)\n\n%!" (String.length audio_data);
          Lwt.return_unit)
        (fun exn ->
          Printf.printf "[ERROR] TTS Error: %s\n\n%!" (Printexc.to_string exn);
          Lwt.return_unit)
    )
  end;

  Printf.printf "=== Examples Complete ===\n%!"