Mix.install(
  [
    {:finch, "~> 0.16.0"},
    {:jason, "~> 1.4"}
  ],
)

Finch.start_link(name: OllamaFinch)


defmodule Ollama do
  @api_endpoint "http://localhost:11434/api/generate"
  @model "llama2-uncensored"

  def stream_response(prompt, fun) do
    payload =
      %{
        model: @model,
        prompt: prompt
      }
      |> Jason.encode!()

    process_fn = fn
      {:status, _status}, acc ->
        acc

      {:headers, _headers}, acc ->
        acc

      {:data, json_string}, acc ->
        data = Jason.decode!(json_string)
        process_chunk(data, acc, fun)
    end

    Finch.build(:post, @api_endpoint, [{"content-type", "application/json"}], payload)
    |> Finch.stream(OllamaFinch, "", process_fn)
  end

  defp process_chunk(%{"response" => new_content, "done" => false}, acc, fun) do
    updated_content = acc <> new_content

    updated_content
    |> tap(fun)
  end

  defp process_chunk(%{"done" => true}, acc, _fun) do
    acc
  end
end