pleroma/lib/pleroma/search/qdrant_search.ex

defmodule Pleroma.Search.QdrantSearch do
  @behaviour Pleroma.Search.SearchBackend
  import Ecto.Query

  alias Pleroma.Activity
  alias Pleroma.Config.Getting, as: Config
  alias Pleroma.Object
  alias Pleroma.Search

  alias __MODULE__.OpenAIClient
  alias __MODULE__.QdrantClient

  import Pleroma.Search.DatabaseSearch, only: [maybe_fetch: 3]

  @impl true
  def create_index do
    payload = Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration])

    with {:ok, %{status: 200}} <- QdrantClient.put("/collections/posts", payload) do
      :ok
    else
      e -> {:error, e}
    end
  end

  @impl true
  def drop_index do
    with {:ok, %{status: 200}} <- QdrantClient.delete("/collections/posts") do
      :ok
    else
      e -> {:error, e}
    end
  end

  def get_embedding(text) do
    with {:ok, %{body: %{"data" => [%{"embedding" => embedding}]}}} <-
           OpenAIClient.post("/v1/embeddings", %{
             input: text,
             model: Config.get([Pleroma.Search.QdrantSearch, :openai_model])
           }) do
      {:ok, embedding}
    else
      _ ->
        {:error, "Failed to get embedding"}
    end
  end

  defp actor_from_activity(%{data: %{"actor" => actor}}) do
    actor
  end

  defp actor_from_activity(_), do: nil

  defp build_index_payload(activity, embedding) do
    actor = actor_from_activity(activity)
    published_at = activity.data["published"]

    %{
      points: [
        %{
          id: activity.id |> FlakeId.from_string() |> Ecto.UUID.cast!(),
          vector: embedding,
          payload: %{actor: actor, published_at: published_at}
        }
      ]
    }
  end

  defp build_search_payload(embedding, options) do
    base = %{
      vector: embedding,
      limit: options[:limit] || 20,
      offset: options[:offset] || 0
    }

    if author = options[:author] do
      Map.put(base, :filter, %{
        must: [%{key: "actor", match: %{value: author.ap_id}}]
      })
    else
      base
    end
  end

  @impl true
  def add_to_index(%Activity{object: %Object{} = object} = activity) do
    search_data = Search.object_to_search_data(object)

    with {:ok, embedding} <- get_embedding(search_data.content),
         {:ok, %{status: 200}} <-
           QdrantClient.put(
             "/collections/posts/points",
             build_index_payload(activity, embedding)
           ) do
      :ok
    else
      e -> {:error, e}
    end
  end

  @impl true
  def remove_from_index(object) do
    activity = Activity.get_by_object_ap_id_with_object(object.data["id"])
    id = activity.id |> FlakeId.from_string() |> Ecto.UUID.cast!()

    with {:ok, %{status: 200}} <-
           QdrantClient.post("/collections/posts/points/delete", %{"points" => [id]}) do
      :ok
    else
      e -> {:error, e}
    end
  end

  @impl true
  def search(user, original_query, options) do
    query = "Represent this sentence for searching relevant passages: #{original_query}"

    with {:ok, embedding} <- get_embedding(query),
         {:ok, %{body: %{"result" => result}}} <-
           QdrantClient.post(
             "/collections/posts/points/search",
             build_search_payload(embedding, options)
           ) do
      ids =
        Enum.map(result, fn %{"id" => id} ->
          Ecto.UUID.dump!(id)
        end)

      from(a in Activity, where: a.id in ^ids)
      |> Activity.with_preloaded_object()
      |> Activity.restrict_deactivated_users()
      |> Ecto.Query.order_by([a], fragment("array_position(?, ?)", ^ids, a.id))
      |> Pleroma.Repo.all()
      |> maybe_fetch(user, original_query)
    else
      _ ->
        []
    end
  end

  @impl true
  def healthcheck_endpoints do
    qdrant_health =
      Config.get([Pleroma.Search.QdrantSearch, :qdrant_url])
      |> URI.parse()
      |> Map.put(:path, "/healthz")
      |> URI.to_string()

    openai_health = Config.get([Pleroma.Search.QdrantSearch, :openai_healthcheck_url])

    [qdrant_health, openai_health] |> Enum.filter(& &1)
  end
end

defmodule Pleroma.Search.QdrantSearch.OpenAIClient do
  alias Pleroma.Config.Getting, as: Config

  def post(path, body) do
    Tesla.post(client(), path, body)
  end

  defp client do
    Tesla.client(middleware(), Pleroma.HTTP.AdapterHelper.tesla_adapter())
  end

  defp middleware do
    [
      {Tesla.Middleware.BaseUrl, Config.get([Pleroma.Search.QdrantSearch, :openai_url])},
      Tesla.Middleware.JSON,
      {Tesla.Middleware.Headers,
       [
         {"Authorization", "Bearer #{Config.get([Pleroma.Search.QdrantSearch, :openai_api_key])}"}
       ]}
    ]
  end
end

defmodule Pleroma.Search.QdrantSearch.QdrantClient do
  alias Pleroma.Config.Getting, as: Config

  def delete(path) do
    Tesla.delete(client(), path)
  end

  def post(path, body) do
    Tesla.post(client(), path, body)
  end

  def put(path, body) do
    Tesla.put(client(), path, body)
  end

  defp client do
    Tesla.client(middleware(), Pleroma.HTTP.AdapterHelper.tesla_adapter())
  end

  defp middleware do
    [
      {Tesla.Middleware.BaseUrl, Config.get([Pleroma.Search.QdrantSearch, :qdrant_url])},
      Tesla.Middleware.JSON,
      {Tesla.Middleware.Headers,
       [
         {"api-key", Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_api_key])}
       ]}
    ]
  end
end