From f4c04e6b2dce6d75d148ca520aaef27005ecaa82 Mon Sep 17 00:00:00 2001
From: Lain Soykaf <lain@lain.com>
Date: Mon, 27 May 2024 14:21:55 +0400
Subject: [PATCH] QdrantSearch: Add health checks.

---
 config/config.exs                          |  3 +++
 lib/pleroma/search/qdrant_search.ex        |  4 +++-
 test/pleroma/search/qdrant_search_test.exs | 20 +++++++++++++++++---
 3 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/config/config.exs b/config/config.exs
index d891a5218..f388dfe52 100644
--- a/config/config.exs
+++ b/config/config.exs
@@ -919,6 +919,9 @@
   qdrant_url: "http://127.0.0.1:6333/",
   qdrant_api_key: "",
   openai_url: "http://127.0.0.1:11345",
+  # The healthcheck url has to be set to nil when used with the real openai
+  # API, as it doesn't have a healthcheck endpoint.
+  openai_healthcheck_url: "http://127.0.0.1:11345/health",
   openai_model: "snowflake/snowflake-arctic-embed-xs",
   openai_api_key: "",
   qdrant_index_configuration: %{
diff --git a/lib/pleroma/search/qdrant_search.ex b/lib/pleroma/search/qdrant_search.ex
index 3c3ffce16..429ae05b8 100644
--- a/lib/pleroma/search/qdrant_search.ex
+++ b/lib/pleroma/search/qdrant_search.ex
@@ -148,7 +148,9 @@ def healthcheck_endpoints do
       |> Map.put(:path, "/healthz")
       |> URI.to_string()
 
-    [qdrant_health]
+    openai_health = Config.get([Pleroma.Search.QdrantSearch, :openai_healthcheck_url])
+
+    [qdrant_health, openai_health] |> Enum.filter(& &1)
   end
 end
 
diff --git a/test/pleroma/search/qdrant_search_test.exs b/test/pleroma/search/qdrant_search_test.exs
index b389aa816..47a77a391 100644
--- a/test/pleroma/search/qdrant_search_test.exs
+++ b/test/pleroma/search/qdrant_search_test.exs
@@ -16,15 +16,29 @@ defmodule Pleroma.Search.QdrantSearchTest do
 
   describe "Qdrant search" do
     test "returns the correct healthcheck endpoints" do
+      # No openai healthcheck URL
       Config
-      |> expect(:get, 1, fn
+      |> expect(:get, 2, fn
         [Pleroma.Search.QdrantSearch, key], nil ->
           %{qdrant_url: "https://qdrant.url"}[key]
       end)
 
-      health_endpoints = QdrantSearch.healthcheck_endpoints()
+      [health_endpoint] = QdrantSearch.healthcheck_endpoints()
 
-      assert "https://qdrant.url/healthz" in health_endpoints
+      assert "https://qdrant.url/healthz" == health_endpoint
+
+      # Set openai healthcheck URL
+      Config
+      |> expect(:get, 2, fn
+        [Pleroma.Search.QdrantSearch, key], nil ->
+          %{qdrant_url: "https://qdrant.url", openai_healthcheck_url: "https://openai.url/health"}[
+            key
+          ]
+      end)
+
+      [_, health_endpoint] = QdrantSearch.healthcheck_endpoints()
+
+      assert "https://openai.url/health" == health_endpoint
     end
 
     test "searches for a term by encoding it and sending it to qdrant" do