118 lines
3.0 KiB
Elixir
118 lines
3.0 KiB
Elixir
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
|