diff --git a/bigframes/ml/llm.py b/bigframes/ml/llm.py index 5c5153e163..69f2f5bc62 100644 --- a/bigframes/ml/llm.py +++ b/bigframes/ml/llm.py @@ -913,6 +913,7 @@ def predict( max_output_tokens: int = 8192, top_k: int = 40, top_p: float = 1.0, + ground_with_google_search: bool = False, ) -> bpd.DataFrame: """Predict the result from input DataFrame. @@ -936,11 +937,20 @@ def predict( Specify a lower value for less random responses and a higher value for more random responses. Default 40. Possible values [1, 40]. - top_p (float, default 0.95):: + top_p (float, default 0.95): Top-P changes how the model selects tokens for output. Tokens are selected from the most (see top-K) to least probable until the sum of their probabilities equals the top-P value. For example, if tokens A, B, and C have a probability of 0.3, 0.2, and 0.1 and the top-P value is 0.5, then the model will select either A or B as the next token by using temperature and excludes C as a candidate. Specify a lower value for less random responses and a higher value for more random responses. Default 1.0. Possible values [0.0, 1.0]. + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the Vertex AI model. When set + to True, the model incorporates relevant information from Google Search + results into its responses, enhancing their accuracy and factualness. + This feature provides an additional column, `ml_generate_text_grounding_result`, + in the response output, detailing the sources used for grounding. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. Returns: bigframes.dataframe.DataFrame: DataFrame of shape (n_samples, n_input_columns + n_prediction_columns). Returns predicted values. @@ -974,6 +984,7 @@ def predict( "top_k": top_k, "top_p": top_p, "flatten_json_output": True, + "ground_with_google_search": ground_with_google_search, } df = self._bqml_model.generate_text(X, options) diff --git a/bigframes/operations/semantics.py b/bigframes/operations/semantics.py index 3e0be74889..46ed5b8b3e 100644 --- a/bigframes/operations/semantics.py +++ b/bigframes/operations/semantics.py @@ -16,6 +16,7 @@ import re import typing from typing import List, Optional +import warnings import numpy as np @@ -39,6 +40,7 @@ def agg( model, cluster_column: typing.Optional[str] = None, max_agg_rows: int = 10, + ground_with_google_search: bool = False, ): """ Performs an aggregation over all rows of the table. @@ -90,6 +92,14 @@ def agg( max_agg_rows (int, default 10): The maxinum number of rows to be aggregated at a time. + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + Returns: bigframes.dataframe.DataFrame: A new DataFrame with the aggregated answers. @@ -119,6 +129,12 @@ def agg( ) column = columns[0] + if ground_with_google_search: + warnings.warn( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + if max_agg_rows <= 1: raise ValueError( f"Invalid value for `max_agg_rows`: {max_agg_rows}." @@ -191,7 +207,12 @@ def agg( # Run model predict_df = typing.cast( - bigframes.dataframe.DataFrame, model.predict(prompt_s, temperature=0.0) + bigframes.dataframe.DataFrame, + model.predict( + prompt_s, + temperature=0.0, + ground_with_google_search=ground_with_google_search, + ), ) agg_df[column] = predict_df["ml_generate_text_llm_result"].combine_first( single_row_df @@ -284,7 +305,7 @@ def cluster_by( df[output_column] = clustered_result["CENTROID_ID"] return df - def filter(self, instruction: str, model): + def filter(self, instruction: str, model, ground_with_google_search: bool = False): """ Filters the DataFrame with the semantics of the user instruction. @@ -305,18 +326,26 @@ def filter(self, instruction: str, model): [1 rows x 2 columns] Args: - instruction: + instruction (str): An instruction on how to filter the data. This value must contain column references by name, which should be wrapped in a pair of braces. For example, if you have a column "food", you can refer to this column in the instructions like: "The {food} is healthy." - model: + model (bigframes.ml.llm.GeminiTextGenerator): A GeminiTextGenerator provided by Bigframes ML package. + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + Returns: - DataFrame filtered by the instruction. + bigframes.pandas.DataFrame: DataFrame filtered by the instruction. Raises: NotImplementedError: when the semantic operator experiment is off. @@ -332,6 +361,12 @@ def filter(self, instruction: str, model): if column not in self._df.columns: raise ValueError(f"Column {column} not found.") + if ground_with_google_search: + warnings.warn( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + df: bigframes.dataframe.DataFrame = self._df[columns].copy() for column in columns: if df[column].dtype != dtypes.STRING_DTYPE: @@ -345,6 +380,7 @@ def filter(self, instruction: str, model): model.predict( self._make_prompt(df, columns, user_instruction, output_instruction), temperature=0.0, + ground_with_google_search=ground_with_google_search, ), ) @@ -352,7 +388,13 @@ def filter(self, instruction: str, model): results["ml_generate_text_llm_result"].str.lower().str.contains("true") ] - def map(self, instruction: str, output_column: str, model): + def map( + self, + instruction: str, + output_column: str, + model, + ground_with_google_search: bool = False, + ): """ Maps the DataFrame with the semantics of the user instruction. @@ -376,21 +418,29 @@ def map(self, instruction: str, output_column: str, model): [2 rows x 3 columns] Args: - instruction: + instruction (str): An instruction on how to map the data. This value must contain column references by name, which should be wrapped in a pair of braces. For example, if you have a column "food", you can refer to this column in the instructions like: "Get the ingredients of {food}." - output_column: + output_column (str): The column name of the mapping result. - model: + model (bigframes.ml.llm.GeminiTextGenerator): A GeminiTextGenerator provided by Bigframes ML package. + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + Returns: - DataFrame with attached mapping results. + bigframes.pandas.DataFrame: DataFrame with attached mapping results. Raises: NotImplementedError: when the semantic operator experiment is off. @@ -406,6 +456,12 @@ def map(self, instruction: str, output_column: str, model): if column not in self._df.columns: raise ValueError(f"Column {column} not found.") + if ground_with_google_search: + warnings.warn( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + df: bigframes.dataframe.DataFrame = self._df[columns].copy() for column in columns: if df[column].dtype != dtypes.STRING_DTYPE: @@ -421,6 +477,7 @@ def map(self, instruction: str, output_column: str, model): model.predict( self._make_prompt(df, columns, user_instruction, output_instruction), temperature=0.0, + ground_with_google_search=ground_with_google_search, )["ml_generate_text_llm_result"], ) @@ -428,7 +485,14 @@ def map(self, instruction: str, output_column: str, model): return concat([self._df, results.rename(output_column)], axis=1) - def join(self, other, instruction: str, model, max_rows: int = 1000): + def join( + self, + other, + instruction: str, + model, + max_rows: int = 1000, + ground_with_google_search: bool = False, + ): """ Joines two dataframes by applying the instruction over each pair of rows from the left and right table. @@ -455,10 +519,10 @@ def join(self, other, instruction: str, model, max_rows: int = 1000): [4 rows x 2 columns] Args: - other: + other (bigframes.pandas.DataFrame): The other dataframe. - instruction: + instruction (str): An instruction on how left and right rows can be joined. This value must contain column references by name. which should be wrapped in a pair of braces. For example: "The {city} belongs to the {country}". @@ -467,15 +531,23 @@ def join(self, other, instruction: str, model, max_rows: int = 1000): self joins. For example: "The {left.employee_name} reports to {right.employee_name}" For unique column names, this prefix is optional. - model: + model (bigframes.ml.llm.GeminiTextGenerator): A GeminiTextGenerator provided by Bigframes ML package. - max_rows: + max_rows (int, default 1000): The maximum number of rows allowed to be sent to the model per call. If the result is too large, the method call will end early with an error. + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + Returns: - The joined dataframe. + bigframes.pandas.DataFrame: The joined dataframe. Raises: ValueError if the amount of data that will be sent for LLM processing is larger than max_rows. @@ -483,6 +555,12 @@ def join(self, other, instruction: str, model, max_rows: int = 1000): self._validate_model(model) columns = self._parse_columns(instruction) + if ground_with_google_search: + warnings.warn( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + joined_table_rows = len(self._df) * len(other) if joined_table_rows > max_rows: @@ -545,7 +623,9 @@ def join(self, other, instruction: str, model, max_rows: int = 1000): joined_df = self._df.merge(other, how="cross", suffixes=("_left", "_right")) - return joined_df.semantics.filter(instruction, model).reset_index(drop=True) + return joined_df.semantics.filter( + instruction, model, ground_with_google_search=ground_with_google_search + ).reset_index(drop=True) def search( self, @@ -644,7 +724,13 @@ def search( return typing.cast(bigframes.dataframe.DataFrame, search_result) - def top_k(self, instruction: str, model, k=10): + def top_k( + self, + instruction: str, + model, + k: int = 10, + ground_with_google_search: bool = False, + ): """ Ranks each tuple and returns the k best according to the instruction. @@ -682,6 +768,14 @@ def top_k(self, instruction: str, model, k=10): k (int, default 10): The number of rows to return. + ground_with_google_search (bool, default False): + Enables Grounding with Google Search for the GeminiTextGenerator model. + When set to True, the model incorporates relevant information from Google + Search results into its responses, enhancing their accuracy and factualness. + Note: Using this feature may impact billing costs. Refer to the pricing + page for details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models + The default is `False`. + Returns: bigframes.dataframe.DataFrame: A new DataFrame with the top k rows. @@ -703,6 +797,12 @@ def top_k(self, instruction: str, model, k=10): "Semantic aggregations are limited to a single column." ) + if ground_with_google_search: + warnings.warn( + "Enables Grounding with Google Search may impact billing cost. See pricing " + "details: https://cloud.google.com/vertex-ai/generative-ai/pricing#google_models" + ) + df: bigframes.dataframe.DataFrame = self._df[columns].copy() column = columns[0] if df[column].dtype != dtypes.STRING_DTYPE: @@ -743,6 +843,7 @@ def top_k(self, instruction: str, model, k=10): user_instruction, model, k - num_selected, + ground_with_google_search, ) num_selected += num_new_selected @@ -757,7 +858,13 @@ def top_k(self, instruction: str, model, k=10): @staticmethod def _topk_partition( - df, column: str, status_column: str, user_instruction: str, model, k + df, + column: str, + status_column: str, + user_instruction: str, + model, + k: int, + ground_with_google_search: bool, ): output_instruction = ( "Given a question and two documents, choose the document that best answers " @@ -784,7 +891,12 @@ def _topk_partition( import bigframes.dataframe predict_df = typing.cast( - bigframes.dataframe.DataFrame, model.predict(prompt_s, temperature=0.0) + bigframes.dataframe.DataFrame, + model.predict( + prompt_s, + temperature=0.0, + ground_with_google_search=ground_with_google_search, + ), ) marks = predict_df["ml_generate_text_llm_result"].str.contains("2") diff --git a/tests/system/load/test_llm.py b/tests/system/load/test_llm.py index 4b0f50973b..da583f41db 100644 --- a/tests/system/load/test_llm.py +++ b/tests/system/load/test_llm.py @@ -68,6 +68,26 @@ def test_llm_gemini_configure_fit(llm_fine_tune_df_default_index, llm_remote_tex # TODO(ashleyxu b/335492787): After bqml rolled out version control: save, load, check parameters to ensure configuration was kept +@pytest.mark.flaky(retries=2) +def test_llm_gemini_w_ground_with_google_search(llm_remote_text_df): + model = llm.GeminiTextGenerator(model_name="gemini-pro", max_iterations=1) + df = model.predict( + llm_remote_text_df["prompt"], + ground_with_google_search=True, + ).to_pandas() + utils.check_pandas_df_schema_and_index( + df, + columns=[ + "ml_generate_text_llm_result", + "ml_generate_text_rai_result", + "ml_generate_text_grounding_result", + "ml_generate_text_status", + "prompt", + ], + index=3, + ) + + # (b/366290533): Claude models are of extremely low capacity. The tests should reside in small tests. Moving these here just to protect BQML's shared capacity(as load test only runs once per day.) and make sure we still have minimum coverage. @pytest.mark.parametrize( "model_name",