Page MenuHomePhorge
Files F85805840

No OneTemporary
Actions

Size
10 KB
Referenced Files
None
Subscribers
None

View Options

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
@@ -1,167 +1,178 @@
defmodule Pleroma.Search.QdrantSearch do
@behaviour Pleroma.Search.SearchBackend
import Ecto.Query
alias Pleroma.Activity
alias Pleroma.Config.Getting, as: Config
alias __MODULE__.OpenAIClient
alias __MODULE__.QdrantClient
import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1]
@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) do
# This will only index public or unlisted notes
maybe_search_data = object_to_search_data(activity.object)
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 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, 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()
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()
+
+ [qdrant_health]
+ end
end
defmodule Pleroma.Search.QdrantSearch.OpenAIClient do
use Tesla
alias Pleroma.Config.Getting, as: Config
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, 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
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
@@ -1,173 +1,185 @@
# Pleroma: A lightweight social networking server
# Copyright © 2017-2021 Pleroma Authors <https://pleroma.social/>
# 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 "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)
{: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", %{author: user})
assert results == [activity]
end
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")
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")
%{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

File Metadata

Mime Type
text/x-diff
Expires
Mon, Oct 12, 5:48 AM (1 d, 20 h)
Storage Engine
blob
Storage Format
Raw Data
Storage Handle
1786057
Default Alt Text
(10 KB)

Event Timeline