Temporarily limit AI Assistant to OpenFn users (#2482)

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.
   """

lib/lightning_web/live/workflow_live/edit.ex

@@ -197,7 +197,10 @@ defmodule LightningWeb.WorkflowLive.Edit do
                     <:tab hash="manual">
                       <span class="inline-block align-middle">Input</span>
                     </:tab>
-                    <:tab hash="aichat">
+                    <:tab
+                      :if={Lightning.AiAssistant.available?(@current_user)}
+                      hash="aichat"
+                    >
                       <span class="inline-block align-middle">AI Assistant</span>
                     </:tab>
                   </LightningWeb.Components.Tabbed.tabs>
@@ -226,7 +229,11 @@ defmodule LightningWeb.WorkflowLive.Edit do
                       />
                     </div>
                   </:panel>
-                  <:panel hash="aichat" class="h-full">
+                  <:panel
+                    :if={Lightning.AiAssistant.available?(@current_user)}
+                    hash="aichat"
+                    class="h-full"
+                  >
                     <%= if @ai_assistant_enabled do %>
                       <div class="grow min-h-0 h-full text-sm">
                         <.live_component

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: "[email protected]")
+      assert AiAssistant.available?(user)
+
+      user = build(:user, email: "[email protected]")
+      refute AiAssistant.available?(user)
+    end
+  end
+
   describe "endpoint_available?" do
     test "availability" do
       Mox.stub(Lightning.MockConfig, :apollo, fn key ->

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: "[email protected]"
     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: "[email protected]"
     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: "[email protected]"
     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: "[email protected]"
     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: "[email protected]"
     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: "[email protected]"
     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: "[email protected]"
     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: "[email protected]"
     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"

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: "[email protected]"
   """
-  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