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)