fix lint and add imports

Author: tswast

bigframes/bigquery/_operations/ai.py

@@ -469,7 +469,7 @@ def generate_embedding(
     # We need to get the SQL for the input data to pass as a subquery to the TVF
     source_sql = data_df.sql
 
-    struct_fields = {}
+    struct_fields: bigframes.core.sql.literals.STRUCT_TYPE = {}
     if output_dimensionality is not None:
         struct_fields["OUTPUT_DIMENSIONALITY"] = output_dimensionality
     if task_type is not None:

bigframes/bigquery/ai.py

@@ -22,6 +22,7 @@
     generate,
     generate_bool,
     generate_double,
+    generate_embedding,
     generate_int,
     if_,
     score,
@@ -33,6 +34,7 @@
     "generate",
     "generate_bool",
     "generate_double",
+    "generate_embedding",
     "generate_int",
     "if_",
     "score",

bigframes/core/sql/literals.py

@@ -16,18 +16,17 @@
 
 import collections.abc
 import json
-from typing import Any, Dict, List, Mapping, Optional, Union
+from typing import Any, List, Mapping, Union
 
-import bigframes.core.compile.googlesql as googlesql
 import bigframes.core.sql
 
+STRUCT_TYPE = Mapping[
+    str,
+    Union[str, int, float, bool, Mapping[str, str], List[str], Mapping[str, Any]],
+]
 
-def struct_literal(
-    struct_options: Mapping[
-        str,
-        Union[str, int, float, bool, Mapping[str, str], List[str], Mapping[str, Any]],
-    ]
-) -> str:
+
+def struct_literal(struct_options: STRUCT_TYPE) -> str:
     rendered_options = []
     for option_name, option_value in struct_options.items():
         if option_name == "model_params":

bigframes/core/sql/ml.py

@@ -14,8 +14,6 @@
 
 from __future__ import annotations
 
-import collections.abc
-import json
 from typing import Any, Dict, List, Mapping, Optional, Union
 
 import bigframes.core.compile.googlesql as googlesql

tests/unit/bigquery/test_ai.py

@@ -17,7 +17,7 @@
 import pandas as pd
 import pytest
 
-import bigframes.bigquery._operations.ai as ai_ops
+import bigframes.bigquery as bbq
 import bigframes.dataframe
 import bigframes.series
 import bigframes.session
@@ -52,7 +52,7 @@ def mock_series(mock_session):
 def test_generate_embedding_with_dataframe(mock_dataframe, mock_session):
     model_name = "project.dataset.model"
 
-    ai_ops.generate_embedding(
+    bbq.ai.generate_embedding(
         model_name,
         mock_dataframe,
         output_dimensionality=256,
@@ -78,7 +78,7 @@ def test_generate_embedding_with_dataframe(mock_dataframe, mock_session):
 def test_generate_embedding_with_series(mock_series, mock_session):
     model_name = "project.dataset.model"
 
-    ai_ops.generate_embedding(
+    bbq.ai.generate_embedding(
         model_name, mock_series, start_second=0.0, end_second=10.0, interval_seconds=5.0
     )
 
@@ -97,7 +97,7 @@ def test_generate_embedding_with_series(mock_series, mock_session):
 def test_generate_embedding_defaults(mock_dataframe, mock_session):
     model_name = "project.dataset.model"
 
-    ai_ops.generate_embedding(
+    bbq.ai.generate_embedding(
         model_name,
         mock_dataframe,
     )
@@ -122,7 +122,7 @@ def test_generate_embedding_with_pandas_dataframe(
 
     pandas_df = pd.DataFrame({"content": ["test"]})
 
-    ai_ops.generate_embedding(
+    bbq.ai.generate_embedding(
         model_name,
         pandas_df,
     )