diff --git a/changelog.d/qdrant_search.add b/changelog.d/qdrant_search.add new file mode 100644 index 000000000..9801131d1 --- /dev/null +++ b/changelog.d/qdrant_search.add @@ -0,0 +1 @@ +Add Qdrant/OpenAI embedding search diff --git a/config/config.exs b/config/config.exs index 6a234db03..64c2a267f 100644 --- a/config/config.exs +++ b/config/config.exs @@ -913,9 +913,10 @@ config :pleroma, Pleroma.Search.QdrantSearch, qdrant_url: "http://127.0.0.1:6333/", - qdrant_api_key: nil, - ollama_url: "http://127.0.0.1:11434", - ollama_model: "snowflake-arctic-embed:xs", + qdrant_api_key: "", + openai_url: "http://127.0.0.1:11345", + openai_model: "snowflake/snowflake-arctic-embed-xs", + openai_api_key: "", qdrant_index_configuration: %{ vectors: %{size: 384, distance: "Cosine"} } diff --git a/docs/configuration/search.md b/docs/configuration/search.md index 682d1e52a..d34f84d4f 100644 --- a/docs/configuration/search.md +++ b/docs/configuration/search.md @@ -12,9 +12,27 @@ While it has no external dependencies, it has problems with performance and rele ## QdrantSearch -This uses the vector search engine [Qdrant](https://qdrant.tech) to search the posts in a vector space. This needs a way to generate embeddings, for now only the [Ollama](Ollama) api is supported. +This uses the vector search engine [Qdrant](https://qdrant.tech) to search the posts in a vector space. This needs a way to generate embeddings and uses the [OpenAI API](https://platform.openai.com/docs/guides/embeddings/what-are-embeddings). This is implemented by several project besides OpenAI itself, including the python-based fastembed-server found in `supplemental/search/fastembed-api`. -The default settings will support a setup where both Ollama and Qdrant run on the same system as pleroma. The embedding model used by Ollama will need to be pulled first (e.g. `ollama pull snowflake-arctic-embed:xs`) for the embedding to work. +The default settings will support a setup where both the fastembed server and Qdrant run on the same system as pleroma. To use it, set the search provider and run the fastembed server, see the README in `supplemental/search/fastembed-api`: + +> config :pleroma, Pleroma.Search, module: Pleroma.Search.QdrantSearch + +Then, start the Qdrant server, see [here](https://qdrant.tech/documentation/quick-start/) for instructions. + +You will also need to create the Qdrant index once by running `mix pleroma.search.indexer create_index`. Running `mix pleroma.search.indexer index` will retroactively index the last 100_000 activities. + +### Indexing and model options + +To see the available configuration options, check out the QdrantSearch section in `config/config.exs`. + +The default indexing option work for the default model (`snowflake-arctic-embed-xs`). To optimize for a low memory footprint, adjust the index configuration as described in the [Qdrant docs](https://qdrant.tech/documentation/guides/optimize/). See also [this blog post](https://qdrant.tech/articles/memory-consumption/) that goes into detail. + +Different embedding models will need different vector size settings. You can see a list of the models supported by the fastembed server [here](https://qdrant.github.io/fastembed/examples/Supported_Models), including their vector dimensions. These vector dimensions need to be set in the `qdrant_index_configuration`. + +E.g, If you want to use `sentence-transformers/all-MiniLM-L6-v2` as a model, you will not need to adjust things, because it and `snowflake-arctic-embed-xs` are both 384 dimensional models. If you want to use `snowflake/snowflake-arctic-embed-l`, you will need to adjust the `size` parameter in the `qdrant_index_configuration` to 1024, as it has a dimension of 1024. + +When using a different model, you will need do drop the index and recreate it (`mix pleroma.search.indexer drop_index` and `mix pleroma.search.indexer create_index`), as the different embeddings are not compatible with each other. ## Meilisearch diff --git a/lib/mix/tasks/pleroma/search/indexer.ex b/lib/mix/tasks/pleroma/search/indexer.ex index 326646b69..81a9fced6 100644 --- a/lib/mix/tasks/pleroma/search/indexer.ex +++ b/lib/mix/tasks/pleroma/search/indexer.ex @@ -9,9 +9,23 @@ defmodule Mix.Tasks.Pleroma.Search.Indexer do alias Pleroma.Workers.SearchIndexingWorker def run(["create_index"]) do - Application.ensure_all_started(:pleroma) + start_pleroma() - Pleroma.Config.get([Pleroma.Search, :module]).create_index() + with :ok <- Pleroma.Config.get([Pleroma.Search, :module]).create_index() do + IO.puts("Index created") + else + e -> IO.puts("Could not create index: #{inspect(e)}") + end + end + + def run(["drop_index"]) do + start_pleroma() + + with :ok <- Pleroma.Config.get([Pleroma.Search, :module]).drop_index() do + IO.puts("Index dropped") + else + e -> IO.puts("Could not drop index: #{inspect(e)}") + end end def run(["index" | options]) do diff --git a/lib/pleroma/search/database_search.ex b/lib/pleroma/search/database_search.ex index 31bfc7e33..24a1ff431 100644 --- a/lib/pleroma/search/database_search.ex +++ b/lib/pleroma/search/database_search.ex @@ -48,6 +48,12 @@ def add_to_index(_activity), do: :ok @impl true def remove_from_index(_object), do: :ok + @impl true + def create_index, do: :ok + + @impl true + def drop_index, do: :ok + def maybe_restrict_author(query, %User{} = author) do Activity.Queries.by_author(query, author) end diff --git a/lib/pleroma/search/meilisearch.ex b/lib/pleroma/search/meilisearch.ex index 2bff663e8..50f5984d6 100644 --- a/lib/pleroma/search/meilisearch.ex +++ b/lib/pleroma/search/meilisearch.ex @@ -10,6 +10,12 @@ defmodule Pleroma.Search.Meilisearch do @behaviour Pleroma.Search.SearchBackend + @impl true + def create_index, do: :ok + + @impl true + def drop_index, do: :ok + defp meili_headers do private_key = Config.get([Pleroma.Search.Meilisearch, :private_key]) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 315262cb3..283c43075 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -1,28 +1,40 @@ defmodule Pleroma.Search.QdrantSearch do @behaviour Pleroma.Search.SearchBackend import Ecto.Query - alias Pleroma.Activity + alias Pleroma.Activity + alias Pleroma.Config.Getting, as: Config + + alias __MODULE__.OpenAIClient alias __MODULE__.QdrantClient - alias __MODULE__.OllamaClient import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] @impl true - def create_index() do - payload = Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration]) - QdrantClient.put("/collections/posts", payload) + 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 - def drop_index() do - QdrantClient.delete("/collections/posts") + @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: %{"embedding" => embedding}}} <- - OllamaClient.post("/api/embeddings", %{ - prompt: text, - model: Pleroma.Config.get([Pleroma.Search.QdrantSearch, :ollama_model]) + with {:ok, %{body: %{"data" => [%{"embedding" => embedding}]}}} <- + OpenAIClient.post("/v1/embeddings", %{ + input: text, + model: Config.get([Pleroma.Search.QdrantSearch, :openai_model]) }) do {:ok, embedding} else @@ -42,10 +54,11 @@ defp build_index_payload(activity, embedding) do } end - defp build_search_payload(embedding) do + defp build_search_payload(embedding, options) do %{ vector: embedding, - limit: 20 + limit: options[:limit] || 20, + offset: options[:offset] || 0 } end @@ -71,12 +84,28 @@ def add_to_index(activity) do end @impl true - def search(_user, query, _options) do + 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, query, options) do query = "Represent this sentence for searching relevant passages: #{query}" with {:ok, embedding} <- get_embedding(query), {:ok, %{body: %{"result" => result}}} <- - QdrantClient.post("/collections/posts/points/search", build_search_payload(embedding)) do + QdrantClient.post( + "/collections/posts/points/search", + build_search_payload(embedding, options) + ) do ids = Enum.map(result, fn %{"id" => id} -> Ecto.UUID.dump!(id) @@ -92,24 +121,26 @@ def search(_user, query, _options) do [] end end - - @impl true - def remove_from_index(_object) do - :ok - end end -defmodule Pleroma.Search.QdrantSearch.OllamaClient do +defmodule Pleroma.Search.QdrantSearch.OpenAIClient do use Tesla + alias Pleroma.Config.Getting, as: Config - plug(Tesla.Middleware.BaseUrl, Pleroma.Config.get([Pleroma.Search.QdrantSearch, :ollama_url])) + plug(Tesla.Middleware.BaseUrl, Config.get([Pleroma.Search.QdrantSearch, :openai_url])) plug(Tesla.Middleware.JSON) + + plug(Tesla.Middleware.Headers, [ + {"Authorization", + "Bearer #{Pleroma.Config.get([Pleroma.Search.QdrantSearch, :openai_api_key])}"} + ]) end defmodule Pleroma.Search.QdrantSearch.QdrantClient do use Tesla + alias Pleroma.Config.Getting, as: Config - plug(Tesla.Middleware.BaseUrl, Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_url])) + plug(Tesla.Middleware.BaseUrl, Config.get([Pleroma.Search.QdrantSearch, :qdrant_url])) plug(Tesla.Middleware.JSON) plug(Tesla.Middleware.Headers, [ diff --git a/lib/pleroma/search/search_backend.ex b/lib/pleroma/search/search_backend.ex index 5be0169d0..9735ab3f4 100644 --- a/lib/pleroma/search/search_backend.ex +++ b/lib/pleroma/search/search_backend.ex @@ -26,4 +26,9 @@ defmodule Pleroma.Search.SearchBackend do Create the index """ @callback create_index() :: :ok | {:error, any()} + + @doc """ + Drop the index + """ + @callback drop_index() :: :ok | {:error, any()} end diff --git a/supplemental/search/fastembed-api/Dockerfile b/supplemental/search/fastembed-api/Dockerfile new file mode 100644 index 000000000..c1e0ef51f --- /dev/null +++ b/supplemental/search/fastembed-api/Dockerfile @@ -0,0 +1,9 @@ +FROM python:3.9 + +WORKDIR /code +COPY fastembed-server.py /workdir/fastembed-server.py +COPY requirements.txt /workdir/requirements.txt + +RUN pip install -r /workdir/requirements.txt + +CMD ["python", "/workdir/fastembed-server.py"] diff --git a/supplemental/search/fastembed-api/README.md b/supplemental/search/fastembed-api/README.md new file mode 100644 index 000000000..63a037207 --- /dev/null +++ b/supplemental/search/fastembed-api/README.md @@ -0,0 +1,6 @@ +# About +This is a minimal implementation of the [OpenAI Embeddings API](https://platform.openai.com/docs/guides/embeddings/what-are-embeddings) meant to be used with the QdrantSearch backend. + +# Usage + +The easiest way to run it is to just use docker compose with `docker compose up`. This starts the server on the default configured port. Different models can be used, for a full list of supported models, check the [fastembed documentation](https://qdrant.github.io/fastembed/examples/Supported_Models/). The first time a model is requested it will be downloaded, which can take a few seconds. diff --git a/supplemental/search/fastembed-api/compose.yml b/supplemental/search/fastembed-api/compose.yml new file mode 100644 index 000000000..d4cb31722 --- /dev/null +++ b/supplemental/search/fastembed-api/compose.yml @@ -0,0 +1,5 @@ +services: + web: + build: . + ports: + - "11345:11345" diff --git a/supplemental/search/fastembed-api/fastembed-server.py b/supplemental/search/fastembed-api/fastembed-server.py new file mode 100644 index 000000000..dd4a7a9c8 --- /dev/null +++ b/supplemental/search/fastembed-api/fastembed-server.py @@ -0,0 +1,23 @@ +from fastembed import TextEmbedding +from fastapi import FastAPI +from pydantic import BaseModel + +models = {} + +app = FastAPI() + +class EmbeddingRequest(BaseModel): + model: str + input: str + +@app.post("/v1/embeddings") +def embeddings(request: EmbeddingRequest): + model = models.get(request.model) or TextEmbedding(request.model) + models[request.model] = model + embeddings = next(model.embed(request.input)).tolist() + return {"data": [{"embedding": embeddings}]} + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=11345) diff --git a/supplemental/search/fastembed-api/requirements.txt b/supplemental/search/fastembed-api/requirements.txt new file mode 100644 index 000000000..db67a8402 --- /dev/null +++ b/supplemental/search/fastembed-api/requirements.txt @@ -0,0 +1,4 @@ +fastapi==0.111.0 +fastembed==0.2.7 +pydantic==1.10.15 +uvicorn==0.29.0 diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs new file mode 100644 index 000000000..a2f9cc7ec --- /dev/null +++ b/test/pleroma/search/qdrant_search_test.exs @@ -0,0 +1,80 @@ +# Pleroma: A lightweight social networking server +# Copyright © 2017-2021 Pleroma Authors +# SPDX-License-Identifier: AGPL-3.0-only + +defmodule Pleroma.Search.QdrantSearchTest do + use Pleroma.DataCase, async: true + use Oban.Testing, repo: Pleroma.Repo + + import Pleroma.Factory + import Mox + + alias Pleroma.Search.QdrantSearch + alias Pleroma.UnstubbedConfigMock, as: Config + alias Pleroma.Web.CommonAPI + alias Pleroma.Workers.SearchIndexingWorker + + describe "Qdrant search" do + test "indexes a public post on creation, deletes from the index on deletion" do + user = insert(:user) + + Tesla.Mock.mock(fn + %{method: :post, url: "https://openai.url/v1/embeddings"} -> + send(self(), "posted_to_openai") + + Tesla.Mock.json(%{ + data: [%{embedding: [1, 2, 3]}] + }) + + %{method: :put, url: "https://qdrant.url/collections/posts/points", body: body} -> + send(self(), "posted_to_qdrant") + + assert match?(%{"points" => [%{"vector" => [1, 2, 3]}]}, Jason.decode!(body)) + + Tesla.Mock.json("ok") + + %{method: :post, url: "https://qdrant.url/collections/posts/points/delete"} -> + send(self(), "deleted_from_qdrant") + Tesla.Mock.json("ok") + end) + + Config + |> expect(:get, 6, fn + [Pleroma.Search, :module], nil -> + QdrantSearch + + [Pleroma.Search.QdrantSearch, key], nil -> + %{ + openai_model: "a_model", + openai_url: "https://openai.url", + qdrant_url: "https://qdrant.url" + }[key] + end) + + {:ok, activity} = + CommonAPI.post(user, %{ + status: "guys i just don't wanna leave the swamp", + visibility: "public" + }) + + args = %{"op" => "add_to_index", "activity" => activity.id} + + assert_enqueued( + worker: SearchIndexingWorker, + args: args + ) + + assert :ok = perform_job(SearchIndexingWorker, args) + assert_received("posted_to_openai") + assert_received("posted_to_qdrant") + + {:ok, _} = CommonAPI.delete(activity.id, user) + + delete_args = %{"op" => "remove_from_index", "object" => activity.object.id} + assert_enqueued(worker: SearchIndexingWorker, args: delete_args) + assert :ok = perform_job(SearchIndexingWorker, delete_args) + + assert_received("deleted_from_qdrant") + end + end +end