AI Assistant: Add Metadata Column to Sessions (#3095)

* Add RAG results column and prompt column

* meta key in chat sessions instead of rag_results and prompt key in chat messages

* Update CL

* Rewording in CL

* When sending user message use the session meta value

Author: elias-ba

CHANGELOG.md

@@ -17,6 +17,9 @@ and this project adheres to
 
 ### Added
 
+- AI Assistant: add metadata column to chat sessions
+  [#3054](https://github.com/OpenFn/lightning/issues/3054)
+  
 ### Changed
 
 ### Fixed

lib/lightning/ai_assistant/ai_assistant.ex

@@ -92,16 +92,18 @@ defmodule Lightning.AiAssistant do
 
   @spec save_message(ChatSession.t(), %{any() => any()}) ::
           {:ok, ChatSession.t()} | {:error, Ecto.Changeset.t()}
-  def save_message(session, message, usage \\ %{}) do
+  def save_message(session, message, usage \\ %{}, meta \\ nil) do
     messages = Enum.map(session.messages, &Map.take(&1, [:id]))
 
     Multi.new()
     |> Multi.put(:usage, usage)
     |> Multi.put(:message, message)
     |> Multi.insert_or_update(
       :upsert,
-      session
-      |> ChatSession.changeset(%{messages: messages ++ [message]})
+      ChatSession.changeset(session, %{
+        messages: messages ++ [message],
+        meta: meta || session.meta
+      })
     )
     |> Multi.merge(&maybe_increment_ai_usage/1)
     |> Repo.transaction()
@@ -131,7 +133,8 @@ defmodule Lightning.AiAssistant do
     ApolloClient.query(
       content,
       %{expression: session.expression, adaptor: session.adaptor},
-      build_history(session)
+      build_history(session),
+      session.meta || %{}
     )
     |> handle_apollo_resp(session)
   end
@@ -142,7 +145,7 @@ defmodule Lightning.AiAssistant do
        )
        when status in 200..299 do
     message = body["history"] |> Enum.reverse() |> hd()
-    save_message(session, message, body["usage"])
+    save_message(session, message, body["usage"], body["meta"])
   end
 
   defp handle_apollo_resp(

lib/lightning/ai_assistant/chat_session.ex

@@ -17,13 +17,15 @@ defmodule Lightning.AiAssistant.ChatSession do
           adaptor: String.t() | nil,
           is_public: boolean(),
           is_deleted: boolean(),
+          meta: map() | nil,
           messages: [ChatMessage.t(), ...] | []
         }
 
   schema "ai_chat_sessions" do
     field :expression, :string, virtual: true
     field :adaptor, :string, virtual: true
     field :title, :string
+    field :meta, :map
     field :is_public, :boolean, default: false
     field :is_deleted, :boolean, default: false
     belongs_to :job, Job
@@ -36,7 +38,7 @@ defmodule Lightning.AiAssistant.ChatSession do
 
   def changeset(chat_session, attrs) do
     chat_session
-    |> cast(attrs, [:title, :is_public, :is_deleted, :job_id, :user_id])
+    |> cast(attrs, [:title, :is_public, :is_deleted, :job_id, :user_id, :meta])
     |> validate_required([:title, :job_id, :user_id])
     |> cast_assoc(:messages)
   end

lib/lightning/apollo_client.ex

@@ -9,13 +9,14 @@ defmodule Lightning.ApolloClient do
           }
           | %{}
 
-  @spec query(String.t(), context(), list()) :: Tesla.Env.result()
-  def query(content, context \\ %{}, history \\ []) do
+  @spec query(String.t(), context(), list(), map()) :: Tesla.Env.result()
+  def query(content, context \\ %{}, history \\ [], meta \\ %{}) do
     payload = %{
       "api_key" => Lightning.Config.apollo(:ai_assistant_api_key),
       "content" => content,
       "context" => context,
-      "history" => history
+      "history" => history,
+      "meta" => meta
     }
 
     client() |> Tesla.post("/services/job_chat", payload)

priv/repo/migrations/20250408161805_add_meta_to_ai_chat_sessions.exs

@@ -0,0 +1,9 @@
+defmodule Lightning.Repo.Migrations.AddMetaToAiChatSessions do
+  use Ecto.Migration
+
+  def change do
+    alter table(:ai_chat_sessions) do
+      add :meta, :jsonb
+    end
+  end
+end

test/lightning/apollo_client_test.exs

@@ -23,7 +23,8 @@ defmodule Lightning.ApolloClientTest do
                "api_key" => "api_key",
                "content" => "foo",
                "context" => %{},
-               "history" => []
+               "history" => [],
+               "meta" => %{}
              }
 
       {:ok,