Basic AI config functions

lib/trento/ai.ex

@@ -0,0 +1,35 @@
+defmodule Trento.AI do
+  @moduledoc """
+  The `Trento.AI` module provides functions to interact with the AI features of the Trento application.
+  """
+
+  alias Trento.AI.ApplicationConfigLoader
+
+  alias Trento.AI.Configurations
+
+  @doc """
+  Checks if the AI features are enabled.
+  """
+  @spec enabled?() :: boolean()
+  def enabled?,
+    do: Keyword.get(ApplicationConfigLoader.load(), :enabled, false)
+
+  @doc """
+  Creates a user configuration for AI.
+
+  See `Trento.AI.Configurations.create_user_configuration/2` for more details.
+  """
+  def create_user_configuration(user, attrs),
+    do: configurations().create_user_configuration(user, attrs)
+
+  @doc """
+  Updates a user configuration for AI.
+
+  See `Trento.AI.Configurations.update_user_configuration/2` for more details.
+  """
+  def update_user_configuration(user, attrs),
+    do: configurations().update_user_configuration(user, attrs)
+
+  defp configurations,
+    do: Keyword.get(ApplicationConfigLoader.load(), :configurations, Configurations)
+end

lib/trento/ai/application_config_loader.ex

@@ -0,0 +1,21 @@
+defmodule Trento.AI.ApplicationConfigLoader do
+  @moduledoc """
+  This module is responsible for loading the AI application configuration and providing access to it.
+  """
+
+  @behaviour Trento.AI.ApplicationConfigLoader
+
+  @callback load_config :: keyword()
+
+  @impl true
+  def load_config, do: Application.get_env(:trento, :ai, [])
+
+  @spec load :: keyword()
+  def load, do: impl().load_config()
+
+  defp impl,
+    do:
+      :trento
+      |> Application.get_env(:ai, [])
+      |> Keyword.get(:application_config_loader, __MODULE__)
+end

lib/trento/ai/configurations.ex

@@ -0,0 +1,58 @@
+defmodule Trento.AI.Configurations do
+  @moduledoc """
+  This module is responsible for managing user AI configurations.
+  """
+
+  alias Trento.Users.User
+
+  alias Trento.AI.UserConfiguration
+
+  alias Trento.Repo
+
+  @doc """
+  Creates a user configuration for AI.
+
+  Only eligible users (not deleted or locked) can have an AI configuration.
+  """
+  @spec create_user_configuration(User.t(), map()) ::
+          {:ok, UserConfiguration.t()} | {:error, Ecto.Changeset.t()} | {:error, :forbidden}
+  def create_user_configuration(
+        %User{id: user_id, deleted_at: nil, locked_at: nil},
+        attrs
+      ) do
+    %UserConfiguration{}
+    |> UserConfiguration.changeset(Map.put(attrs, :user_id, user_id))
+    |> Repo.insert()
+  end
+
+  def create_user_configuration(%User{}, _), do: {:error, :forbidden}
+
+  @doc """
+  Updates a user configuration for AI.
+
+  Only eligible users (not deleted or locked) can update their AI configuration.
+
+  If the user does not have an existing configuration, an error will be returned.
+  """
+  @spec update_user_configuration(User.t(), map()) ::
+          {:ok, UserConfiguration.t()}
+          | {:error, Ecto.Changeset.t()}
+          | {:error, :forbidden | :not_found}
+  def update_user_configuration(
+        %User{id: user_id, deleted_at: nil, locked_at: nil},
+        attrs
+      )
+      when not is_nil(user_id) do
+    case Repo.get_by(UserConfiguration, user_id: user_id) do
+      nil ->
+        {:error, :not_found}
+
+      %UserConfiguration{} = user_configuration ->
+        user_configuration
+        |> UserConfiguration.changeset(attrs)
+        |> Repo.update()
+    end
+  end
+
+  def update_user_configuration(%User{}, _), do: {:error, :forbidden}
+end

lib/trento/ai/llm_registry.ex

@@ -0,0 +1,53 @@
+defmodule Trento.AI.LLMRegistry do
+  @moduledoc """
+  This module is responsible for managing the registry of available LLM providers and their models.
+  """
+
+  alias Trento.AI.ApplicationConfigLoader
+
+  @doc """
+  Returns the list of configured LLM providers.
+  """
+  @spec providers :: [atom()]
+  def providers, do: Keyword.keys(get_ai_providers_config())
+
+  @doc """
+  Returns the list of models for a given provider or all models if `:all` is passed.
+  """
+  @spec get_provider_models(atom() | :all) :: [bitstring()]
+  def get_provider_models(:all),
+    do:
+      Enum.flat_map(get_ai_providers_config(), fn {_provider, config} ->
+        Keyword.get(config, :models, [])
+      end)
+
+  def get_provider_models(provider) when is_atom(provider) do
+    get_ai_providers_config()
+    |> Keyword.get(provider, [])
+    |> Keyword.get(:models, [])
+  end
+
+  def get_provider_models(_), do: []
+
+  @doc """
+  Returns the provider for a given model or nil if the model is not supported.
+  """
+  @spec get_model_provider(bitstring()) :: atom() | nil
+  def get_model_provider(model) do
+    Enum.find_value(get_ai_providers_config(), fn {provider, config} ->
+      if model in Keyword.get(config, :models, []) do
+        provider
+      else
+        nil
+      end
+    end)
+  end
+
+  @doc """
+  Checks if a given model is supported by any provider.
+  """
+  @spec model_supported?(bitstring()) :: boolean()
+  def model_supported?(model), do: model in get_provider_models(:all)
+
+  defp get_ai_providers_config, do: Keyword.get(ApplicationConfigLoader.load(), :providers, [])
+end

lib/trento/ai/user_configuration.ex

@@ -0,0 +1,50 @@
+defmodule Trento.AI.UserConfiguration do
+  @moduledoc false
+
+  use Ecto.Schema
+  import Ecto.Changeset
+
+  alias Trento.Support.Ecto.EncryptedBinary
+
+  alias Trento.AI.LLMRegistry
+
+  @type t :: %__MODULE__{}
+
+  @primary_key false
+  schema "ai_configurations" do
+    field :model, :string
+    field :provider, Ecto.Enum, values: LLMRegistry.providers()
+    field :api_key, EncryptedBinary, redact: true
+
+    belongs_to :user, Trento.Users.User, primary_key: true
+
+    timestamps(type: :utc_datetime_usec)
+  end
+
+  def changeset(ai_configuration, attrs) do
+    updated_attrs = maybe_update_provider(attrs)
+
+    ai_configuration
+    |> cast(updated_attrs, [:user_id, :model, :provider, :api_key])
+    |> validate_required([:user_id, :model, :api_key])
+    |> validate_change(:model, &validate_model/2)
+    |> unique_constraint(:user_id,
+      name: :ai_configurations_pkey,
+      message: "User already has a configuration"
+    )
+    |> foreign_key_constraint(:user_id, message: "User does not exist")
+  end
+
+  defp maybe_update_provider(%{model: model} = attrs),
+    do: Map.put(attrs, :provider, LLMRegistry.get_model_provider(model))
+
+  defp maybe_update_provider(attrs), do: attrs
+
+  defp validate_model(_model_field_atom, model) do
+    if LLMRegistry.model_supported?(model) do
+      []
+    else
+      [model: {"is not supported", validation: :ai_model_validity}]
+    end
+  end
+end

lib/trento/users/user.ex

@@ -51,6 +51,8 @@ defmodule Trento.Users.User do
 
     has_many :personal_access_tokens, PersonalAccessToken, preload_order: [desc: :created_at]
 
+    has_one :ai_configuration, Trento.AI.UserConfiguration
+
     timestamps(type: :utc_datetime_usec)
   end
 

priv/repo/migrations/20260326100000_add_ai_user_configurations.exs

@@ -0,0 +1,14 @@
+defmodule Trento.Repo.Migrations.AddAiUserConfigurationsTable do
+  use Ecto.Migration
+
+  def change do
+    create table(:ai_configurations, primary_key: false) do
+      add :user_id, references(:users, on_delete: :delete_all), primary_key: true
+      add :provider, :string
+      add :model, :string
+      add :api_key, :binary
+
+      timestamps(type: :utc_datetime_usec)
+    end
+  end
+end

test/support/ai/ai_case.ex

@@ -0,0 +1,18 @@
+defmodule Trento.AI.AICase do
+  @moduledoc false
+
+  use ExUnit.CaseTemplate
+
+  setup _ do
+    stub_config_loader()
+
+    :ok
+  end
+
+  def stub_config_loader do
+    Mox.stub_with(
+      Trento.AI.ApplicationConfigLoader.Mock,
+      Trento.AI.ApplicationConfigLoader
+    )
+  end
+end

test/support/factory.ex

@@ -182,6 +182,8 @@ defmodule Trento.Factory do
 
   alias TrentoWeb.Auth.PersonalAccessToken, as: PAT
 
+  alias Trento.AI.{LLMRegistry, UserConfiguration}
+
   use ExMachina.Ecto, repo: Trento.Repo
 
   def host_registered_event_factory do
@@ -1458,4 +1460,25 @@ defmodule Trento.Factory do
       target_type: Enum.random(["host", "cluster"])
     }
   end
+
+  def random_ai_model do
+    Trento.AI.AICase.stub_config_loader()
+
+    :all
+    |> LLMRegistry.get_provider_models()
+    |> Enum.random()
+  end
+
+  def ai_user_configuration_factory(attrs) do
+    user_id = Map.get(attrs, :user_id, 1)
+    model = Map.get(attrs, :model, random_ai_model())
+
+    %UserConfiguration{}
+    |> UserConfiguration.changeset(%{
+      user_id: user_id,
+      model: model,
+      api_key: Faker.String.base64(32)
+    })
+    |> Ecto.Changeset.apply_changes()
+  end
 end

test/test_helper.exs

@@ -53,6 +53,19 @@ Mox.defmock(Trento.Support.DateService.Mock, for: Trento.Support.DateService)
 Mox.defmock(Joken.CurrentTime.Mock, for: Joken.CurrentTime)
 Application.put_env(:joken, :current_time_adapter, Joken.CurrentTime.Mock)
 
+Mox.defmock(Trento.AI.ApplicationConfigLoader.Mock,
+  for: Trento.AI.ApplicationConfigLoader
+)
+
+test_ai_config =
+  Keyword.put(
+    Application.get_env(:trento, :ai),
+    :application_config_loader,
+    Trento.AI.ApplicationConfigLoader.Mock
+  )
+
+Application.put_env(:trento, :ai, test_ai_config)
+
 Application.ensure_all_started(:ex_machina, :faker)
 
 if Application.get_env(:trento, :flaky_tests_detection)[:enabled?] == true do