From cd7e2138d11901fc7a0c8c2f22b7a5d57383a555 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Tue, 14 May 2024 14:13:37 +0400 Subject: [PATCH 01/30] Search: Basic Qdrant/Ollama search --- config/config.exs | 9 ++ lib/mix/tasks/pleroma/search/indexer.ex | 60 ++++++++++++ lib/pleroma/search/qdrant_search.ex | 117 ++++++++++++++++++++++++ 3 files changed, 186 insertions(+) create mode 100644 lib/mix/tasks/pleroma/search/indexer.ex create mode 100644 lib/pleroma/search/qdrant_search.ex diff --git a/config/config.exs b/config/config.exs index b69044a2b..f74eda6b2 100644 --- a/config/config.exs +++ b/config/config.exs @@ -915,6 +915,15 @@ config :pleroma, Pleroma.Uploaders.Uploader, timeout: 30_000 +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_index_configuration: %{ + vectors: %{size: 384, distance: "Cosine"} + } + # Import environment specific config. This must remain at the bottom # of this file so it overrides the configuration defined above. import_config "#{Mix.env()}.exs" diff --git a/lib/mix/tasks/pleroma/search/indexer.ex b/lib/mix/tasks/pleroma/search/indexer.ex new file mode 100644 index 000000000..ffa2f3c94 --- /dev/null +++ b/lib/mix/tasks/pleroma/search/indexer.ex @@ -0,0 +1,60 @@ +# Pleroma: A lightweight social networking server +# Copyright © 2017-2021 Pleroma Authors +# SPDX-License-Identifier: AGPL-3.0-only + +defmodule Mix.Tasks.Pleroma.Search.Indexer do + import Mix.Pleroma + import Ecto.Query + + alias Pleroma.Workers.SearchIndexingWorker + + def run(["index" | options]) do + {options, [], []} = + OptionParser.parse( + options, + strict: [ + limit: :integer + ] + ) + + start_pleroma() + + limit = Keyword.get(options, :limit, 100_000) + + per_step = 1000 + chunks = max(div(limit, per_step), 1) + + 1..chunks + |> Enum.each(fn step -> + q = + from(a in Pleroma.Activity, + limit: ^per_step, + offset: ^per_step * (^step - 1), + select: [:id], + order_by: [desc: :id] + ) + + {:ok, ids} = + Pleroma.Repo.transaction(fn -> + Pleroma.Repo.stream(q, timeout: :infinity) + |> Enum.map(fn a -> + a.id + end) + end) + + IO.puts("Got #{length(ids)} activities, adding to indexer") + + ids + |> Enum.chunk_every(100) + |> Enum.each(fn chunk -> + IO.puts("Adding #{length(chunk)} activities to indexing queue") + + chunk + |> Enum.map(fn id -> + SearchIndexingWorker.new(%{"op" => "add_to_index", "activity" => id}) + end) + |> Oban.insert_all() + end) + end) + end +end diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex new file mode 100644 index 000000000..fcaa9e686 --- /dev/null +++ b/lib/pleroma/search/qdrant_search.ex @@ -0,0 +1,117 @@ +defmodule Pleroma.Search.QdrantSearch do + @behaviour Pleroma.Search.SearchBackend + import Ecto.Query + alias Pleroma.Activity + + alias __MODULE__.QdrantClient + alias __MODULE__.OllamaClient + + import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] + + def initialize_index() do + payload = Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration]) + QdrantClient.put("/collections/posts", payload) + end + + def drop_index() do + QdrantClient.delete("/collections/posts") + 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]) + }) + |> IO.inspect() do + {:ok, embedding} + else + _ -> + {:error, "Failed to get embedding"} + end + end + + defp build_index_payload(activity, embedding) do + %{ + points: [ + %{ + id: activity.id |> FlakeId.from_string() |> Ecto.UUID.cast!(), + vector: embedding + } + ] + } + end + + defp build_search_payload(embedding) do + %{ + vector: embedding, + limit: 20 + } + end + + @impl true + def add_to_index(activity) do + # This will only index public or unlisted notes + maybe_search_data = object_to_search_data(activity.object) + IO.puts("TRYING TO INDEX\n\n") + + if activity.data["type"] == "Create" and maybe_search_data do + with {:ok, embedding} <- get_embedding(maybe_search_data.content), + {:ok, %{status: 200}} <- + QdrantClient.put( + "/collections/posts/points", + build_index_payload(activity, embedding) + ) do + :ok + else + e -> {:error, e} + end + else + :ok + end + end + + @impl true + def search(_user, query, _options) do + with {:ok, embedding} <- get_embedding(query), + {:ok, %{body: %{"result" => result}}} <- + QdrantClient.post("/collections/posts/points/search", build_search_payload(embedding)) 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() + else + _ -> + [] + end + end + + @impl true + def remove_from_index(_object) do + :ok + end +end + +defmodule Pleroma.Search.QdrantSearch.OllamaClient do + use Tesla + + plug(Tesla.Middleware.BaseUrl, Pleroma.Config.get([Pleroma.Search.QdrantSearch, :ollama_url])) + plug(Tesla.Middleware.JSON) +end + +defmodule Pleroma.Search.QdrantSearch.QdrantClient do + use Tesla + + plug(Tesla.Middleware.BaseUrl, Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_url])) + plug(Tesla.Middleware.JSON) + + plug(Tesla.Middleware.Headers, [ + {"api-key", Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_api_key])} + ]) +end From bb08a766f4c1bd84c98e245c1871c46fcc7c7a8d Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Tue, 14 May 2024 14:26:41 +0400 Subject: [PATCH 02/30] QdrantSearch: Remove debugging stuff --- lib/pleroma/search/qdrant_search.ex | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index fcaa9e686..726a30b3b 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -22,8 +22,7 @@ def get_embedding(text) do OllamaClient.post("/api/embeddings", %{ prompt: text, model: Pleroma.Config.get([Pleroma.Search.QdrantSearch, :ollama_model]) - }) - |> IO.inspect() do + }) do {:ok, embedding} else _ -> @@ -53,7 +52,6 @@ defp build_search_payload(embedding) do def add_to_index(activity) do # This will only index public or unlisted notes maybe_search_data = object_to_search_data(activity.object) - IO.puts("TRYING TO INDEX\n\n") if activity.data["type"] == "Create" and maybe_search_data do with {:ok, embedding} <- get_embedding(maybe_search_data.content), From 1490ff30af7001adc386b4fec54c62e1a524d7d6 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Tue, 14 May 2024 15:09:38 +0400 Subject: [PATCH 03/30] QdrantSearch: Add query prefix. --- lib/pleroma/search/qdrant_search.ex | 2 ++ 1 file changed, 2 insertions(+) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 726a30b3b..31e7754ae 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -71,6 +71,8 @@ def add_to_index(activity) do @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 From c50f0f31f418037063bd97efcdc0f60b89594212 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Tue, 14 May 2024 16:56:58 +0400 Subject: [PATCH 04/30] Docs/Search: Add basic documentation of the qdrant search --- docs/configuration/search.md | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/docs/configuration/search.md b/docs/configuration/search.md index 0316c9bf4..682d1e52a 100644 --- a/docs/configuration/search.md +++ b/docs/configuration/search.md @@ -10,6 +10,12 @@ To use built-in search that has no external dependencies, set the search module While it has no external dependencies, it has problems with performance and relevancy. +## 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. + +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. + ## Meilisearch Note that it's quite a bit more memory hungry than PostgreSQL (around 4-5G for ~1.2 million From 1261c43a7af48ed6e6753461944659391c4c58cc Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Tue, 14 May 2024 17:19:36 +0400 Subject: [PATCH 05/30] SearchBackend: Add create_index --- lib/mix/tasks/pleroma/search/indexer.ex | 6 ++++++ lib/pleroma/search/qdrant_search.ex | 3 ++- lib/pleroma/search/search_backend.ex | 5 +++++ 3 files changed, 13 insertions(+), 1 deletion(-) diff --git a/lib/mix/tasks/pleroma/search/indexer.ex b/lib/mix/tasks/pleroma/search/indexer.ex index ffa2f3c94..326646b69 100644 --- a/lib/mix/tasks/pleroma/search/indexer.ex +++ b/lib/mix/tasks/pleroma/search/indexer.ex @@ -8,6 +8,12 @@ defmodule Mix.Tasks.Pleroma.Search.Indexer do alias Pleroma.Workers.SearchIndexingWorker + def run(["create_index"]) do + Application.ensure_all_started(:pleroma) + + Pleroma.Config.get([Pleroma.Search, :module]).create_index() + end + def run(["index" | options]) do {options, [], []} = OptionParser.parse( diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 31e7754ae..315262cb3 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -8,7 +8,8 @@ defmodule Pleroma.Search.QdrantSearch do import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] - def initialize_index() do + @impl true + def create_index() do payload = Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration]) QdrantClient.put("/collections/posts", payload) end diff --git a/lib/pleroma/search/search_backend.ex b/lib/pleroma/search/search_backend.ex index 68bc48cec..5be0169d0 100644 --- a/lib/pleroma/search/search_backend.ex +++ b/lib/pleroma/search/search_backend.ex @@ -21,4 +21,9 @@ defmodule Pleroma.Search.SearchBackend do from index. """ @callback remove_from_index(object :: Pleroma.Object.t()) :: :ok | {:error, any()} + + @doc """ + Create the index + """ + @callback create_index() :: :ok | {:error, any()} end From a9be4907c0d7b34e5564584d2d040632c32f2aa3 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Thu, 16 May 2024 10:47:24 +0400 Subject: [PATCH 06/30] SearchBackend: Add drop_index --- lib/mix/tasks/pleroma/search/indexer.ex | 18 ++++++++++++++++-- lib/pleroma/search/database_search.ex | 6 ++++++ lib/pleroma/search/meilisearch.ex | 6 ++++++ lib/pleroma/search/qdrant_search.ex | 14 ++++++++++++-- lib/pleroma/search/search_backend.ex | 5 +++++ 5 files changed, 45 insertions(+), 4 deletions(-) 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..4bd35c17c 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -11,11 +11,21 @@ defmodule Pleroma.Search.QdrantSearch do @impl true def create_index() do payload = Pleroma.Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration]) - QdrantClient.put("/collections/posts", payload) + + with {:ok, %{status: 200}} <- QdrantClient.put("/collections/posts", payload) do + :ok + else + e -> {:error, e} + end end + @impl true def drop_index() do - QdrantClient.delete("/collections/posts") + with {:ok, %{status: 200}} <- QdrantClient.delete("/collections/posts") do + :ok + else + e -> {:error, e} + end end def get_embedding(text) do 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 From 069ce4448c556af90293cde9b9872c3d53eb894b Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 11:55:17 +0400 Subject: [PATCH 07/30] Add basic fastembed server --- python/fastembed-server.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) create mode 100644 python/fastembed-server.py diff --git a/python/fastembed-server.py b/python/fastembed-server.py new file mode 100644 index 000000000..fa3f7c82b --- /dev/null +++ b/python/fastembed-server.py @@ -0,0 +1,21 @@ +from fastembed import TextEmbedding +from fastapi import FastAPI +from pydantic import BaseModel + +model = TextEmbedding("snowflake/snowflake-arctic-embed-xs") + +app = FastAPI() + +class EmbeddingRequest(BaseModel): + model: str + prompt: str + +@app.post("/api/embeddings") +def embeddings(request: EmbeddingRequest): + embeddings = next(model.embed(request.prompt)).tolist() + return {"embedding": embeddings} + +if __name__ == "__main__": + import uvicorn + + uvicorn.run(app, host="0.0.0.0", port=11345) From 769773a500d4c6ec021776b493f7d98c9f87e81e Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 12:08:42 +0400 Subject: [PATCH 08/30] Add dockerfile --- python/Dockerfile | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 python/Dockerfile diff --git a/python/Dockerfile b/python/Dockerfile new file mode 100644 index 000000000..f83c1c1b3 --- /dev/null +++ b/python/Dockerfile @@ -0,0 +1,8 @@ +FROM python:3.9 + +WORKDIR /code +COPY fastembed-server.py /workdir/fastembed-server.py + +RUN pip install --no-cache-dir --upgrade fastembed fastapi uvicorn + +CMD ["python", "/workdir/fastembed-server.py"] From 61e9027131843858b017d3b7c18c3a396d5656a9 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 12:19:42 +0400 Subject: [PATCH 09/30] Add docker compose file for fastembed server --- python/compose.yml | 5 +++++ 1 file changed, 5 insertions(+) create mode 100644 python/compose.yml diff --git a/python/compose.yml b/python/compose.yml new file mode 100644 index 000000000..d4cb31722 --- /dev/null +++ b/python/compose.yml @@ -0,0 +1,5 @@ +services: + web: + build: . + ports: + - "11345:11345" From 933117785fb1b5b671c61d09671cf6418b105187 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 13:43:47 +0400 Subject: [PATCH 10/30] QdrantSearch: Add basic test --- lib/pleroma/search/qdrant_search.ex | 11 ++-- test/pleroma/search/qdrant_search_test.exs | 65 ++++++++++++++++++++++ 2 files changed, 72 insertions(+), 4 deletions(-) create mode 100644 test/pleroma/search/qdrant_search_test.exs diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 4bd35c17c..2d9315a2f 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -5,12 +5,13 @@ defmodule Pleroma.Search.QdrantSearch do alias __MODULE__.QdrantClient alias __MODULE__.OllamaClient + alias Pleroma.Config.Getting, as: Config 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]) + payload = Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration]) with {:ok, %{status: 200}} <- QdrantClient.put("/collections/posts", payload) do :ok @@ -32,7 +33,7 @@ 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]) + model: Config.get([Pleroma.Search.QdrantSearch, :ollama_model]) }) do {:ok, embedding} else @@ -111,15 +112,17 @@ def remove_from_index(_object) do defmodule Pleroma.Search.QdrantSearch.OllamaClient 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, :ollama_url])) plug(Tesla.Middleware.JSON) 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/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs new file mode 100644 index 000000000..9be246a9a --- /dev/null +++ b/test/pleroma/search/qdrant_search_test.exs @@ -0,0 +1,65 @@ +# 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.Web.CommonAPI + alias Pleroma.UnstubbedConfigMock, as: Config + alias Pleroma.Search.QdrantSearch + alias Pleroma.Workers.SearchIndexingWorker + + describe "Qdrant search" do + test "indexes a public post on creation" do + user = insert(:user) + + Tesla.Mock.mock(fn + %{method: :post, url: "https://ollama.url/api/embeddings"} -> + send(self(), "posted_to_ollama") + Tesla.Mock.json(%{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") + end) + + Config + |> expect(:get, 4, fn + [Pleroma.Search, :module], nil -> + QdrantSearch + + [Pleroma.Search.QdrantSearch, key], nil -> + %{ + ollama_model: "a_model", + ollama_url: "https://ollama.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_ollama") + assert_received("posted_to_qdrant") + end + end +end From e3933a067feae1f087616f675657d6ff99b2782b Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 14:04:32 +0400 Subject: [PATCH 11/30] QdrantSearch: Implement post deletion --- lib/pleroma/search/qdrant_search.ex | 18 +++++++++++++----- test/pleroma/search/qdrant_search_test.exs | 16 ++++++++++++++-- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 2d9315a2f..acfaaff52 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -81,6 +81,19 @@ def add_to_index(activity) do 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, query, _options) do query = "Represent this sentence for searching relevant passages: #{query}" @@ -103,11 +116,6 @@ 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 diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index 9be246a9a..e816311aa 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -15,7 +15,7 @@ defmodule Pleroma.Search.QdrantSearchTest do alias Pleroma.Workers.SearchIndexingWorker describe "Qdrant search" do - test "indexes a public post on creation" do + test "indexes a public post on creation, deletes from the index on deletion" do user = insert(:user) Tesla.Mock.mock(fn @@ -29,10 +29,14 @@ test "indexes a public post on creation" do 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, 4, fn + |> expect(:get, 6, fn [Pleroma.Search, :module], nil -> QdrantSearch @@ -60,6 +64,14 @@ test "indexes a public post on creation" do assert :ok = perform_job(SearchIndexingWorker, args) assert_received("posted_to_ollama") 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 From 39525bcec7c685cb28ca4702b6e145a78e733fee Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 14:07:47 +0400 Subject: [PATCH 12/30] Add qdrant changelog --- changelog.d/qdrant_search.add | 1 + 1 file changed, 1 insertion(+) create mode 100644 changelog.d/qdrant_search.add diff --git a/changelog.d/qdrant_search.add b/changelog.d/qdrant_search.add new file mode 100644 index 000000000..6f9e39e23 --- /dev/null +++ b/changelog.d/qdrant_search.add @@ -0,0 +1 @@ +Add Qdrant/Ollama search From 3345ddd2d4ef380929cc231118a5fb6486c0bd5c Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sat, 18 May 2024 15:02:22 +0400 Subject: [PATCH 13/30] Linting --- lib/pleroma/search/qdrant_search.ex | 11 ++++++----- test/pleroma/search/qdrant_search_test.exs | 4 ++-- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index acfaaff52..a6c6c6a0d 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -1,16 +1,17 @@ defmodule Pleroma.Search.QdrantSearch do @behaviour Pleroma.Search.SearchBackend import Ecto.Query - alias Pleroma.Activity - alias __MODULE__.QdrantClient - alias __MODULE__.OllamaClient + alias Pleroma.Activity alias Pleroma.Config.Getting, as: Config + alias __MODULE__.OllamaClient + alias __MODULE__.QdrantClient + import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] @impl true - def create_index() do + def create_index do payload = Config.get([Pleroma.Search.QdrantSearch, :qdrant_index_configuration]) with {:ok, %{status: 200}} <- QdrantClient.put("/collections/posts", payload) do @@ -21,7 +22,7 @@ def create_index() do end @impl true - def drop_index() do + def drop_index do with {:ok, %{status: 200}} <- QdrantClient.delete("/collections/posts") do :ok else diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index e816311aa..698894cdb 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -9,9 +9,9 @@ defmodule Pleroma.Search.QdrantSearchTest do import Pleroma.Factory import Mox - alias Pleroma.Web.CommonAPI - alias Pleroma.UnstubbedConfigMock, as: Config alias Pleroma.Search.QdrantSearch + alias Pleroma.UnstubbedConfigMock, as: Config + alias Pleroma.Web.CommonAPI alias Pleroma.Workers.SearchIndexingWorker describe "Qdrant search" do From 72ec261a69a7dda7ab95667e425824ab7758b636 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 12:17:46 +0400 Subject: [PATCH 14/30] B QdrantSearch: Switch to OpenAI api --- changelog.d/qdrant_search.add | 2 +- config/config.exs | 7 ++++--- lib/pleroma/search/qdrant_search.ex | 19 ++++++++++++------- test/pleroma/search/qdrant_search_test.exs | 15 +++++++++------ 4 files changed, 26 insertions(+), 17 deletions(-) diff --git a/changelog.d/qdrant_search.add b/changelog.d/qdrant_search.add index 6f9e39e23..9801131d1 100644 --- a/changelog.d/qdrant_search.add +++ b/changelog.d/qdrant_search.add @@ -1 +1 @@ -Add Qdrant/Ollama search +Add Qdrant/OpenAI embedding search diff --git a/config/config.exs b/config/config.exs index f74eda6b2..dd0150c66 100644 --- a/config/config.exs +++ b/config/config.exs @@ -917,9 +917,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", + openai_api_key: "", qdrant_index_configuration: %{ vectors: %{size: 384, distance: "Cosine"} } diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index a6c6c6a0d..5ae04be78 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -5,7 +5,7 @@ defmodule Pleroma.Search.QdrantSearch do alias Pleroma.Activity alias Pleroma.Config.Getting, as: Config - alias __MODULE__.OllamaClient + alias __MODULE__.OpenAIClient alias __MODULE__.QdrantClient import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] @@ -31,10 +31,10 @@ def drop_index do end def get_embedding(text) do - with {:ok, %{body: %{"embedding" => embedding}}} <- - OllamaClient.post("/api/embeddings", %{ - prompt: text, - model: 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 @@ -119,12 +119,17 @@ def search(_user, query, _options) do 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, 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 diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index 698894cdb..a2f9cc7ec 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -19,9 +19,12 @@ 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://ollama.url/api/embeddings"} -> - send(self(), "posted_to_ollama") - Tesla.Mock.json(%{embedding: [1, 2, 3]}) + %{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") @@ -42,8 +45,8 @@ test "indexes a public post on creation, deletes from the index on deletion" do [Pleroma.Search.QdrantSearch, key], nil -> %{ - ollama_model: "a_model", - ollama_url: "https://ollama.url", + openai_model: "a_model", + openai_url: "https://openai.url", qdrant_url: "https://qdrant.url" }[key] end) @@ -62,7 +65,7 @@ test "indexes a public post on creation, deletes from the index on deletion" do ) assert :ok = perform_job(SearchIndexingWorker, args) - assert_received("posted_to_ollama") + assert_received("posted_to_openai") assert_received("posted_to_qdrant") {:ok, _} = CommonAPI.delete(activity.id, user) From b9af017a4cf1025c7d8245fa4f1dbcb678ddd4b9 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 12:33:49 +0400 Subject: [PATCH 15/30] B FastembedServer: Switch to OpenAI api, support changing models --- python/fastembed-server.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/python/fastembed-server.py b/python/fastembed-server.py index fa3f7c82b..dd4a7a9c8 100644 --- a/python/fastembed-server.py +++ b/python/fastembed-server.py @@ -2,18 +2,20 @@ from fastembed import TextEmbedding from fastapi import FastAPI from pydantic import BaseModel -model = TextEmbedding("snowflake/snowflake-arctic-embed-xs") +models = {} app = FastAPI() class EmbeddingRequest(BaseModel): model: str - prompt: str + input: str -@app.post("/api/embeddings") +@app.post("/v1/embeddings") def embeddings(request: EmbeddingRequest): - embeddings = next(model.embed(request.prompt)).tolist() - return {"embedding": embeddings} + 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 From c139a9f38c06ab4485b98b56b9ad4cce4d57be12 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 12:39:54 +0400 Subject: [PATCH 16/30] B Config: Set default Qdrant embedder to our fastembed-api server --- config/config.exs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/config/config.exs b/config/config.exs index dd0150c66..c3b20947d 100644 --- a/config/config.exs +++ b/config/config.exs @@ -919,7 +919,7 @@ qdrant_url: "http://127.0.0.1:6333/", qdrant_api_key: "", openai_url: "http://127.0.0.1:11345", - openai_model: "snowflake", + openai_model: "snowflake/snowflake-arctic-embed-xs", openai_api_key: "", qdrant_index_configuration: %{ vectors: %{size: 384, distance: "Cosine"} From e142ea400a9ed3595f8d432edd90ea26fc7d2eb5 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 12:42:08 +0400 Subject: [PATCH 17/30] Docs: Switch docs from Ollama to OpenAI. --- docs/configuration/search.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/configuration/search.md b/docs/configuration/search.md index 682d1e52a..388f5acd1 100644 --- a/docs/configuration/search.md +++ b/docs/configuration/search.md @@ -12,9 +12,9 @@ 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. ## Meilisearch From dd48810186e3b4ee14e1d3727f37bd470d0711a4 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 12:47:08 +0400 Subject: [PATCH 18/30] B FastembedAPI: Move to more appropriate folder --- {python => supplemental/search/fastembed-api}/Dockerfile | 0 {python => supplemental/search/fastembed-api}/compose.yml | 0 {python => supplemental/search/fastembed-api}/fastembed-server.py | 0 3 files changed, 0 insertions(+), 0 deletions(-) rename {python => supplemental/search/fastembed-api}/Dockerfile (100%) rename {python => supplemental/search/fastembed-api}/compose.yml (100%) rename {python => supplemental/search/fastembed-api}/fastembed-server.py (100%) diff --git a/python/Dockerfile b/supplemental/search/fastembed-api/Dockerfile similarity index 100% rename from python/Dockerfile rename to supplemental/search/fastembed-api/Dockerfile diff --git a/python/compose.yml b/supplemental/search/fastembed-api/compose.yml similarity index 100% rename from python/compose.yml rename to supplemental/search/fastembed-api/compose.yml diff --git a/python/fastembed-server.py b/supplemental/search/fastembed-api/fastembed-server.py similarity index 100% rename from python/fastembed-server.py rename to supplemental/search/fastembed-api/fastembed-server.py From 8329ad521419119f89e3e2577269475190cfe921 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 12:59:03 +0400 Subject: [PATCH 19/30] B FastembedAPI: Add requirements.txt --- supplemental/search/fastembed-api/Dockerfile | 3 ++- supplemental/search/fastembed-api/requirements.txt | 4 ++++ 2 files changed, 6 insertions(+), 1 deletion(-) create mode 100644 supplemental/search/fastembed-api/requirements.txt diff --git a/supplemental/search/fastembed-api/Dockerfile b/supplemental/search/fastembed-api/Dockerfile index f83c1c1b3..c1e0ef51f 100644 --- a/supplemental/search/fastembed-api/Dockerfile +++ b/supplemental/search/fastembed-api/Dockerfile @@ -2,7 +2,8 @@ FROM python:3.9 WORKDIR /code COPY fastembed-server.py /workdir/fastembed-server.py +COPY requirements.txt /workdir/requirements.txt -RUN pip install --no-cache-dir --upgrade fastembed fastapi uvicorn +RUN pip install -r /workdir/requirements.txt CMD ["python", "/workdir/fastembed-server.py"] 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 From 23881842ae33a294e344cef0cc2f1385ea6819f9 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 13:04:27 +0400 Subject: [PATCH 20/30] B FastembedAPI: Add readme --- supplemental/search/fastembed-api/README.md | 6 ++++++ 1 file changed, 6 insertions(+) create mode 100644 supplemental/search/fastembed-api/README.md 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. From 6a3a0cc0f5995185428c92f3c53e9c8524ea6856 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 13:20:37 +0400 Subject: [PATCH 21/30] Docs: Write docs for the QdrantSearch --- docs/configuration/search.md | 20 +++++++++++++++++++- 1 file changed, 19 insertions(+), 1 deletion(-) diff --git a/docs/configuration/search.md b/docs/configuration/search.md index 388f5acd1..6598e533f 100644 --- a/docs/configuration/search.md +++ b/docs/configuration/search.md @@ -14,7 +14,25 @@ While it has no external dependencies, it has problems with performance and rele 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 the fastembed server and Qdrant run on the same system as pleroma. +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`: + +https://qdrant.github.io/fastembed/examples/Supported_Models/ + +> config :pleroma, Pleroma.Search, module: Pleroma.Search.QdrantSearch + +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/). + +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 From 6ec306d0684f3c5c05d768a3c431008925f21f15 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 13:24:24 +0400 Subject: [PATCH 22/30] Docs: Add more information about index memory consumption. --- docs/configuration/search.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/configuration/search.md b/docs/configuration/search.md index 6598e533f..ed85acd2a 100644 --- a/docs/configuration/search.md +++ b/docs/configuration/search.md @@ -26,7 +26,7 @@ You will also need to create the Qdrant index once by running `mix pleroma.searc 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/). +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`. From dbaab6f54e306e5fb930ce1ed0699631c8aeaae1 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 13:38:31 +0400 Subject: [PATCH 23/30] Docs: Mention running the Qdrant server --- docs/configuration/search.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/configuration/search.md b/docs/configuration/search.md index ed85acd2a..d34f84d4f 100644 --- a/docs/configuration/search.md +++ b/docs/configuration/search.md @@ -16,10 +16,10 @@ This uses the vector search engine [Qdrant](https://qdrant.tech) to search the p 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`: -https://qdrant.github.io/fastembed/examples/Supported_Models/ - > 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 From 1b4f1db9b2990f725a06f0dff41980c51853c5e9 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Sun, 19 May 2024 14:41:05 +0400 Subject: [PATCH 24/30] QdrantSearch: Support pagination. --- lib/pleroma/search/qdrant_search.ex | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 5ae04be78..283c43075 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -54,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 @@ -96,12 +97,15 @@ def remove_from_index(object) do end @impl true - def search(_user, query, _options) do + 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) From 94e4f215896dc7976a54fd146daf3e286602925a Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Thu, 23 May 2024 14:38:30 +0400 Subject: [PATCH 25/30] QdrantSearch: Deal with actor restrictions --- lib/pleroma/search/qdrant_search.ex | 22 ++++- test/pleroma/search/qdrant_search_test.exs | 95 +++++++++++++++++++++- 2 files changed, 114 insertions(+), 3 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 283c43075..9cb34ef71 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -43,23 +43,41 @@ def get_embedding(text) do 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 + 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 options[:actor] do + Map.put(base, :filter, %{ + must: [%{key: "actor", match: %{value: options[:actor].ap_id}}] + }) + else + base + end end @impl true diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index a2f9cc7ec..371074dcf 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -15,6 +15,94 @@ defmodule Pleroma.Search.QdrantSearchTest do alias Pleroma.Workers.SearchIndexingWorker describe "Qdrant search" do + test "searches for a term by encoding it and sending it to qdrant" do + user = insert(:user) + + {:ok, activity} = + CommonAPI.post(user, %{ + status: "guys i just don't wanna leave the swamp", + visibility: "public" + }) + + Config + |> expect(:get, 3, 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) + + Tesla.Mock.mock(fn + %{url: "https://openai.url/v1/embeddings", method: :post} -> + Tesla.Mock.json(%{ + data: [%{embedding: [1, 2, 3]}] + }) + + %{url: "https://qdrant.url/collections/posts/points/search", method: :post, body: body} -> + data = Jason.decode!(body) + refute data["filter"] + + Tesla.Mock.json(%{ + result: [%{"id" => activity.id |> FlakeId.from_string() |> Ecto.UUID.cast!()}] + }) + end) + + results = QdrantSearch.search(nil, "guys i just don't wanna leave the swamp", %{}) + + assert results == [activity] + end + + test "for a given actor, ask for only relevant matches" do + user = insert(:user) + + {:ok, activity} = + CommonAPI.post(user, %{ + status: "guys i just don't wanna leave the swamp", + visibility: "public" + }) + + Config + |> expect(:get, 3, 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) + + Tesla.Mock.mock(fn + %{url: "https://openai.url/v1/embeddings", method: :post} -> + Tesla.Mock.json(%{ + data: [%{embedding: [1, 2, 3]}] + }) + + %{url: "https://qdrant.url/collections/posts/points/search", method: :post, body: body} -> + data = Jason.decode!(body) + + assert data["filter"] == %{ + "must" => [%{"key" => "actor", "match" => %{"value" => user.ap_id}}] + } + + Tesla.Mock.json(%{ + result: [%{"id" => activity.id |> FlakeId.from_string() |> Ecto.UUID.cast!()}] + }) + end) + + results = + QdrantSearch.search(nil, "guys i just don't wanna leave the swamp", %{actor: user}) + + assert results == [activity] + end + test "indexes a public post on creation, deletes from the index on deletion" do user = insert(:user) @@ -29,7 +117,12 @@ test "indexes a public post on creation, deletes from the index on deletion" do %{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)) + data = Jason.decode!(body) + %{"points" => [%{"vector" => vector, "payload" => payload}]} = data + + assert vector == [1, 2, 3] + assert payload["actor"] + assert payload["published_at"] Tesla.Mock.json("ok") From a566ad56e1434715d00067b1e49be66b6787f5ba Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Thu, 23 May 2024 18:55:16 +0400 Subject: [PATCH 26/30] QdrantSearch: Fix actor / author restriction --- lib/pleroma/search/qdrant_search.ex | 4 ++-- test/pleroma/search/qdrant_search_test.exs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 9cb34ef71..19e8cd4bf 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -71,9 +71,9 @@ defp build_search_payload(embedding, options) do offset: options[:offset] || 0 } - if options[:actor] do + if author = options[:author] do Map.put(base, :filter, %{ - must: [%{key: "actor", match: %{value: options[:actor].ap_id}}] + must: [%{key: "actor", match: %{value: author.ap_id}}] }) else base diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index 371074dcf..46485392e 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -98,7 +98,7 @@ test "for a given actor, ask for only relevant matches" do end) results = - QdrantSearch.search(nil, "guys i just don't wanna leave the swamp", %{actor: user}) + QdrantSearch.search(nil, "guys i just don't wanna leave the swamp", %{author: user}) assert results == [activity] end From 8b76f56050a609bf562053cb7201a9204901490e Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Mon, 27 May 2024 13:57:42 +0400 Subject: [PATCH 27/30] QdrantSearch: Add healthcheck for qdrant --- lib/pleroma/search/qdrant_search.ex | 11 +++++++++++ test/pleroma/search/qdrant_search_test.exs | 12 ++++++++++++ 2 files changed, 23 insertions(+) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 19e8cd4bf..3c3ffce16 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -139,6 +139,17 @@ def search(_user, query, options) do [] 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() + + [qdrant_health] + end end defmodule Pleroma.Search.QdrantSearch.OpenAIClient do diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index 46485392e..b389aa816 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -15,6 +15,18 @@ defmodule Pleroma.Search.QdrantSearchTest do alias Pleroma.Workers.SearchIndexingWorker describe "Qdrant search" do + test "returns the correct healthcheck endpoints" do + Config + |> expect(:get, 1, fn + [Pleroma.Search.QdrantSearch, key], nil -> + %{qdrant_url: "https://qdrant.url"}[key] + end) + + health_endpoints = QdrantSearch.healthcheck_endpoints() + + assert "https://qdrant.url/healthz" in health_endpoints + end + test "searches for a term by encoding it and sending it to qdrant" do user = insert(:user) From ec3f3fef7798111641f08020d5fd7ae16e407b89 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Mon, 27 May 2024 14:15:04 +0400 Subject: [PATCH 28/30] Fastembed Server: Add health check endpoint --- supplemental/search/fastembed-api/fastembed-server.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/supplemental/search/fastembed-api/fastembed-server.py b/supplemental/search/fastembed-api/fastembed-server.py index dd4a7a9c8..02da69db2 100644 --- a/supplemental/search/fastembed-api/fastembed-server.py +++ b/supplemental/search/fastembed-api/fastembed-server.py @@ -17,6 +17,10 @@ def embeddings(request: EmbeddingRequest): embeddings = next(model.embed(request.input)).tolist() return {"data": [{"embedding": embeddings}]} +@app.get("/health") +def health(): + return {"status": "ok"} + if __name__ == "__main__": import uvicorn From f4c04e6b2dce6d75d148ca520aaef27005ecaa82 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Mon, 27 May 2024 14:21:55 +0400 Subject: [PATCH 29/30] QdrantSearch: Add health checks. --- config/config.exs | 3 +++ lib/pleroma/search/qdrant_search.ex | 4 +++- test/pleroma/search/qdrant_search_test.exs | 20 +++++++++++++++++--- 3 files changed, 23 insertions(+), 4 deletions(-) diff --git a/config/config.exs b/config/config.exs index d891a5218..f388dfe52 100644 --- a/config/config.exs +++ b/config/config.exs @@ -919,6 +919,9 @@ qdrant_url: "http://127.0.0.1:6333/", qdrant_api_key: "", openai_url: "http://127.0.0.1:11345", + # The healthcheck url has to be set to nil when used with the real openai + # API, as it doesn't have a healthcheck endpoint. + openai_healthcheck_url: "http://127.0.0.1:11345/health", openai_model: "snowflake/snowflake-arctic-embed-xs", openai_api_key: "", qdrant_index_configuration: %{ diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 3c3ffce16..429ae05b8 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -148,7 +148,9 @@ def healthcheck_endpoints do |> Map.put(:path, "/healthz") |> URI.to_string() - [qdrant_health] + openai_health = Config.get([Pleroma.Search.QdrantSearch, :openai_healthcheck_url]) + + [qdrant_health, openai_health] |> Enum.filter(& &1) end end diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs index b389aa816..47a77a391 100644 --- a/test/pleroma/search/qdrant_search_test.exs +++ b/test/pleroma/search/qdrant_search_test.exs @@ -16,15 +16,29 @@ defmodule Pleroma.Search.QdrantSearchTest do describe "Qdrant search" do test "returns the correct healthcheck endpoints" do + # No openai healthcheck URL Config - |> expect(:get, 1, fn + |> expect(:get, 2, fn [Pleroma.Search.QdrantSearch, key], nil -> %{qdrant_url: "https://qdrant.url"}[key] end) - health_endpoints = QdrantSearch.healthcheck_endpoints() + [health_endpoint] = QdrantSearch.healthcheck_endpoints() - assert "https://qdrant.url/healthz" in health_endpoints + assert "https://qdrant.url/healthz" == health_endpoint + + # Set openai healthcheck URL + Config + |> expect(:get, 2, fn + [Pleroma.Search.QdrantSearch, key], nil -> + %{qdrant_url: "https://qdrant.url", openai_healthcheck_url: "https://openai.url/health"}[ + key + ] + end) + + [_, health_endpoint] = QdrantSearch.healthcheck_endpoints() + + assert "https://openai.url/health" == health_endpoint end test "searches for a term by encoding it and sending it to qdrant" do From ddf103eca04c9571ba8310915556cc51cd4a9af8 Mon Sep 17 00:00:00 2001 From: Lain Soykaf Date: Mon, 27 May 2024 14:35:08 +0400 Subject: [PATCH 30/30] QdrantSearch: Fetch a post in search if possible. --- lib/pleroma/search/qdrant_search.ex | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex index 429ae05b8..b659bb682 100644 --- a/lib/pleroma/search/qdrant_search.ex +++ b/lib/pleroma/search/qdrant_search.ex @@ -9,6 +9,7 @@ defmodule Pleroma.Search.QdrantSearch do alias __MODULE__.QdrantClient import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] + import Pleroma.Search.DatabaseSearch, only: [maybe_fetch: 3] @impl true def create_index do @@ -115,8 +116,8 @@ def remove_from_index(object) do end @impl true - def search(_user, query, options) do - query = "Represent this sentence for searching relevant passages: #{query}" + 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}}} <- @@ -134,6 +135,7 @@ def search(_user, query, options) do |> Activity.restrict_deactivated_users() |> Ecto.Query.order_by([a], fragment("array_position(?, ?)", ^ids, a.id)) |> Pleroma.Repo.all() + |> maybe_fetch(user, original_query) else _ -> []