B QdrantSearch: Switch to OpenAI api

This commit is contained in:
Lain Soykaf 2024-05-19 12:17:46 +04:00
parent cc1321ea2e
commit 72ec261a69
4 changed files with 26 additions and 17 deletions

View File

@ -1 +1 @@
Add Qdrant/Ollama search Add Qdrant/OpenAI embedding search

View File

@ -917,9 +917,10 @@ config :pleroma, Pleroma.Uploaders.Uploader, timeout: 30_000
config :pleroma, Pleroma.Search.QdrantSearch, config :pleroma, Pleroma.Search.QdrantSearch,
qdrant_url: "http://127.0.0.1:6333/", qdrant_url: "http://127.0.0.1:6333/",
qdrant_api_key: nil, qdrant_api_key: "",
ollama_url: "http://127.0.0.1:11434", openai_url: "http://127.0.0.1:11345",
ollama_model: "snowflake-arctic-embed:xs", openai_model: "snowflake",
openai_api_key: "",
qdrant_index_configuration: %{ qdrant_index_configuration: %{
vectors: %{size: 384, distance: "Cosine"} vectors: %{size: 384, distance: "Cosine"}
} }

View File

@ -5,7 +5,7 @@ defmodule Pleroma.Search.QdrantSearch do
alias Pleroma.Activity alias Pleroma.Activity
alias Pleroma.Config.Getting, as: Config alias Pleroma.Config.Getting, as: Config
alias __MODULE__.OllamaClient alias __MODULE__.OpenAIClient
alias __MODULE__.QdrantClient alias __MODULE__.QdrantClient
import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1] import Pleroma.Search.Meilisearch, only: [object_to_search_data: 1]
@ -31,10 +31,10 @@ defmodule Pleroma.Search.QdrantSearch do
end end
def get_embedding(text) do def get_embedding(text) do
with {:ok, %{body: %{"embedding" => embedding}}} <- with {:ok, %{body: %{"data" => [%{"embedding" => embedding}]}}} <-
OllamaClient.post("/api/embeddings", %{ OpenAIClient.post("/v1/embeddings", %{
prompt: text, input: text,
model: Config.get([Pleroma.Search.QdrantSearch, :ollama_model]) model: Config.get([Pleroma.Search.QdrantSearch, :openai_model])
}) do }) do
{:ok, embedding} {:ok, embedding}
else else
@ -119,12 +119,17 @@ defmodule Pleroma.Search.QdrantSearch do
end end
end end
defmodule Pleroma.Search.QdrantSearch.OllamaClient do defmodule Pleroma.Search.QdrantSearch.OpenAIClient do
use Tesla use Tesla
alias Pleroma.Config.Getting, as: Config 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.JSON)
plug(Tesla.Middleware.Headers, [
{"Authorization",
"Bearer #{Pleroma.Config.get([Pleroma.Search.QdrantSearch, :openai_api_key])}"}
])
end end
defmodule Pleroma.Search.QdrantSearch.QdrantClient do defmodule Pleroma.Search.QdrantSearch.QdrantClient do

View File

@ -19,9 +19,12 @@ defmodule Pleroma.Search.QdrantSearchTest do
user = insert(:user) user = insert(:user)
Tesla.Mock.mock(fn Tesla.Mock.mock(fn
%{method: :post, url: "https://ollama.url/api/embeddings"} -> %{method: :post, url: "https://openai.url/v1/embeddings"} ->
send(self(), "posted_to_ollama") send(self(), "posted_to_openai")
Tesla.Mock.json(%{embedding: [1, 2, 3]})
Tesla.Mock.json(%{
data: [%{embedding: [1, 2, 3]}]
})
%{method: :put, url: "https://qdrant.url/collections/posts/points", body: body} -> %{method: :put, url: "https://qdrant.url/collections/posts/points", body: body} ->
send(self(), "posted_to_qdrant") send(self(), "posted_to_qdrant")
@ -42,8 +45,8 @@ defmodule Pleroma.Search.QdrantSearchTest do
[Pleroma.Search.QdrantSearch, key], nil -> [Pleroma.Search.QdrantSearch, key], nil ->
%{ %{
ollama_model: "a_model", openai_model: "a_model",
ollama_url: "https://ollama.url", openai_url: "https://openai.url",
qdrant_url: "https://qdrant.url" qdrant_url: "https://qdrant.url"
}[key] }[key]
end) end)
@ -62,7 +65,7 @@ defmodule Pleroma.Search.QdrantSearchTest do
) )
assert :ok = perform_job(SearchIndexingWorker, args) assert :ok = perform_job(SearchIndexingWorker, args)
assert_received("posted_to_ollama") assert_received("posted_to_openai")
assert_received("posted_to_qdrant") assert_received("posted_to_qdrant")
{:ok, _} = CommonAPI.delete(activity.id, user) {:ok, _} = CommonAPI.delete(activity.id, user)