test_ai.py

# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest

from bigframes.bigquery import ai, ml
import bigframes.pandas as bpd


@pytest.fixture(scope="session")
def embedding_model(bq_connection, dataset_id):
    model_name = f"{dataset_id}.embedding_model"
    return ml.create_model(
        model_name=model_name,
        options={"endpoint": "gemini-embedding-001"},
        connection_name=bq_connection,
    )


@pytest.fixture(scope="session")
def text_model(bq_connection, dataset_id):
    model_name = f"{dataset_id}.text_model"
    return ml.create_model(
        model_name=model_name,
        options={"endpoint": "gemini-2.5-flash"},
        connection_name=bq_connection,
    )


def test_generate_embedding(embedding_model):
    df = bpd.DataFrame(
        {
            "content": [
                "What is BigQuery?",
                "What is BQML?",
            ]
        }
    )

    result = ai.generate_embedding(embedding_model, df)

    assert len(result) == 2
    assert "embedding" in result.columns
    assert "statistics" in result.columns
    assert "status" in result.columns


def test_generate_embedding_with_options(embedding_model):
    df = bpd.DataFrame(
        {
            "content": [
                "What is BigQuery?",
                "What is BQML?",
            ]
        }
    )

    result = ai.generate_embedding(
        embedding_model, df, task_type="RETRIEVAL_DOCUMENT", output_dimensionality=256
    )

    assert len(result) == 2
    embedding = result["embedding"].to_pandas()
    assert len(embedding[0]) == 256


def test_generate_text(text_model):
    df = bpd.DataFrame({"prompt": ["Dog", "Cat"]})

    result = ai.generate_text(text_model, df)

    assert len(result) == 2
    assert "result" in result.columns
    assert "statistics" in result.columns
    assert "full_response" in result.columns
    assert "status" in result.columns


def test_generate_text_with_options(text_model):
    df = bpd.DataFrame({"prompt": ["Dog", "Cat"]})

    result = ai.generate_text(text_model, df, max_output_tokens=1)

    # It basically asserts that the results are still returned.
    assert len(result) == 2


def test_generate_table(text_model):
    df = bpd.DataFrame(
        {"prompt": ["Generate a table of 2 programming languages and their creators."]}
    )

    result = ai.generate_table(
        text_model,
        df,
        output_schema="language STRING, creator STRING",
    )

    assert "language" in result.columns
    assert "creator" in result.columns
    # The model may not always return the exact number of rows requested.
    assert len(result) > 0