From a531f41ee52f6a3f9bff1641df07628e1316c19a Mon Sep 17 00:00:00 2001
From: Bram <bram2w@users.noreply.github.com>
Date: Thu, 23 Oct 2025 11:15:15 +0200
Subject: [PATCH] Creating formula for AI field with choice type fails hard
 (#4079)

Allow formula field to reference AI field with choice output type
---
 ...reference_ai_field_with_choice_output.json |  8 ++
 .../src/baserow_premium/fields/field_types.py | 11 +++
 .../fields/test_ai_field_type.py              | 73 +++++++++++++++++++
 3 files changed, 92 insertions(+)
 create mode 100644 changelog/entries/unreleased/bug/3851_allow_formula_field_to_reference_ai_field_with_choice_output.json

diff --git a/changelog/entries/unreleased/bug/3851_allow_formula_field_to_reference_ai_field_with_choice_output.json b/changelog/entries/unreleased/bug/3851_allow_formula_field_to_reference_ai_field_with_choice_output.json
new file mode 100644
index 0000000000..d173499aee
--- /dev/null
+++ b/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"
+}
\ No newline at end of file
diff --git a/premium/backend/src/baserow_premium/fields/field_types.py b/premium/backend/src/baserow_premium/fields/field_types.py
index f0d189903c..c960a09967 100644
--- a/premium/backend/src/baserow_premium/fields/field_types.py
+++ b/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,
diff --git a/premium/backend/tests/baserow_premium_tests/fields/test_ai_field_type.py b/premium/backend/tests/baserow_premium_tests/fields/test_ai_field_type.py
index 6390187f71..2433272901 100644
--- a/premium/backend/tests/baserow_premium_tests/fields/test_ai_field_type.py
+++ b/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"