From 72ec261a69a7dda7ab95667e425824ab7758b636 Mon Sep 17 00:00:00 2001
From: Lain Soykaf <lain@lain.com>
Date: Sun, 19 May 2024 12:17:46 +0400
Subject: [PATCH] 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)