From cc44439e128f6d2b5c209c4e94e93c0ce720228e Mon Sep 17 00:00:00 2001 From: Stuart Corbishley Date: Mon, 9 Sep 2024 17:43:37 +0200 Subject: [PATCH] Temporarily limit AI Assistant to OpenFn users (#2482) --- lib/lightning/ai_assistant/ai_assistant.ex | 4 ++ lib/lightning_web/live/workflow_live/edit.ex | 11 ++++- .../ai_assistant/ai_assistant_test.exs | 10 +++++ .../live/workflow_live/edit_test.exs | 41 +++++++++++++++---- test/support/conn_case.ex | 10 ++++- 5 files changed, 65 insertions(+), 11 deletions(-) diff --git a/lib/lightning/ai_assistant/ai_assistant.ex b/lib/lightning/ai_assistant/ai_assistant.ex index f0fc809bd9..a1598441f9 100644 --- a/lib/lightning/ai_assistant/ai_assistant.ex +++ b/lib/lightning/ai_assistant/ai_assistant.ex @@ -142,6 +142,10 @@ defmodule Lightning.AiAssistant do is_binary(endpoint) && is_binary(api_key) end + def available?(user) do + String.match?(user.email, ~r/@openfn\.org/i) + end + @doc """ Checks if the Apollo endpoint is available. """ diff --git a/lib/lightning_web/live/workflow_live/edit.ex b/lib/lightning_web/live/workflow_live/edit.ex index a80afd072b..2b972a732b 100644 --- a/lib/lightning_web/live/workflow_live/edit.ex +++ b/lib/lightning_web/live/workflow_live/edit.ex @@ -197,7 +197,10 @@ defmodule LightningWeb.WorkflowLive.Edit do <:tab hash="manual"> Input - <:tab hash="aichat"> + <:tab + :if={Lightning.AiAssistant.available?(@current_user)} + hash="aichat" + > AI Assistant @@ -226,7 +229,11 @@ defmodule LightningWeb.WorkflowLive.Edit do /> - <:panel hash="aichat" class="h-full"> + <:panel + :if={Lightning.AiAssistant.available?(@current_user)} + hash="aichat" + class="h-full" + > <%= if @ai_assistant_enabled do %>
<.live_component diff --git a/test/lightning/ai_assistant/ai_assistant_test.exs b/test/lightning/ai_assistant/ai_assistant_test.exs index d6d392e459..761f565210 100644 --- a/test/lightning/ai_assistant/ai_assistant_test.exs +++ b/test/lightning/ai_assistant/ai_assistant_test.exs @@ -13,6 +13,16 @@ defmodule Lightning.AiAssistantTest do [user: user, project: project, workflow: workflow] end + describe "available?/1" do + test "is not available for users without openfn.org email" do + user = build(:user, email: "test@openFn.org") + assert AiAssistant.available?(user) + + user = build(:user, email: "test@example.net") + refute AiAssistant.available?(user) + end + end + describe "endpoint_available?" do test "availability" do Mox.stub(Lightning.MockConfig, :apollo, fn key -> diff --git a/test/lightning_web/live/workflow_live/edit_test.exs b/test/lightning_web/live/workflow_live/edit_test.exs index 89256dfe8d..ca1e8185bb 100644 --- a/test/lightning_web/live/workflow_live/edit_test.exs +++ b/test/lightning_web/live/workflow_live/edit_test.exs @@ -1375,6 +1375,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do describe "AI Assistant:" do setup :create_workflow + @tag email: "user@openfn.org" test "correct information is displayed when the assistant is not configured", %{ conn: conn, @@ -1418,6 +1419,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do "AI Assistant has not been configured for your instance" end + @tag email: "user@openfn.org" test "onboarding ui is displayed when no session exists for the project", %{ conn: conn, project: project, @@ -1470,6 +1472,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do assert has_element?(view, "#ai-assistant-form") end + @tag email: "user@openfn.org" test "authorized users can send a message", %{ conn: conn, project: project, @@ -1504,8 +1507,18 @@ defmodule LightningWeb.WorkflowLive.EditTest do # insert session so that the onboarding flow is not displayed insert(:chat_session, user: user, job: job_1) - for {conn, _user} <- - setup_project_users(conn, project, [:owner, :admin, :editor]) do + [:owner, :admin, :editor] + |> Enum.map(fn role -> + user = + insert(:user, email: "email-#{Enum.random(1..1_000)}@openfn.org") + + insert(:project_user, project: project, user: user, role: role) + + user + end) + |> Enum.each(fn user -> + conn = log_in_user(conn, user) + {:ok, view, _html} = live( conn, @@ -1535,9 +1548,20 @@ defmodule LightningWeb.WorkflowLive.EditTest do refute html =~ "You are not authorized to use the Ai Assistant" assert_patch(view) - end + end) + + [:viewer] + |> Enum.map(fn role -> + user = + insert(:user, email: "email-#{Enum.random(1..1_000)}@openfn.org") + + insert(:project_user, project: project, user: user, role: role) + + user + end) + |> Enum.each(fn user -> + conn = log_in_user(conn, user) - for {conn, _user} <- setup_project_users(conn, project, [:viewer]) do {:ok, view, _html} = live( conn, @@ -1565,9 +1589,10 @@ defmodule LightningWeb.WorkflowLive.EditTest do |> render_submit(%{content: "Hello"}) assert html =~ "You are not authorized to use the Ai Assistant" - end + end) end + @tag email: "user@openfn.org" test "users can start a new session", %{ conn: conn, project: project, @@ -1640,6 +1665,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do assert html =~ "Pong" end + @tag email: "user@openfn.org" test "users can resume a session", %{ conn: conn, project: project, @@ -1719,6 +1745,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do assert html =~ expected_answer end + @tag email: "user@openfn.org" test "an error is displayed incase the assistant does not return 200", %{ conn: conn, project: project, @@ -1770,6 +1797,7 @@ defmodule LightningWeb.WorkflowLive.EditTest do "Oops! Could not reach the Ai Server. Please try again later." end + @tag email: "user@openfn.org" test "an error is displayed incase the assistant query process crashes", %{ conn: conn, project: project, @@ -1821,13 +1849,12 @@ defmodule LightningWeb.WorkflowLive.EditTest do "Oops! Something went wrong. Please try again." end + @tag email: "user@openfn.org" test "shows a flash error when limit has reached", %{ conn: conn, project: %{id: project_id} = project, workflow: %{jobs: [job_1 | _]} = workflow } do - [{conn, _user}] = setup_project_users(conn, project, [:owner]) - Mox.stub(Lightning.MockConfig, :apollo, fn :endpoint -> "http://localhost:4001/health_check" :openai_api_key -> "openai_api_key" diff --git a/test/support/conn_case.ex b/test/support/conn_case.ex index 01892aebcd..5235ee8d19 100644 --- a/test/support/conn_case.ex +++ b/test/support/conn_case.ex @@ -79,9 +79,15 @@ defmodule LightningWeb.ConnCase do It stores an updated connection and a registered user in the test context. + + Optionally you may provide the current user's email address by providing + an `email` tag in the context. + + @tag email: "my@user.com" """ - def register_and_log_in_user(%{conn: conn}) do - user = Lightning.AccountsFixtures.user_fixture() + def register_and_log_in_user(%{conn: conn} = tags) do + attrs = Map.take(tags, [:email]) + user = Lightning.Factories.insert(:user, attrs) %{conn: log_in_user(conn, user), user: user} end