lib/medical_transcription/transcription_server.ex

defmodule Medicode.TranscriptionServer do
  @moduledoc """
  GenServer responsible for transcribing audio files
  """
  use GenServer

  alias Medicode.Transcriptions
  alias Medicode.Transcriptions.Transcription

  @registry :transcription_registry

  def start_link(%{transcription: transcription, name: name}) do
    GenServer.start_link(__MODULE__, {:transcription, transcription}, name: via_tuple(name))
  end

  @doc """
  This function will be called by the supervisor to retrieve the specification 
  of the child process.The child process is configured to restart only if it 
  terminates abnormally.
  """
  def child_spec(process_name) do
    %{
      id: __MODULE__,
      start: {__MODULE__, :start_link, [process_name]},
      restart: :transient
    }
  end

  @impl GenServer
  def init(init_arg) do
    # NOTE: The rule is: don't do anything slow or risky in your GenServer's init function.
    # But that isn't always practical. GenServers have a reasonably elegant solution to this: handle_continue/2.
    # We can change our init function to return {:ok, INTIAL_STATE, {:continue, CONTINUE_TYPE}} which will both
    # unblock the initialization and guarantee that handle_continue/2 is called before any other message is processed.
    #
    # Explanation from: https://www.openmymind.net/Elixir-A-Little-Beyond-The-Basics-Part-8-genservers
    {:ok, init_arg, {:continue, :start}}
  end

  @impl GenServer
  def handle_continue(:start, {:transcription, transcription}) do
    {:ok, transcription} =
      Transcriptions.update_transcription(transcription, %{status: :transcribing})

    Phoenix.PubSub.broadcast(
      :medicode_pubsub,
      "transcriptions:#{transcription.id}",
      {:transcription_started, transcription.id}
    )

    stream_transcription_and_search(transcription.filename)

    {:noreply, {:transcription, transcription}}
  end

  @impl GenServer
  def handle_info({:chunk, result}, state) do
    {:transcription, transcription} = state

    %Transcription{id: id} = transcription

    {:ok, chunk} =
      Transcriptions.create_chunk(%{
        transcription_id: id,
        text: String.trim(result.text),
        text_vector: Medicode.Coding.compute_vector_as_list(result.text),
        start_mark: result.start_mark,
        end_mark: result.end_mark
      })

    Phoenix.PubSub.broadcast(
      :medicode_pubsub,
      "transcriptions:#{id}",
      {:transcription_updated, chunk}
    )

    Medicode.ClassificationSupervisor.start_classification(chunk)

    {:noreply, state}
  end

  def handle_info(:finished, state) do
    {:stop, :normal, state}
  end

  @impl GenServer
  def terminate(reason, state) do
    {:transcription, transcription} = state

    {:ok, transcription} =
      Transcriptions.update_transcription(transcription, %{status: :finished})

    %Transcription{id: id} = transcription

    Phoenix.PubSub.broadcast(
      :medicode_pubsub,
      "transcriptions:#{id}",
      {:transcription_finished, reason}
    )

    reason
  end

  # Ideas for future exploration:
  # - Instead of storing the long description vectors in a binary file on disk, we could store them within a vector DB
  #   (such as pgvector or Pinecone.io)
  # - A potential improvement would be to not code each chunk of transcribed audio separately, but to instead gather
  #   complete sentences based on punctuation. We may want to suggest codes for the entire audio as a single piece as
  #   well
  defp stream_transcription_and_search(audio_file_path) do
    # audio transcription + semantic search
    Medicode.TranscriptionServing
    |> Nx.Serving.batched_run({:file, audio_file_path})
    |> Enum.each(fn chunk ->
      result = %{
        start_mark: format_timestamp(chunk.start_timestamp_seconds),
        end_mark: format_timestamp(chunk.end_timestamp_seconds),
        text: chunk.text
      }

      send(self(), {:chunk, result})
    end)

    send(self(), :finished)
  end

  defp format_timestamp(seconds) when is_nil(seconds), do: nil

  defp format_timestamp(seconds) do
    seconds |> round() |> Time.from_seconds_after_midnight() |> Time.to_string()
  end

  defp via_tuple(name),
    do: {:via, Registry, {@registry, name}}
end