Skip to content

Commit 3e7dc97

Browse files
shuoweilshuoweil
committed
revert: restore ai_ops.py and AI snapshots to main
1 parent b25a3e8 commit 3e7dc97

23 files changed

Lines changed: 49 additions & 56 deletions

File tree

  • packages/bigframes
    • bigframes/core/compile/sqlglot/expressions
    • tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops
      • test_ai_classify_with_params
      • test_ai_classify
      • test_ai_generate_bool_with_connection_id
      • test_ai_generate_bool_with_model_param
      • test_ai_generate_bool
      • test_ai_generate_double_with_connection_id
      • test_ai_generate_double_with_model_param
      • test_ai_generate_double
      • test_ai_generate_int_with_connection_id
      • test_ai_generate_int_with_model_param
      • test_ai_generate_int
      • test_ai_generate_with_connection_id
      • test_ai_generate_with_model_param
      • test_ai_generate_with_output_schema
      • test_ai_generate
      • test_ai_if_with_endpoint
      • test_ai_if
      • test_ai_score_with_endpoint_and_max_error_ratio
      • test_ai_score

‎packages/bigframes/bigframes/core/compile/sqlglot/expressions/ai_ops.py‎

Lines changed: 5 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -101,22 +101,17 @@ def _construct_prompt(
101101
prompt_context: tuple[str | None, ...],
102102
param_name: str = "prompt",
103103
) -> sge.Kwarg:
104-
expressions = []
104+
prompt: list[str | sge.Expression] = []
105105
column_ref_idx = 0
106106

107-
for idx, elem in enumerate(prompt_context):
108-
field_name = f"_field_{idx + 1}"
107+
for elem in prompt_context:
109108
if elem is None:
110-
value = exprs[column_ref_idx].expr
109+
prompt.append(exprs[column_ref_idx].expr)
111110
column_ref_idx += 1
112111
else:
113-
value = sge.Literal.string(elem)
112+
prompt.append(sge.Literal.string(elem))
114113

115-
expressions.append(
116-
sge.PropertyEQ(this=sge.to_identifier(field_name), expression=value)
117-
)
118-
119-
return sge.Kwarg(this=param_name, expression=sge.Struct(expressions=expressions))
114+
return sge.Kwarg(this=param_name, expression=sge.Tuple(expressions=prompt))
120115

121116

122117
def _construct_named_args(op: ops.ScalarOp) -> list[sge.Kwarg]:

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/None/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,3 @@
11
SELECT
2-
AI.CLASSIFY(input => STRUCT(`string_col` AS _field_1), categories => ['greeting', 'rejection']) AS `result`
3-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
2+
AI.CLASSIFY(input => (`string_col`), categories => ['greeting', 'rejection']) AS `result`
3+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify/bigframes-dev.us.bigframes-default-connection/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
SELECT
22
AI.CLASSIFY(
3-
input => STRUCT(`string_col` AS _field_1),
3+
input => (`string_col`),
44
categories => ['greeting', 'rejection'],
55
connection_id => 'bigframes-dev.us.bigframes-default-connection'
66
) AS `result`
7-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
7+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_classify_with_params/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
SELECT
22
AI.CLASSIFY(
3-
input => STRUCT(`string_col` AS _field_1),
3+
input => (`string_col`),
44
categories => ['greeting', 'rejection'],
55
examples => [('hi', 'greeting'), ('bye', 'rejection')],
66
endpoint => 'gemini-2.5-flash',
77
max_error_ratio => 0.1
88
) AS `result`
9-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
9+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
SELECT
22
AI.GENERATE(
3-
prompt => STRUCT(`string_col` AS _field_1, ' is the same as ' AS _field_2, `string_col` AS _field_3),
3+
prompt => (`string_col`, ' is the same as ', `string_col`),
44
endpoint => 'gemini-2.5-flash',
55
request_type => 'SHARED'
66
) AS `result`
7-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
7+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
SELECT
22
AI.GENERATE_BOOL(
3-
prompt => STRUCT(`string_col` AS _field_1, ' is the same as ' AS _field_2, `string_col` AS _field_3),
3+
prompt => (`string_col`, ' is the same as ', `string_col`),
44
endpoint => 'gemini-2.5-flash'
55
) AS `result`
6-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_connection_id/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
SELECT
22
AI.GENERATE_BOOL(
3-
prompt => STRUCT(`string_col` AS _field_1, ' is the same as ' AS _field_2, `string_col` AS _field_3),
3+
prompt => (`string_col`, ' is the same as ', `string_col`),
44
connection_id => 'bigframes-dev.us.bigframes-default-connection',
55
endpoint => 'gemini-2.5-flash'
66
) AS `result`
7-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
7+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_bool_with_model_param/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
SELECT
22
AI.GENERATE_BOOL(
3-
prompt => STRUCT(`string_col` AS _field_1, ' is the same as ' AS _field_2, `string_col` AS _field_3),
3+
prompt => (`string_col`, ' is the same as ', `string_col`),
44
model_params => JSON '{}'
55
) AS `result`
6-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
SELECT
22
AI.GENERATE_DOUBLE(
3-
prompt => STRUCT(`string_col` AS _field_1, ' is the same as ' AS _field_2, `string_col` AS _field_3),
3+
prompt => (`string_col`, ' is the same as ', `string_col`),
44
endpoint => 'gemini-2.5-flash'
55
) AS `result`
6-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
6+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

‎packages/bigframes/tests/unit/core/compile/sqlglot/expressions/snapshots/test_ai_ops/test_ai_generate_double_with_connection_id/out.sql‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
SELECT
22
AI.GENERATE_DOUBLE(
3-
prompt => STRUCT(`string_col` AS _field_1, ' is the same as ' AS _field_2, `string_col` AS _field_3),
3+
prompt => (`string_col`, ' is the same as ', `string_col`),
44
connection_id => 'bigframes-dev.us.bigframes-default-connection',
55
endpoint => 'gemini-2.5-flash'
66
) AS `result`
7-
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`
7+
FROM `bigframes-dev`.`sqlglot_test`.`scalar_types` AS `bft_0`

0 commit comments

Comments
 (0)