Creating formula for AI field with choice type fails hard (baserow#4079)

Allow formula field to reference AI field with choice output type

Author: bram2w

changelog/entries/unreleased/bug/3851_allow_formula_field_to_reference_ai_field_with_choice_output.json

@@ -0,0 +1,8 @@
+{
+    "type": "bug",
+    "message": "Allow formula field to reference AI field with choice output type",
+    "domain": "database",
+    "issue_number": 3851,
+    "bullet_points": [],
+    "created_at": "2025-10-08"
+}

premium/backend/src/baserow_premium/fields/field_types.py

@@ -249,6 +249,17 @@ def get_group_by_field_filters_and_annotations(
             field, field_name, base_queryset, value, cte, rows
         )
 
+    def get_formula_reference_to_model_field(
+        self, model_field, db_column, already_in_subquery
+    ):
+        instance = model_field.model.get_field_object(
+            model_field.name, include_trash=True
+        )["field"]
+        baserow_field_type = self.get_baserow_field_type(instance)
+        return baserow_field_type.get_formula_reference_to_model_field(
+            model_field, db_column, already_in_subquery
+        )
+
     def get_export_serialized_value(
         self,
         row,

premium/backend/tests/baserow_premium_tests/fields/test_ai_field_type.py

@@ -1120,3 +1120,76 @@ def test_link_row_field_can_be_sorted_when_linking_an_ai_field(premium_data_fixt
             .values_list("id", flat=True)
         )
     assert result == [row_b2.id, row_b1.id]
+
+
+@pytest.mark.django_db
+@pytest.mark.field_ai
+def test_formula_field_can_reference_ai_choice_output_without_error(
+    premium_data_fixture,
+):
+    premium_data_fixture.register_fake_generate_ai_type()
+    user = premium_data_fixture.create_user()
+    table = premium_data_fixture.create_database_table(user=user)
+
+    ai_choice_field = premium_data_fixture.create_ai_field(
+        table=table,
+        order=0,
+        name="ai_choice",
+        ai_output_type="choice",
+        ai_generative_ai_type="test_generative_ai",
+        ai_generative_ai_model="test_1",
+        ai_prompt="'pick one'",
+    )
+    premium_data_fixture.create_select_option(
+        field=ai_choice_field, value="Red", color="red", order=0
+    )
+    premium_data_fixture.create_select_option(
+        field=ai_choice_field, value="Blue", color="blue", order=1
+    )
+
+    formula_field = FieldHandler().create_field(
+        user=user,
+        table=table,
+        type_name="formula",
+        name="formula_from_ai_choice",
+        formula="field('ai_choice')",
+    )
+
+    assert formula_field is not None
+
+
+@pytest.mark.django_db
+@pytest.mark.field_ai
+def test_ai_field_can_be_used_in_lookup_expression(premium_data_fixture):
+    premium_data_fixture.register_fake_generate_ai_type()
+    user = premium_data_fixture.create_user()
+
+    source_table = premium_data_fixture.create_database_table(user=user, name="Source")
+    ai_field = premium_data_fixture.create_ai_field(
+        table=source_table,
+        order=0,
+        name="ai_text",
+        ai_output_type="text",
+        ai_generative_ai_type="test_generative_ai",
+        ai_generative_ai_model="test_1",
+        ai_prompt="'Hello World'",
+    )
+
+    target_table = premium_data_fixture.create_database_table(user=user, name="Target")
+    link_row_field = premium_data_fixture.create_link_row_field(
+        table=target_table,
+        order=0,
+        name="link_to_source",
+        link_row_table=source_table,
+    )
+
+    formula_field = FieldHandler().create_field(
+        user=user,
+        table=target_table,
+        type_name="formula",
+        name="lookup_ai_field",
+        formula="lookup('link_to_source', 'ai_text')",
+    )
+
+    assert formula_field is not None
+    assert formula_field.formula_type == "array"