Add Conversational Analytics support in EDW service layer.

PiperOrigin-RevId: 930873387

Author: ScottLinnn

perfkitbenchmarker/edw_service.py

@@ -17,12 +17,17 @@
 directory as a subclass of BaseEdwService.
 """
 
+import csv
+import dataclasses
 import datetime
+import io
 import os
 from typing import Any, Dict, List
 
 from absl import flags
 from absl import logging
+from perfkitbenchmarker import data
+from perfkitbenchmarker import errors
 from perfkitbenchmarker import resource
 from perfkitbenchmarker.resources.container_service import kubernetes_cluster
 
@@ -223,13 +228,41 @@
     'Cloud directory of bucket to source ongoing load data '
     'for EDW search benchmarks (without rare token).'
 )
+EDW_CONVERSATIONAL_ANALYTICS_QUESTIONS_FILE = flags.DEFINE_string(
+    'conversational_analytics_questions_file',
+    None,
+    'Filename of the conversational analytics questions CSV (resolved via'
+    ' data_search_paths).',
+)
 EDW_SEARCH_INDEX_NAME = flags.DEFINE_string(
     'edw_search_index_name',
     None,
     'Name of index for edw search index benchmarks',
 )
 
 
+@dataclasses.dataclass
+class ConversationalAnalyticsQuestion:
+  """Representation of a Conversational Analytics Question.
+
+  Attributes:
+    question: The natural language question to be asked.
+    db_id: The database targeted by the question(ecomm or call_center).
+    ground_truth_sql: The ground truth SQL query for the question.
+  """
+
+  def __init__(
+      self,
+      *,
+      question: str,
+      db_id: str,
+      ground_truth_sql: str,
+  ):
+    self.question = question
+    self.db_id = db_id
+    self.ground_truth_sql = ground_truth_sql
+
+
 FLAGS = flags.FLAGS
 
 EDW_PYTHON_DRIVER_LIB_FILE = 'edw_python_driver_lib.py'
@@ -483,6 +516,11 @@ def __init__(self, edw_service_spec):
     self.supports_wait_on_delete = True
     self.client_interface: EdwClientInterface
     self.container_cluster: kubernetes_cluster.KubernetesCluster | None = None
+    self._conversational_analytics_questions: (
+        list[ConversationalAnalyticsQuestion] | None
+    ) = None
+    if EDW_CONVERSATIONAL_ANALYTICS_QUESTIONS_FILE.value:
+      self._LoadConversationalAnalyticsQuestions()
 
   def GetClientInterface(self) -> EdwClientInterface:
     """Gets the active Client Interface."""
@@ -863,6 +901,65 @@ def InjectTokenIntoTable(
     """
     raise NotImplementedError
 
+  def _LoadConversationalAnalyticsQuestions(self) -> None:
+    """Reads the spreadsheet and loads all data into memory."""
+    if not EDW_CONVERSATIONAL_ANALYTICS_QUESTIONS_FILE.value:
+      raise ValueError(
+          'conversational_analytics_questions_file flag must be set.'
+      )
+
+    logging.info(
+        'Loading dataset from resource: %s',
+        EDW_CONVERSATIONAL_ANALYTICS_QUESTIONS_FILE.value,
+    )
+    csv_path = data.ResourcePath(
+        EDW_CONVERSATIONAL_ANALYTICS_QUESTIONS_FILE.value
+    )
+    with open(csv_path, 'r', encoding='utf-8') as f:
+      csv_content = f.read()
+
+    new_question_list = []
+    f = io.StringIO(csv_content)
+    reader = csv.DictReader(f)
+    for row in reader:
+      db_id = row.get('DB ID', '').strip()
+      question = row.get('Question', '').strip()
+      ground_truth_sql = row.get('Ground Truth SQL', '').strip()
+
+      if question and db_id and ground_truth_sql:
+        q = ConversationalAnalyticsQuestion(
+            question=question,
+            db_id=db_id,
+            ground_truth_sql=ground_truth_sql,
+        )
+        new_question_list.append(q)
+
+    if not new_question_list:
+      raise errors.Benchmarks.RunError(
+          f'Loaded 0 questions from {csv_path}.'
+      )
+
+    self._conversational_analytics_questions = new_question_list
+    logging.info(
+        'Loaded %d questions.', len(self._conversational_analytics_questions)
+    )
+
+  def GetConversationalAnalyticsQuestionList(
+      self,
+  ) -> List[ConversationalAnalyticsQuestion]:
+    """Returns the loaded question list."""
+    if self._conversational_analytics_questions is None:
+      raise ValueError('Conversational analytics questions not loaded.')
+    return self._conversational_analytics_questions
+
+  def GetConversationalAnalyticsClientInterface(
+      self,
+  ) -> EdwClientInterface:
+    """Returns the Conversational Analytics Client Interface instance."""
+    raise NotImplementedError(
+        'Conversational Analytics is not supported for this service.'
+    )
+
   @staticmethod
   def ColsToRows(col_res: dict[str, list[Any]]) -> list[dict[str, Any]]:
     """Converts a dictionary of columns to a list of rows.

perfkitbenchmarker/providers/gcp/bigquery.py

@@ -15,6 +15,7 @@
 
 import copy
 import datetime
+import hashlib
 import json
 import logging
 import os
@@ -25,6 +26,7 @@
 from absl import flags
 from perfkitbenchmarker import data
 from perfkitbenchmarker import edw_service
+from perfkitbenchmarker import errors
 from perfkitbenchmarker import provider_info
 from perfkitbenchmarker import vm_util
 from perfkitbenchmarker.linux_packages import google_cloud_sdk
@@ -40,6 +42,13 @@
     ' benchmarking.',
 )
 
+BQ_CA_AGENT = flags.DEFINE_string(
+    'bq_ca_agent',
+    None,
+    'The full resource name of the agent to use, e.g. '
+    'projects/{project}/locations/{location}/dataAgents/{agent_id}',
+)
+
 BQ_CLIENT_FILE = 'bq-jdbc-simba-client-1.8-temp-labels.jar'
 BQ_PYTHON_CLIENT_FILE = 'bq_python_driver.py'
 BQ_PYTHON_CLIENT_DIR = 'edw/bigquery/clients/python'
@@ -508,6 +517,116 @@ def RunQueryWithResults(self, query_name: str) -> str:
     return stdout
 
 
+class ConversationalAnalyticsClientInterface(PythonClientInterface):
+  """Conversational Analytics Client Interface subclassing PythonClientInterface."""
+
+  def _GetQueryFileName(self, query_name: str) -> str:
+    """Generates a filename from a query name."""
+    sanitized = re.sub(r'[^a-zA-Z0-9_-]', '_', query_name)
+    # Limit sanitized part to 30 chars
+    sanitized = sanitized[:30]
+    # Add hash of full query name to ensure uniqueness
+    query_hash = hashlib.md5(query_name.encode('utf-8')).hexdigest()[:8]
+    return f'./{sanitized}_{query_hash}.txt'
+
+  def _ParseConversationalAnalyticsResults(
+      self, results: dict[str, Any], query_name: str
+  ) -> tuple[float, dict[str, Any]]:
+    """Parses the results from Conversational Analytics query execution."""
+    execution_time = results.get('query_wall_time_in_secs', -1.0)
+    details = results.get('details', {})
+    query_results = details.get('query_results', {})
+
+    # Extract essential fields
+    text_response = query_results.get('text_response')
+    generated_sql = query_results.get('generated_sql')
+    retrieved_data = query_results.get('retrieved_data')
+    # Fail fast validation
+    error_msg = None
+    if not text_response:
+      error_msg = f"'text_response' is missing or empty. Got: {text_response!r}"
+    elif not generated_sql:
+      error_msg = f"'generated_sql' is missing or empty. Got: {generated_sql!r}"
+    elif not retrieved_data:
+      error_msg = (
+          f"'retrieved_data' is missing or empty. Got: {retrieved_data!r}"
+      )
+
+    metadata = {
+        'question': query_name,
+        'text_response': text_response or '',
+        'generated_sql': generated_sql or '',
+        'retrieved_data': retrieved_data or [],
+        'thoughts': query_results.get('thoughts', []),
+        'progress_messages': query_results.get('progress_messages', []),
+        'time_to_first_token_secs': query_results.get(
+            'time_to_first_token_secs', 0.0
+        ),
+        'total_stream_time_secs': query_results.get(
+            'total_stream_time_secs', 0.0
+        ),
+        'job_id': details.get('job_id', ''),
+    }
+
+    if error_msg:
+      logging.warning('Conversational Analytics query failed: %s', error_msg)
+      metadata['error'] = f'Conversational Analytics query failed: {error_msg}'
+      return -1.0, metadata
+
+    logging.info(
+        'Conversational Analytics Response: %s',
+        metadata.get('text_response', ''),
+    )
+    return execution_time, metadata
+
+  @override
+  def Prepare(self, package_name: str) -> None:
+    """Prepares the client vm by installing geminidataanalytics package and driver."""
+    super().Prepare(package_name)
+
+    # Install BQ CA specific SDK
+    self.client_vm.RemoteCommand(
+        'source .venv/bin/activate && pip install'
+        ' google-cloud-geminidataanalytics'
+    )
+
+    # Push Conversational Analytics driver script.
+    driver_local_path = data.ResourcePath(
+        os.path.join('edw/bigquery/clients/python', 'bq_ca_driver.py')
+    )
+    self.client_vm.PushFile(driver_local_path, 'bq_ca_driver.py')
+
+  @override
+  def ExecuteQuery(
+      self, query_name: str, print_results: bool = False
+  ) -> tuple[float, dict[str, Any]]:
+    """Executes a single conversational analytics question."""
+    logging.info('Conversational Analytics Question: %s', query_name)
+    key_file = os.path.basename(FLAGS.gcp_service_account_key_file)
+    cmd = (
+        'source .venv/bin/activate && python3 bq_ca_driver.py single '
+        f'--project={self.project_id} '
+        f'--agent={BQ_CA_AGENT.value} '
+        f'--credentials_file={key_file} '
+        '--print_results '
+    )
+
+    remote_query_file = self._GetQueryFileName(query_name)
+    # Write question to remote file to handle quotes securely
+    vm_util.CreateRemoteFile(self.client_vm, query_name, remote_query_file)
+    cmd += f'--query_file={remote_query_file}'
+
+    stdout, _ = self.client_vm.RemoteCommand(cmd)
+
+    try:
+      results = json.loads(stdout)
+    except ValueError as e:
+      raise errors.Benchmarks.RunError(
+          f'Failed to parse Conversational Analytics response: {stdout}'
+      ) from e
+    return self._ParseConversationalAnalyticsResults(results, query_name)
+
+
 class Bigquery(edw_service.EdwService):
   """Object representing a Bigquery cluster.
 
@@ -526,6 +645,14 @@ def __init__(self, edw_service_spec):
     project_id, dataset_id = _SplitClusterIdentifier(self.cluster_identifier)
     self.client_interface = GetBigQueryClientInterface(project_id, dataset_id)
 
+  @override
+  def GetConversationalAnalyticsClientInterface(
+      self,
+  ) -> ConversationalAnalyticsClientInterface:
+    """Returns the Conversational Analytics Client Interface instance."""
+    project_id, dataset_id = _SplitClusterIdentifier(self.cluster_identifier)
+    return ConversationalAnalyticsClientInterface(project_id, dataset_id)
+
   def _Create(self):
     """Create a BigQuery cluster.