init

enterprise/backend/src/metabase_enterprise/metabot_v3/agent/profiles.clj

@@ -204,6 +204,7 @@
 (defn get-tools-for-profile
   "Get tool registry filtered by profile configuration and user capabilities."
   [profile-id capabilities]
+  (def tsp-caps capabilities)
   (some-> (get-profile profile-id)
           :tools
           (filter-by-capabilities capabilities)

enterprise/backend/src/metabase_enterprise/metabot_v3/agent/prompts.clj

@@ -176,7 +176,9 @@
   [tools]
   (vec
    (for [[tool-name tool] tools
-         :let [fname  (or (-> tool meta :prompt)
+         :let [fname  (or (if (map? tool)
+                            (:prompt tool)
+                            (-> tool meta :prompt))
                           (str tool-name ".md"))
                prompt (some-> (io/resource (str "metabot/prompts/tools/" fname))
                               slurp)]
@@ -204,6 +206,7 @@
             dialect-instructions (when sql-dialect
                                    (get-cached-dialect-instructions sql-dialect))
             tool-instructions    (extract-tool-instructions tools)
+            _ (def tsp-tool-instructions tool-instructions)
             current-user-info    (or (get context :current_user_info)
                                      (get context :current-user-info))
             current-time         (or (get context :current_time)

enterprise/backend/src/metabase_enterprise/metabot_v3/agent/tools/transforms.clj

@@ -87,6 +87,7 @@
               format-python-library-output))
 
 (mu/defn ^{:tool-name "write_transform_sql"
+           :prompt "write_transform_sql.md"
            :capabilities #{:feature-transforms :permission-write-transforms}}
   write-transform-sql-tool
   "Write new SQL queries or edit existing queries for transforms.
@@ -97,17 +98,16 @@
 
   For edit mode, provide edits as an array of {old_string, new_string, replace_all} objects.
   For replace mode, provide new_content with the complete SQL."
-  [{:keys [transform_id edit_action thinking transform_name transform_description
+  [{:keys [transform_id mode edits new_content thinking transform_name transform_description
            source_database source_tables]}
    :- [:map {:closed true}
-       [:transform_id {:optional true} [:maybe :int]]
-       [:edit_action [:map
-                      [:mode [:enum "edit" "replace"]]
-                      [:edits {:optional true} [:maybe [:sequential [:map
-                                                                     [:old_string :string]
-                                                                     [:new_string :string]
-                                                                     [:replace_all {:optional true} [:maybe :boolean]]]]]]
-                      [:new_content {:optional true} [:maybe :string]]]]
+       [:transform_id {:optional true} [:maybe [:or :int :string]]]
+       [:mode [:enum "edit" "replace"]]
+       [:edits {:optional true} [:maybe [:sequential [:map
+                                                      [:old_string :string]
+                                                      [:new_string :string]
+                                                      [:replace_all {:optional true} [:maybe :boolean]]]]]]
+       [:new_content {:optional true} [:maybe :string]]
        [:thinking {:optional true} [:maybe :string]]
        [:transform_name {:optional true} [:maybe :string]]
        [:transform_description {:optional true} [:maybe :string]]
@@ -117,7 +117,9 @@
     (let [result (add-output
                   (transforms-write-tools/write-transform-sql
                    {:transform_id transform_id
-                    :edit_action edit_action
+                    :mode mode
+                    :edits edits
+                    :new_content new_content
                     :thinking thinking
                     :transform_name transform_name
                     :transform_description transform_description
@@ -136,6 +138,7 @@
         {:output (str "Failed to write SQL transform: " (or (ex-message e) "Unknown error"))}))))
 
 (mu/defn ^{:tool-name "write_transform_python"
+           :prompt "write_transform_python.md"
            :capabilities #{:feature-transforms :feature-transforms-python :permission-write-transforms}}
   write-transform-python-tool
   "Write new Python code or edit existing code for transforms.
@@ -150,17 +153,16 @@
   Use `get_transform_python_library_details` before writing any Python code to inspect the shared library.
   Use the shared library in your code by adding `import common` at the top of the file.
   Keep `import common` at the top of the file even if it is currently unused."
-  [{:keys [transform_id edit_action thinking transform_name transform_description
+  [{:keys [transform_id mode edits new_content thinking transform_name transform_description
            source_database source_tables]}
    :- [:map {:closed true}
-       [:transform_id {:optional true} [:maybe :int]]
-       [:edit_action [:map
-                      [:mode [:enum "edit" "replace"]]
-                      [:edits {:optional true} [:maybe [:sequential [:map
-                                                                     [:old_string :string]
-                                                                     [:new_string :string]
-                                                                     [:replace_all {:optional true} [:maybe :boolean]]]]]]
-                      [:new_content {:optional true} [:maybe :string]]]]
+       [:transform_id {:optional true} [:maybe [:or :int :string]]]
+       [:mode [:enum "edit" "replace"]]
+       [:edits {:optional true} [:maybe [:sequential [:map
+                                                      [:old_string :string]
+                                                      [:new_string :string]
+                                                      [:replace_all {:optional true} [:maybe :boolean]]]]]]
+       [:new_content {:optional true} [:maybe :string]]
        [:thinking {:optional true} [:maybe :string]]
        [:transform_name {:optional true} [:maybe :string]]
        [:transform_description {:optional true} [:maybe :string]]
@@ -170,7 +172,9 @@
     (add-output
      (transforms-write-tools/write-transform-python
       {:transform_id transform_id
-       :edit_action edit_action
+       :mode mode
+       :edits edits
+       :new_content new_content
        :thinking thinking
        :transform_name transform_name
        :transform_description transform_description

enterprise/backend/src/metabase_enterprise/metabot_v3/context.clj

@@ -9,6 +9,7 @@
    [metabase.lib-be.core :as lib-be]
    [metabase.lib.core :as lib]
    [metabase.transforms-base.util :as transforms-base.u]
+   [metabase.transforms.feature-gating :as transforms.gating]
    [metabase.util.json :as json]
    [metabase.util.log :as log]
    [metabase.util.malli :as mu]
@@ -202,10 +203,24 @@
       (assoc context :user_is_viewing annotated-viewing))
     context))
 
+(defn- feature-capabilities
+  "Returns a set of feature capability strings based on enabled features.
+  These correspond to the capability keywords used in tool :capabilities metadata:
+  - \"feature:transforms\" -> :feature-transforms
+  - \"feature:transforms-python\" -> :feature-transforms-python
+  - \"feature:snippets\" -> :feature-snippets"
+  []
+  (cond-> #{"feature:snippets"}  ; snippets are always available
+    (transforms.gating/any-transforms-enabled?)    (conj "feature:transforms")
+    (transforms.gating/python-transforms-enabled?) (conj "feature:transforms-python")))
+
 (defn- add-backend-capabilities
-  "Add backend capabilities to context, merging with any existing capabilities."
+  "Add backend capabilities to context, merging with any existing capabilities.
+  Includes both API endpoint capabilities and feature flag capabilities."
   [context]
-  (update context :capabilities (fnil into #{}) (backend-metabot-capabilities)))
+  (update context :capabilities (fnil into #{})
+          (concat (backend-metabot-capabilities)
+                  (feature-capabilities))))
 
 (defn- add-recent-views
   "Add user's recent views to the context since these have a higher likelihood of being relevant to a user's query.

enterprise/backend/src/metabase_enterprise/metabot_v3/self/openrouter.clj

@@ -231,6 +231,7 @@
                                       :description "Output structured data"
                                       :parameters  schema}}])
                        (seq (mapv tool->openai-chat tools)))
+        _ (def tsp-tools all-tools)
         req        (cond-> {:model             model
                             :stream            true
                             :stream_options    {:include_usage true}

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/transforms_write.clj

@@ -4,6 +4,7 @@
   (:require
    [clojure.string :as str]
    [metabase-enterprise.metabot-v3.agent.streaming :as streaming]
+   [metabase.util :as u]
    [metabase.util.log :as log]))
 
 (set! *warn-on-reflection* true)
@@ -63,21 +64,37 @@ def transform():
     return pd.DataFrame([{\"message\": \"Hello from Python transform!\"}])
 ")
 
-(defn- create-fresh-transform
-  "Create a fresh transform structure with the given source type."
-  [source-type transform-name transform-description source-database source-tables]
-  (let [template (case source-type
-                   :sql fresh-sql-template
-                   :python fresh-python-template)]
-    {:id nil
-     :name (or transform-name (str "New " (name source-type) " Transform"))
-     :description (or transform-description "")
-     :source {:type (name source-type)
-              :query template
-              :source-database source-database
-              :source-tables (or source-tables {})}
-     :target {:type "table"
-              :name ""}}))
+(defn- create-fresh-sql-transform
+  "Create a fresh SQL transform structure.
+  The source uses type 'query' with a DatasetQuery object containing native SQL."
+  [transform-name transform-description source-database]
+  {:id nil
+   :name (or transform-name "New SQL Transform")
+   :description (or transform-description "")
+   :source {:type "query"
+            :query {:type "native"
+                    :database source-database
+                    :native {:query fresh-sql-template}}}
+   :target {:type "table"
+            :name ""
+            :schema nil
+            :database source-database}})
+
+(defn- create-fresh-python-transform
+  "Create a fresh Python transform structure.
+  The source uses type 'python' with the code in the 'body' field."
+  [transform-name transform-description source-database source-tables]
+  {:id nil
+   :name (or transform-name "New Python Transform")
+   :description (or transform-description "")
+   :source {:type "python"
+            :body fresh-python-template
+            :source-database source-database
+            :source-tables (or source-tables [])}
+   :target {:type "table"
+            :name ""
+            :schema nil
+            :database source-database}})
 
 ;;; Write Transform SQL Tool
 
@@ -86,7 +103,9 @@ def transform():
 
   Parameters:
   - transform-id: Optional ID of existing transform to edit
-  - edit-action: Either {:mode \"edit\" :edits [...]} or {:mode \"replace\" :new-content \"...\"}
+  - mode: Either \"edit\" or \"replace\"
+  - edits: For edit mode, array of {:old_string :new_string :replace_all} maps
+  - new_content: For replace mode, the complete content
   - thinking: Explanation of what changes are being made
   - transform-name: Name for new transform (required if no transform-id)
   - transform-description: Description for new transform
@@ -98,26 +117,30 @@ def transform():
   Returns map with:
   - :structured-output - The suggested transform
   - :data-parts - Transform suggestion data part for streaming"
-  [{:keys [transform_id edit_action thinking transform_name transform_description
+  [{:keys [transform_id mode edits new_content thinking transform_name transform_description
            source_database source_tables memory-atom context]}]
   (log/info "Writing SQL transform" {:transform-id transform_id
-                                     :edit-mode (:mode edit_action)
+                                     :edit-mode mode
                                      :has-context (some? context)})
+  (def tsp-memory memory-atom)
+  (def tsp-context context)
+  (def tsp-transform-id transform_id)
+  (let [ ;; Get current transform from context or memory, or create fresh one
+        ;; For new transforms, generate a temporary ID so the LLM can reference it later
+        is-new-transform? (nil? transform_id)
+        effective-transform-id (or transform_id (u/generate-nano-id))
 
-  (let [;; Get current transform from context or memory, or create fresh one
-        current-transform (cond
-                            ;; Try to get from context first
-                            (and transform_id context)
-                            (get-in context [:transforms (str transform_id)])
-
-                            ;; Then try memory
-                            (and transform_id memory-atom)
-                            (get-in @memory-atom [:state :transforms (str transform_id)])
-
-                            ;; Create fresh transform
-                            :else
-                            (create-fresh-transform :sql transform_name transform_description
-                                                    source_database source_tables))
+        current-transform (or
+                           ;; Try to get from context first (for transforms user is viewing)
+                           (when (and transform_id context)
+                             (get-in context [:transforms (str transform_id)]))
+                           ;; Then try memory (for transforms created in session)
+                           (when (and transform_id memory-atom)
+                             (get-in @memory-atom [:state :transforms (str transform_id)]))
+                           ;; Create fresh transform if no transform_id
+                           (when-not transform_id
+                             (create-fresh-sql-transform transform_name transform_description
+                                                         source_database)))
 
         _ (when (and transform_id (nil? current-transform))
             (throw (ex-info (str "Transform with ID " transform_id " not found. "
@@ -125,39 +148,48 @@ def transform():
                             {:agent-error? true
                              :transform-id transform_id})))
 
-        current-sql (get-in current-transform [:source :query] "")
+        ;; Extract current SQL from the DatasetQuery structure
+        ;; Path is [:source :query :native :query] for native SQL transforms
+        current-sql (get-in current-transform [:source :query :native :query] "")
 
         ;; Apply edits based on mode
-        new-sql (case (:mode edit_action)
-                  "edit" (apply-edits current-sql (:edits edit_action))
-                  "replace" (:new_content edit_action)
+        new-sql (case mode
+                  "edit" (apply-edits current-sql edits)
+                  "replace" new_content
                   ;; Default to replace if no mode specified but new_content provided
-                  (if (:new_content edit_action)
-                    (:new_content edit_action)
-                    (throw (ex-info "Invalid edit_action: must specify mode 'edit' with edits or mode 'replace' with new_content"
+                  (if new_content
+                    new_content
+                    (throw (ex-info "Invalid input: must specify mode 'edit' with edits or mode 'replace' with new_content"
                                     {:agent-error? true
-                                     :edit-action edit_action}))))
+                                     :mode mode}))))
 
-        ;; Build suggested transform
+        ;; Build suggested transform - update SQL in the nested DatasetQuery structure
         suggested-transform (cond-> current-transform
-                              true (assoc-in [:source :query] new-sql)
+                              true (assoc-in [:source :query :native :query] new-sql)
                               transform_name (assoc :name transform_name)
                               transform_description (assoc :description transform_description)
-                              source_database (assoc-in [:source :source-database] source_database)
-                              source_tables (assoc-in [:source :source-tables] source_tables))]
+                              source_database (assoc-in [:source :query :database] source_database)
+                              source_database (assoc-in [:target :database] source_database))]
 
-    ;; Store in memory if we have an ID
-    (when (and transform_id memory-atom)
-      (swap! memory-atom assoc-in [:state :transforms (str transform_id)] suggested-transform))
+    ;; Always store in memory so the LLM can reference this transform later
+    (when memory-atom
+      (swap! memory-atom assoc-in [:state :transforms (str effective-transform-id)] suggested-transform))
 
-    (log/debug "SQL transform written" {:transform-id transform_id
+    (log/debug "SQL transform written" {:transform-id effective-transform-id
+                                        :is-new? is-new-transform?
                                         :sql-length (count new-sql)})
 
     {:structured-output {:transform suggested-transform
+                         :transform-id effective-transform-id
                          :thinking thinking
-                         :message "Transform SQL updated successfully."}
+                         :message (if is-new-transform?
+                                    (str "Transform SQL created successfully. Transform ID: " effective-transform-id)
+                                    "Transform SQL updated successfully.")}
      :data-parts [(streaming/transform-suggestion-part suggested-transform)]
-     :instructions "The transform suggestion has been created and displayed to the user. Do not repeat the SQL content."}))
+     :instructions (str "The transform suggestion has been created and displayed to the user. "
+                        (when is-new-transform?
+                          (str "The transform ID is: " effective-transform-id ". Use this ID if you need to edit this transform later. "))
+                        "Do not repeat the SQL content.")}))
 
 ;;; Write Transform Python Tool
 
@@ -166,7 +198,9 @@ def transform():
 
   Parameters:
   - transform-id: Optional ID of existing transform to edit
-  - edit-action: Either {:mode \"edit\" :edits [...]} or {:mode \"replace\" :new-content \"...\"}
+  - mode: Either \"edit\" or \"replace\"
+  - edits: For edit mode, array of {:old_string :new_string :replace_all} maps
+  - new_content: For replace mode, the complete content
   - thinking: Explanation of what changes are being made
   - transform-name: Name for new transform (required if no transform-id)
   - transform-description: Description for new transform
@@ -178,26 +212,28 @@ def transform():
   Returns map with:
   - :structured-output - The suggested transform
   - :data-parts - Transform suggestion data part for streaming"
-  [{:keys [transform_id edit_action thinking transform_name transform_description
+  [{:keys [transform_id mode edits new_content thinking transform_name transform_description
            source_database source_tables memory-atom context]}]
   (log/info "Writing Python transform" {:transform-id transform_id
-                                        :edit-mode (:mode edit_action)
+                                        :edit-mode mode
                                         :has-context (some? context)})
 
   (let [;; Get current transform from context or memory, or create fresh one
-        current-transform (cond
-                            ;; Try to get from context first
-                            (and transform_id context)
-                            (get-in context [:transforms (str transform_id)])
+        ;; For new transforms, generate a temporary ID so the LLM can reference it later
+        is-new-transform? (nil? transform_id)
+        effective-transform-id (or transform_id (u/generate-nano-id))
 
-                            ;; Then try memory
-                            (and transform_id memory-atom)
-                            (get-in @memory-atom [:state :transforms (str transform_id)])
-
-                            ;; Create fresh transform
-                            :else
-                            (create-fresh-transform :python transform_name transform_description
-                                                    source_database source_tables))
+        current-transform (or
+                            ;; Try to get from context first (for transforms user is viewing)
+                           (when (and transform_id context)
+                             (get-in context [:transforms (str transform_id)]))
+                            ;; Then try memory (for transforms created in session)
+                           (when (and transform_id memory-atom)
+                             (get-in @memory-atom [:state :transforms (str transform_id)]))
+                            ;; Create fresh transform if no transform_id
+                           (when-not transform_id
+                             (create-fresh-python-transform transform_name transform_description
+                                                            source_database source_tables)))
 
         _ (when (and transform_id (nil? current-transform))
             (throw (ex-info (str "Transform with ID " transform_id " not found. "
@@ -205,37 +241,45 @@ def transform():
                             {:agent-error? true
                              :transform-id transform_id})))
 
-        current-python (get-in current-transform [:source :query] "")
+        ;; Extract current Python code from the body field
+        current-python (get-in current-transform [:source :body] "")
 
         ;; Apply edits based on mode
-        new-python (case (:mode edit_action)
-                     "edit" (apply-edits current-python (:edits edit_action))
-                     "replace" (:new_content edit_action)
+        new-python (case mode
+                     "edit" (apply-edits current-python edits)
+                     "replace" new_content
                      ;; Default to replace if no mode specified but new_content provided
-                     (if (:new_content edit_action)
-                       (:new_content edit_action)
-                       (throw (ex-info "Invalid edit_action: must specify mode 'edit' with edits or mode 'replace' with new_content"
+                     (if new_content
+                       new_content
+                       (throw (ex-info "Invalid input: must specify mode 'edit' with edits or mode 'replace' with new_content"
                                        {:agent-error? true
-                                        :edit-action edit_action}))))
+                                        :mode mode}))))
 
-        ;; Build suggested transform
+        ;; Build suggested transform - update Python code in the body field
         suggested-transform (cond-> current-transform
-                              true (assoc-in [:source :query] new-python)
-                              true (assoc-in [:source :type] "python")
+                              true (assoc-in [:source :body] new-python)
                               transform_name (assoc :name transform_name)
                               transform_description (assoc :description transform_description)
                               source_database (assoc-in [:source :source-database] source_database)
+                              source_database (assoc-in [:target :database] source_database)
                               source_tables (assoc-in [:source :source-tables] source_tables))]
 
-    ;; Store in memory if we have an ID
-    (when (and transform_id memory-atom)
-      (swap! memory-atom assoc-in [:state :transforms (str transform_id)] suggested-transform))
+    ;; Always store in memory so the LLM can reference this transform later
+    (when memory-atom
+      (swap! memory-atom assoc-in [:state :transforms (str effective-transform-id)] suggested-transform))
 
-    (log/debug "Python transform written" {:transform-id transform_id
+    (log/debug "Python transform written" {:transform-id effective-transform-id
+                                           :is-new? is-new-transform?
                                            :python-length (count new-python)})
 
     {:structured-output {:transform suggested-transform
+                         :transform-id effective-transform-id
                          :thinking thinking
-                         :message "Transform Python code updated successfully."}
+                         :message (if is-new-transform?
+                                    (str "Transform Python code created successfully. Transform ID: " effective-transform-id)
+                                    "Transform Python code updated successfully.")}
      :data-parts [(streaming/transform-suggestion-part suggested-transform)]
-     :instructions "The transform suggestion has been created and displayed to the user. Do not repeat the Python content."}))
+     :instructions (str "The transform suggestion has been created and displayed to the user. "
+                        (when is-new-transform?
+                          (str "The transform ID is: " effective-transform-id ". Use this ID if you need to edit this transform later. "))
+                        "Do not repeat the Python content.")}))

enterprise/backend/test/metabase_enterprise/metabot_v3/context_test.clj

@@ -185,6 +185,14 @@
                        :capabilities
                        (filter #(str/starts-with? % "backend:"))))))))
 
+(deftest feature-capabilities-test
+  (testing "Feature capabilities are added to context"
+    (let [capabilities (:capabilities (context/create-context {}))]
+      (testing "snippets capability is always present"
+        (is (contains? capabilities "feature:snippets")))
+      (testing "transforms capability is present when transforms are enabled"
+        (is (contains? capabilities "feature:transforms"))))))
+
 (deftest annotate-transform-source-types-native-transform-test
   (testing "Annotates draft native transform with source_type :native"
     (let [input {:user_is_viewing [{:type "transform"

enterprise/backend/test/metabase_enterprise/metabot_v3/tools/transforms_write_test.clj

@@ -9,86 +9,102 @@
   (testing "replace mode replaces entire content"
     (let [memory-atom (atom {:state {}})
           result (transforms-write/write-transform-sql
-                  {:edit_action {:mode "replace"
-                                 :new_content "SELECT * FROM users"}
+                  {:mode "replace"
+                   :new_content "SELECT * FROM users"
                    :transform_name "Test Transform"
+                   :source_database 1
                    :memory-atom memory-atom})]
       (is (= "SELECT * FROM users"
-             (get-in result [:structured-output :transform :source :query])))))
+             (get-in result [:structured-output :transform :source :query :native :query])))))
 
   (testing "edit mode with single edit"
     (let [existing-transform {:id 1
                               :name "Existing"
-                              :source {:query "SELECT id FROM orders"}}
+                              :source {:type "query"
+                                       :query {:type "native"
+                                               :database 1
+                                               :native {:query "SELECT id FROM orders"}}}}
           memory-atom (atom {:state {:transforms {"1" existing-transform}}})
           result (transforms-write/write-transform-sql
                   {:transform_id 1
-                   :edit_action {:mode "edit"
-                                 :edits [{:old_string "orders"
-                                          :new_string "customers"}]}
+                   :mode "edit"
+                   :edits [{:old_string "orders"
+                            :new_string "customers"}]
                    :memory-atom memory-atom})]
       (is (= "SELECT id FROM customers"
-             (get-in result [:structured-output :transform :source :query])))))
+             (get-in result [:structured-output :transform :source :query :native :query])))))
 
   (testing "edit mode with multiple edits"
     (let [existing-transform {:id 1
                               :name "Existing"
-                              :source {:query "SELECT col_a, col_b FROM table1"}}
+                              :source {:type "query"
+                                       :query {:type "native"
+                                               :database 1
+                                               :native {:query "SELECT col_a, col_b FROM table1"}}}}
           memory-atom (atom {:state {:transforms {"1" existing-transform}}})
           result (transforms-write/write-transform-sql
                   {:transform_id 1
-                   :edit_action {:mode "edit"
-                                 :edits [{:old_string "col_a" :new_string "col_x"}
-                                         {:old_string "col_b" :new_string "col_y"}
-                                         {:old_string "table1" :new_string "table2"}]}
+                   :mode "edit"
+                   :edits [{:old_string "col_a" :new_string "col_x"}
+                           {:old_string "col_b" :new_string "col_y"}
+                           {:old_string "table1" :new_string "table2"}]
                    :memory-atom memory-atom})]
       (is (= "SELECT col_x, col_y FROM table2"
-             (get-in result [:structured-output :transform :source :query])))))
+             (get-in result [:structured-output :transform :source :query :native :query])))))
 
   (testing "edit mode fails when text not found"
     (let [existing-transform {:id 1
                               :name "Existing"
-                              :source {:query "SELECT * FROM users"}}
+                              :source {:type "query"
+                                       :query {:type "native"
+                                               :database 1
+                                               :native {:query "SELECT * FROM users"}}}}
           memory-atom (atom {:state {:transforms {"1" existing-transform}}})]
       (is (thrown-with-msg?
            clojure.lang.ExceptionInfo
            #"Could not find text"
            (transforms-write/write-transform-sql
             {:transform_id 1
-             :edit_action {:mode "edit"
-                           :edits [{:old_string "nonexistent"
-                                    :new_string "replacement"}]}
+             :mode "edit"
+             :edits [{:old_string "nonexistent"
+                      :new_string "replacement"}]
              :memory-atom memory-atom})))))
 
   (testing "edit mode fails for ambiguous matches without replace_all"
     (let [existing-transform {:id 1
                               :name "Existing"
-                              :source {:query "SELECT foo, foo FROM mytable"}}
+                              :source {:type "query"
+                                       :query {:type "native"
+                                               :database 1
+                                               :native {:query "SELECT foo, foo FROM mytable"}}}}
           memory-atom (atom {:state {:transforms {"1" existing-transform}}})]
       (is (thrown-with-msg?
            clojure.lang.ExceptionInfo
            #"Found 2 matches"
            (transforms-write/write-transform-sql
             {:transform_id 1
-             :edit_action {:mode "edit"
-                           :edits [{:old_string "foo"
-                                    :new_string "bar"}]}
+             :mode "edit"
+             :edits [{:old_string "foo"
+                      :new_string "bar"}]
              :memory-atom memory-atom})))))
 
   (testing "edit mode with replace_all replaces all occurrences"
     (let [existing-transform {:id 1
                               :name "Existing"
-                              :source {:query "SELECT xyz, xyz, xyz FROM mytable"}}
+                              :source {:type "query"
+                                       :query {:type "native"
+                                               :database 1
+                                               :native {:query "SELECT xyz, xyz, xyz FROM mytable"}}}}
           memory-atom (atom {:state {:transforms {"1" existing-transform}}})
           result (transforms-write/write-transform-sql
                   {:transform_id 1
-                   :edit_action {:mode "edit"
-                                 :edits [{:old_string "xyz"
-                                          :new_string "abc"
-                                          :replace_all true}]}
+                   :mode "edit"
+                   :edits [{:old_string "xyz"
+                            :new_string "abc"
+                            :replace_all true}]
                    :memory-atom memory-atom})]
       (is (= "SELECT abc, abc, abc FROM mytable"
-             (get-in result [:structured-output :transform :source :query]))))))
+             (get-in result [:structured-output :transform :source :query :native :query]))))))
 
 ;;; Transform Creation Tests
 
@@ -96,31 +112,34 @@
   (testing "creates fresh SQL transform when no transform_id"
     (let [memory-atom (atom {:state {}})
           result (transforms-write/write-transform-sql
-                  {:edit_action {:mode "replace"
-                                 :new_content "SELECT 1"}
+                  {:mode "replace"
+                   :new_content "SELECT 1"
                    :transform_name "New Transform"
                    :transform_description "A test transform"
+                   :source_database 1
                    :memory-atom memory-atom})]
       (is (= "New Transform" (get-in result [:structured-output :transform :name])))
       (is (= "A test transform" (get-in result [:structured-output :transform :description])))
-      (is (= "SELECT 1" (get-in result [:structured-output :transform :source :query])))))
+      (is (= "query" (get-in result [:structured-output :transform :source :type])))
+      (is (= "SELECT 1" (get-in result [:structured-output :transform :source :query :native :query])))))
 
   (testing "creates fresh Python transform"
     (let [memory-atom (atom {:state {}})
           result (transforms-write/write-transform-python
-                  {:edit_action {:mode "replace"
-                                 :new_content "def transform():\n    return pd.DataFrame()"}
+                  {:mode "replace"
+                   :new_content "def transform():\n    return pd.DataFrame()"
                    :transform_name "Python Transform"
+                   :source_database 1
                    :memory-atom memory-atom})]
       (is (= "Python Transform" (get-in result [:structured-output :transform :name])))
       (is (= "python" (get-in result [:structured-output :transform :source :type]))))))
 
 (deftest create-fresh-python-template-test
   (testing "fresh Python transforms include common import"
-    (let [result (#'transforms-write/create-fresh-transform
-                  :python "Python Transform" nil nil nil)]
+    (let [result (#'transforms-write/create-fresh-python-transform
+                  "Python Transform" nil 1 nil)]
       (is (= "import common\nimport pandas as pd\n\ndef transform():\n    # Your transformation logic here\n    return pd.DataFrame([{\"message\": \"Hello from Python transform!\"}])\n"
-             (get-in result [:source :query]))))))
+             (get-in result [:source :body]))))))
 
 ;;; Data Parts Tests
 
@@ -128,8 +147,10 @@
   (testing "returns transform_suggestion data part"
     (let [memory-atom (atom {:state {}})
           result (transforms-write/write-transform-sql
-                  {:edit_action {:mode "replace" :new_content "SELECT 1"}
+                  {:mode "replace"
+                   :new_content "SELECT 1"
                    :transform_name "Test"
+                   :source_database 1
                    :memory-atom memory-atom})]
       (is (contains? result :data-parts))
       (is (= 1 (count (:data-parts result))))
@@ -142,21 +163,33 @@
 
 (deftest memory-storage-test
   (testing "stores updated transform in memory when transform_id provided"
-    (let [existing-transform {:id 1 :name "Existing" :source {:query "SELECT 1"}}
+    (let [existing-transform {:id 1
+                              :name "Existing"
+                              :source {:type "query"
+                                       :query {:type "native"
+                                               :database 1
+                                               :native {:query "SELECT 1"}}}}
           memory-atom (atom {:state {:transforms {"1" existing-transform}}})
           _ (transforms-write/write-transform-sql
              {:transform_id 1
-              :edit_action {:mode "replace" :new_content "SELECT 2"}
+              :mode "replace"
+              :new_content "SELECT 2"
               :memory-atom memory-atom})]
-      (is (= "SELECT 2" (get-in @memory-atom [:state :transforms "1" :source :query])))))
+      (is (= "SELECT 2" (get-in @memory-atom [:state :transforms "1" :source :query :native :query])))))
 
-  (testing "does not store in memory when no transform_id"
+  (testing "stores new transform with generated ID when no transform_id"
     (let [memory-atom (atom {:state {:transforms {}}})
-          _ (transforms-write/write-transform-sql
-             {:edit_action {:mode "replace" :new_content "SELECT 1"}
-              :transform_name "New"
-              :memory-atom memory-atom})]
-      (is (empty? (get-in @memory-atom [:state :transforms]))))))
+          result (transforms-write/write-transform-sql
+                  {:mode "replace"
+                   :new_content "SELECT 1"
+                   :transform_name "New"
+                   :source_database 1
+                   :memory-atom memory-atom})
+          generated-id (get-in result [:structured-output :transform-id])]
+      (is (some? generated-id) "Should generate a transform ID")
+      (is (= 1 (count (get-in @memory-atom [:state :transforms]))) "Should store one transform")
+      (is (= "SELECT 1" (get-in @memory-atom [:state :transforms generated-id :source :query :native :query]))
+          "Should store the transform with the generated ID"))))
 
 ;;; Error Handling Tests
 
@@ -168,14 +201,16 @@
            #"Transform with ID 999 not found"
            (transforms-write/write-transform-sql
             {:transform_id 999
-             :edit_action {:mode "replace" :new_content "SELECT 1"}
+             :mode "replace"
+             :new_content "SELECT 1"
              :memory-atom memory-atom})))))
 
-  (testing "fails when edit_action invalid"
+  (testing "fails when mode is invalid"
     (let [memory-atom (atom {:state {}})]
       (is (thrown-with-msg?
            clojure.lang.ExceptionInfo
-           #"Invalid edit_action"
+           #"Invalid input"
            (transforms-write/write-transform-sql
-            {:edit_action {:mode "invalid"}
+            {:mode "invalid"
+             :source_database 1
              :memory-atom memory-atom}))))))