mediawiki_client

defmodule Wiki.LiftWing do
  @moduledoc """
  This module provides an adapter for the [Lift Wing API](https://api.wikimedia.org/wiki/Lift_Wing_API)
  for machine learning predictions.

  See this page for a catalog of models, https://wikitech.wikimedia.org/wiki/Machine_Learning/LiftWing\#Current_Inference_Services

  Specific API methods have only been written for a handful of the models,
  please follow the pattern below to add additional models, and please
  contribute back here!

  The server currently rate-limits requests to 50,000 per hour per IP address.
  """

  alias Wiki.{Error, Util}

  @type client_option ::
          {:adapter, module()}
          | {:debug, true}
          | {:endpoint, binary()}
          | {:user_agent, binary()}
  @typedoc """
    - `:adapter` - Override the HTTP adapter
    - `:debug` - Turn on verbose logging by setting to `true`
    - `:endpoint` - Override the base URL to query
    - `:user_agent` - Override the user-agent header string
  """
  @type client_options :: [client_option()]
  @default_adapter Tesla.Adapter.Hackney
  @default_endpoint "https://api.wikimedia.org/service/lw/inference/v1/models/"

  @doc """
  Create a new Lift Wing client.

  ## Arguments

  - `opts` - Configuration options that can change client behavior

  ## Return value

  Returns an opaque client object, which should be passed to `request/2`.
  """
  @spec new(client_options()) :: Tesla.Client.t()
  def new(opts \\ []) do
    endpoint = opts[:endpoint] || @default_endpoint

    client(endpoint, opts)
  end

  @doc """
  Revscoring inference for the specific wikis listed on https://api.wikimedia.org/wiki/Lift_Wing_API/Reference/Get_revscoring_articlequality_prediction

  ```
  Wiki.LiftWing.new()
  |> Wiki.LiftWing.revscoring_articlequality("enwiki", 12345)
  # {:ok,
  #  %{
  #    "enwiki" => %{
  #      "models" => %{"articlequality" => %{"version" => "0.9.2"}},
  #      "scores" => %{
  #        "12345" => %{
  #          "articlequality" => %{
  #            "score" => %{
  #              "prediction" => "Stub",
  #              "probability" => %{
  #                "B" => 0.04109666150848109,
  #                "C" => 0.02060738177009099,
  #                "FA" => 0.0029400688910391592,
  #                "GA" => 0.004970401857774162,
  #                "Start" => 0.17362493306327959,
  #                "Stub" => 0.7567605529093352
  #              }
  #            }
  #          }
  #        }
  #      }
  #    }
  #  }}
  ```
  """
  @spec revscoring_articlequality(Tesla.Client.t(), binary(), integer(), boolean()) ::
          {:ok, map} | {:error, any()}
  def revscoring_articlequality(client, wiki, revid, with_features \\ false) do
    model =
      case wiki do
        "wikidatawiki" -> "itemquality"
        _ -> "articlequality"
      end

    request(client,
      method: :post,
      url: "#{wiki}-#{model}:predict",
      body: Jason.encode!(Enum.into([rev_id: revid, extended_output: with_features], %{}))
    )
  end

  @doc """
  Language-agnostic article quality as documented on https://api.wikimedia.org/wiki/Lift_Wing_API/Reference/Get_language_agnostic_articlequality_prediction

  ```
  Wiki.LiftWing.new()
  |> Wiki.LiftWing.articlequality("en", 123456)
  # {:ok,
  #  %{
  #    "model_name" => "articlequality",
  #    "model_version" => "1",
  #    "revision_id" => 123456,
  #    "score" => 0.10391570042067477,
  #    "wiki_db" => "enwiki"
  #  }}
  ```
  """
  @spec articlequality(Tesla.Client.t(), binary(), integer(), boolean()) ::
          {:ok, map} | {:error, any()}
  def articlequality(client, lang, revid, with_features \\ false) do
    request(client,
      method: :post,
      url: "/articlequality:predict",
      body: Jason.encode!(%{lang: lang, rev_id: revid, extended_output: to_string(with_features)})
    )
  end

  @doc """
  Article topic prediction

  ```
  Wiki.LiftWing.new()
  |> Wiki.LiftWing.articletopic("enwiki", 12345)
  # {:ok,
  #  %{
  #    "enwiki" => %{
  #      "models" => %{"articletopic" => %{"version" => "1.3.0"}},
  #      "scores" => %{
  #        "12345" => %{
  #          "articletopic" => %{
  #            "score" => %{
  #              "prediction" => ["STEM.STEM*"],
  #              "probability" => %{
  #                "Culture.Internet culture" => 0.0027448342044452587,
  #                "Culture.Media.Books" => 7.42678990345212e-4,
  #                "Geography.Regions.Europe.Western Europe" => 0.001866188163881686,
  #                "Culture.Visual arts.Comics and Anime" => 3.67749742655391e-4,
  #                "STEM.STEM*" => 0.9049354239597514,
  #                ....
  ```
  """
  @spec articletopic(Tesla.Client.t(), binary(), integer(), boolean()) ::
          {:ok, map} | {:error, any()}
  def articletopic(client, wiki, revid, with_features \\ false) do
    request(client,
      method: :post,
      url: "/#{wiki}-articletopic:predict",
      body: Jason.encode!(%{rev_id: revid, extended_output: with_features})
    )
  end

  @doc """
  Make a Lift Wing request.

  ## Arguments

  - `client` - Client object as returned by `new/1`.
  - Remaining parameters are passed to `Tesla.request`
  """
  @spec request(Tesla.Client.t(), keyword) :: {:ok, map} | {:error, any()}
  def request(client, opts) do
    with {:ok, response} <- Tesla.request(client, opts),
         {:ok, result} <- validate(response) do
      {:ok, result.body}
    end
  end

  defp validate(result) do
    with nil <- validate_body_type(result.body),
         nil <- validate_api_errors(result.body),
         nil <- validate_http_status(result.status) do
      {:ok, result}
    end
  end

  defp validate_http_status(status) do
    case status do
      status when status >= 200 and status < 300 -> nil
      status -> {:error, %Error{message: "Error received with HTTP status #{status}"}}
    end
  end

  defp validate_body_type(body) do
    with body when is_map(body) <- body,
         body when body != %{} <- body do
      nil
    else
      _ -> {:error, %Error{message: "Empty response"}}
    end
  end

  defp validate_api_errors(body) do
    case body["error"] do
      nil -> nil
      error -> {:error, %Error{message: summarize_error(error)}}
    end
  end

  defp summarize_error(error) when is_binary(error),
    do: error

  @spec client(binary(), client_options()) :: Tesla.Client.t()
  defp client(url, opts) do
    adapter = opts[:adapter] || @default_adapter
    user_agent = opts[:user_agent] || Util.default_user_agent()

    ([
       {Tesla.Middleware.BaseUrl, url},
       {Tesla.Middleware.Headers,
        [
          {"user-agent", user_agent}
        ]},
       Tesla.Middleware.FollowRedirects,
       Tesla.Middleware.JSON
     ] ++
       if(opts[:debug], do: [Tesla.Middleware.Logger], else: []))
    |> Tesla.client(adapter)
  end
end