diff --git a/server/app/controllers/api/v1/connectors_controller.rb b/server/app/controllers/api/v1/connectors_controller.rb
index 821a99f7..5d8e9e67 100644
--- a/server/app/controllers/api/v1/connectors_controller.rb
+++ b/server/app/controllers/api/v1/connectors_controller.rb
@@ -9,7 +9,7 @@ class ConnectorsController < ApplicationController
       # TODO: Enable this once we have query validation implemented for all the connectors
       # before_action :validate_query, only: %i[query_source]
       # TODO: Enable this for ai_ml sources
-      # before_action :validate_catalog, only: %i[query_source]
+      before_action :validate_catalog, only: %i[query_source]
       after_action :event_logger
 
       def index
@@ -126,6 +126,7 @@ def set_connector
       end
 
       def validate_catalog
+        return unless @connector.ai_model?
         return if @connector.catalog.present?
 
         render_error(
diff --git a/server/app/models/connector.rb b/server/app/models/connector.rb
index c2ee3c36..d48bc955 100644
--- a/server/app/models/connector.rb
+++ b/server/app/models/connector.rb
@@ -133,4 +133,8 @@ def set_category
   rescue StandardError => e
     Rails.logger.error("Failed to set category for connector ##{id}: #{e.message}")
   end
+
+  def ai_model?
+    connector_category == "AI Model"
+  end
 end
diff --git a/server/spec/models/connector_spec.rb b/server/spec/models/connector_spec.rb
index fd814aaa..8a1c1981 100644
--- a/server/spec/models/connector_spec.rb
+++ b/server/spec/models/connector_spec.rb
@@ -196,6 +196,11 @@
       expect(result).to include(ai_ml_connector)
       expect(result).not_to include(non_ai_ml_connector)
     end
+
+    it "check whether connector is ai model or not" do
+      expect(ai_ml_connector.ai_model?).to eq(true)
+      expect(non_ai_ml_connector.ai_model?).to eq(false)
+    end
   end
 
   describe ".data" do
diff --git a/server/spec/requests/api/v1/connectors_controller_spec.rb b/server/spec/requests/api/v1/connectors_controller_spec.rb
index 02f9da27..693988f0 100644
--- a/server/spec/requests/api/v1/connectors_controller_spec.rb
+++ b/server/spec/requests/api/v1/connectors_controller_spec.rb
@@ -481,19 +481,39 @@
         expect(response_hash[:data]).to eq([record1.record.data, record2.record.data])
       end
 
-      # it "returns an error message for missing catalog" do
-      #   catalog = connector.catalog
-      #   catalog.connector_id = connectors.second.id
-      #   catalog.save
-
-      #   allow(Connectors::QuerySource).to receive(:call).and_return(double(:context, success?: true,
-      #                                                                                records: [record1, record2]))
-      #   post "/api/v1/connectors/#{connector.id}/query_source", params: request_body.to_json, headers:
-      #     { "Content-Type": "application/json" }.merge(auth_headers(user, workspace.id))
-      #   expect(response).to have_http_status(:unprocessable_entity)
-      #   response_hash = JSON.parse(response.body).with_indifferent_access
-      #   expect(response_hash.dig(:errors, 0, :detail)).to eq("Catalog is not present for the connector")
-      # end
+      it "returns an error message for missing catalog for ai connectors" do
+        catalog = connector.catalog
+        catalog.connector_id = connectors.second.id
+        catalog.save
+        # rubocop:disable Rails/SkipsModelValidations
+        connector.update_column(:connector_category, "AI Model")
+        # rubocop:enable Rails/SkipsModelValidations
+
+        allow(Connectors::QuerySource).to receive(:call).and_return(double(:context, success?: true,
+                                                                                     records: [record1, record2]))
+        post "/api/v1/connectors/#{connector.id}/query_source", params: request_body.to_json, headers:
+          { "Content-Type": "application/json" }.merge(auth_headers(user, workspace.id))
+        expect(response).to have_http_status(:unprocessable_entity)
+        response_hash = JSON.parse(response.body).with_indifferent_access
+        expect(response_hash.dig(:errors, 0, :detail)).to eq("Catalog is not present for the connector")
+      end
+
+      it "should not return error message for missing catalog for data connectors" do
+        catalog = connector.catalog
+        catalog.connector_id = connectors.second.id
+        # rubocop:disable Rails/SkipsModelValidations
+        connector.update_column(:connector_category, "Data Warehouse")
+        # rubocop:enable Rails/SkipsModelValidations
+        catalog.save
+
+        allow(Connectors::QuerySource).to receive(:call).and_return(double(:context, success?: true,
+                                                                                     records: [record1, record2]))
+        post "/api/v1/connectors/#{connector.id}/query_source", params: request_body.to_json, headers:
+          { "Content-Type": "application/json" }.merge(auth_headers(user, workspace.id))
+        expect(response).to have_http_status(:ok)
+        response_hash = JSON.parse(response.body).with_indifferent_access
+        expect(response_hash[:data]).to eq([record1.record.data, record2.record.data])
+      end
 
       it "returns success status for a valid query for viewer role" do
         workspace.workspace_users.first.update(role: viewer_role)