llamacpp_embedding_nats/main.jl

using JSON, Dates, UUIDs, PrettyPrinting, Base64, DataFrames, NATS, HTTP, Base.Threads
using GeneralUtils, msghandler

# ---------------------------------------------- 100 --------------------------------------------- #
const config = JSON.parsefile("./appconfig.json")

# Initialize NATS client
nats_conn  = NATS.connect(config["nats_server_info"]["url"])

es = GeneralUtils.eventSink(natsConn=nats_conn, 
                            topic=config["thisServiceOutputChannel"]["status"], 
                            service_name=config["thisServiceName"])

function _log(msg::String)
  es(msg)
end

# -------------------------------------------- storage --------------------------------------------- #
const storage1 = GeneralUtils.GarageStorage(
  config["externalService"]["objectStorage"]["endpoint"],
  config["externalService"]["objectStorage"]["access_key_id"],
  config["externalService"]["objectStorage"]["secret_access_key"],
  config["externalService"]["objectStorage"]["bucket"]
)

function fileserverUploadHandler(_::String, dataname::String, data::Vector{UInt8})::Dict{String, Any}
  _, _ = GeneralUtils.s3PutFile(storage1, dataname, data)
  return Dict("status" => 200, "uploadid" => "", "fileid" => "", "url" => dataname)
end

function fileserverDownloadHandler(url::String, max_retries::Int, base_delay_ms::Int,
  max_delay_ms::Int, correlation_id::String)::Union{Vector{UInt8}, Nothing}
  delay = base_delay_ms
  for attempt in 1:max_retries
    data = GeneralUtils.s3GetFile(storage1, url)
    if data !== nothing
      GeneralUtils.s3DeleteFile(storage1, url)
      return data
    else
      @info "correlation_id: $correlation_id, Attempt $attempt"
      if attempt < max_retries
        sleep(delay / 1000.0)
        delay = min(delay * 2, max_delay_ms)
      end
    end
  end
  return nothing
end

# ----------------------------------------- healthserver ----------------------------------------- #
# function healthcheck(req::HTTP.Request)
#   return HTTP.Response(200, JSON.json(Dict("status"=>"ok")))
# end
# const ROUTER = HTTP.Router()
# HTTP.register!(ROUTER, "GET", "/ready", healthcheck)
# HTTP.register!(ROUTER, "GET", "/health", healthcheck)
# HTTP.register!(ROUTER, "GET", "/", healthcheck)
# healthserver = HTTP.serve!(ROUTER, "0.0.0.0", 3000)


# ------------------------------------------------------------------------------------------------ #
#                                    registration helper function                                  #
# ------------------------------------------------------------------------------------------------ #
"""
    register_service(nats_conn, service_name, topic, config)

Register a service with the load balancer via NATS request.

# Arguments
- `nats_conn`: NATS connection object
- `service_name`: Name of the service being registered
- `topic`: NATS topic/channel of the registering service
- `config`: Configuration dictionary containing service URLs

# Returns
- Response from the load balancer NATS request

Registers the service with the servicesloadbalancer topic using msghandler.smartpack
to create the registration message envelope.
"""
function register_service(nats_conn, service_name, topic, config)
  register_dict = Dict(
    "register_service_name" => service_name,
    "topic" => topic,
    "free_or_paid" => "free"
  )
  register_payloads = [("register", register_dict, "dictionary")]
  _, register_msg_envelope_json_str = msghandler.smartpack(
              config["externalService"]["servicesloadbalancer"]["url"], 
              register_payloads; 
              msg_purpose="register",   
              broker_url=config["nats_server_info"]["url"])
  try
    reply = NATS.request(nats_conn, 
                      config["externalService"]["servicesloadbalancer"]["url"], 
                      register_msg_envelope_json_str, 
                      timeout=60)
    _log("service $service_name is registered")  
  catch e
    _log("registering service $service_name failed: $(sprint(showerror, e))")    
  end
end

ch = Channel(1)
flag = Channel(1)

# use reply() here so that when this_service scale to >1 instance, 
# NATS will randomly select only one receiving instance
sub1 = NATS.subscribe(nats_conn, 
                config["thisServiceInputChannel"]["inbox"]; 
                queue_group="group1", spawn=true) do msg
      put!(ch, msg)
  end

_log("service started")

# ------------------------------------------------------------------------------------------------ #
#                                      this service main loop                                      #
# ------------------------------------------------------------------------------------------------ #
Threads.@spawn begin  # Register service periodically
  while true
    if isready(flag)
      register_service(nats_conn, 
                      config["thisServiceName"],
                      config["thisServiceInputChannel"]["inbox"], 
                      config)
      # _log("periodic registration sent")
    end
    sleep(60) # this line should make spawned thread sleep while allowing main thread running 
  end
end

function main()

# ---------------------------------- register service initially ---------------------------------- #
  register_service(nats_conn, 
                  config["thisServiceName"],
                  config["thisServiceInputChannel"]["inbox"], 
                  config)
# ---------------------------------------------- 100 --------------------------------------------- #
  _log("ready to accept NATS messages at broker=$(config["nats_server_info"]["url"]), topic=$(config["thisServiceInputChannel"]["inbox"])")
  while true
    try
      put!(flag, "ready") # ready for new request. do not send register message during llm working
      incoming_nats_msg = take!(ch) # keep this service running.

      # start working, the work may take more than a few minutes. pause async loop registeration to 
      # prevent a loadbalancer assigning new work to this service.
      _ = take!(flag)

      incoming_envelope_json_str = String(incoming_nats_msg.payload)
      incoming_envelope = msghandler.smartunpack(
        incoming_envelope_json_str; 
        fileserver_download_handler=fileserverDownloadHandler)
      if incoming_envelope["msg_purpose"] == "ping"
        _log("ping response")
        outgoing_envelope, outgoing_envelope_json_str = 
          msghandler.smartpack("requester inbox", 
                              [("response", "available", "text")]; 
                              broker_url=config["nats_server_info"]["url"],
                              fileserver_upload_handler=fileserverUploadHandler)
        NATS.publish(nats_conn, incoming_nats_msg.reply_to, outgoing_envelope_json_str)
      else
        _log("processing new request, msg_purpose=$(incoming_envelope["msg_purpose"])")
        dataname1 = incoming_envelope["payloads"][1][1]
        documents_dict = incoming_envelope["payloads"][1][2]
        documents = documents_dict["documents"]
        _log("received $(length(documents)) documents")
        body_data = Dict(
            "model" => "",
            "input" => documents
        )
        headers = [
            "Content-Type" => "application/json",
            "Authorization" => "Bearer no-key"
        ]
        http_response = HTTP.post(config["externalService"]["llamacpp_embedding"]["url"], 
                                headers, JSON.json(body_data))
        _log("HTTP response status=$(http_response.status)")

        # HTTP.jl doesn't immediately close connections after post() - it pools them for reuse.
        # Without explicit cleanup, connections accumulate and the llama.cpp server times out
        # on subsequent requests. This forces connection pool cleanup between requests.
        HTTP.close_idle_connections!()
        http_response_body = JSON.parse(http_response.body)
        dataname2 = "transferjob/$(UUIDs.uuid4())"
        outgoing_envelope, outgoing_envelope_json_str = 
          msghandler.smartpack("requester inbox", 
                              [(dataname2, http_response_body, "dictionary")]; 
                              msg_purpose=incoming_envelope["msg_purpose"],
                              sender_id=incoming_envelope["sender_id"],
                              broker_url=config["nats_server_info"]["url"],
                              fileserver_upload_handler=fileserverUploadHandler)
        # re-register this service to let a loadbalancer knows that this service is ready for next request
        register_service(nats_conn, 
                        config["thisServiceName"],
                        config["thisServiceInputChannel"]["inbox"], 
                        config)
        _log("re-registered for next request")
        reply_topic = 
        if incoming_nats_msg.reply_to !== nothing
          incoming_nats_msg.reply_to
        elseif incoming_envelope["reply_to"] != ""
          incoming_envelope["reply_to"]
        else
          _log("no reply_to found, using fallback topic")
          "sommpanion.unknown_subject"
        end
        NATS.publish(nats_conn, reply_topic, outgoing_envelope_json_str)
        _log("reply published to $reply_topic, dataname2=$dataname2")
      end
    catch e
      bt = catch_backtrace()
      err_msg = sprint() do io
          showerror(io, e, bt)
          println(io)
      end

      _log(err_msg)
      # println(err_msg)
    end
  end
end

main()