google-evalbench 1.7.1__py3-none-any.whl

evalbench/__init__.py

@@ -0,0 +1,12 @@
+import os
+import sys
+
+# Append package root to sys.path for legacy absolute imports.
+# Using append (rather than insert) prevents namespace collisions in spawned child processes.
+sys.path.append(os.path.dirname(__file__))
+
+
+from . import reporting
+from . import util
+from . import dataset
+from . import evaluator

evalbench/client/eval_client.py

@@ -0,0 +1,149 @@
+import asyncio
+import os
+from aiologger import Logger
+import grpc
+from evalproto import eval_request_pb2, eval_connect_pb2, eval_config_pb2
+from evalproto import eval_service_pb2_grpc
+import random
+import argparse
+import google.oauth2.id_token
+import google.auth.transport.requests
+
+
+def get_id_token(audience):
+    """
+    Fetches an ID token for the specified audience.
+    """
+    # The request object is used to make the HTTP call to fetch the token
+    request = google.auth.transport.requests.Request()
+
+    # This will search for credentials in:
+    # 1. Environment variables (GOOGLE_APPLICATION_CREDENTIALS)
+    # 2. Metadata server (if running on GCP)
+    token = google.oauth2.id_token.fetch_id_token(request, audience)
+
+    return token
+
+
+class EvalbenchClient:
+    def __init__(self, endpoint: str):
+        self.endpoint = endpoint
+        if self.endpoint == "local":
+            host = os.getenv("EVALBENCH_HOST", "localhost")
+            port = os.getenv("PORT", "50051")
+            address = f"{host}:{port}"
+            if os.getenv("EVALBENCH_INSECURE", "").lower() == "true":
+                self.channel = grpc.aio.insecure_channel(address)
+            else:
+                channel_creds = grpc.alts_channel_credentials()
+                self.channel = grpc.aio.secure_channel(address, channel_creds)
+        else:
+            address = f"{self.endpoint}:443"
+            id_token = get_id_token(f"https://{self.endpoint}")
+            # 2. Create Call Credentialss
+            # This puts the token in the 'authorization: Bearer <token>' header
+            call_creds = grpc.access_token_call_credentials(id_token)
+
+            # 3. Create SSL Channel Credentials
+            # ID tokens require a secure (TLS) channel
+            channel_creds = grpc.ssl_channel_credentials()
+
+            # 4. Composite Credentials
+            # Combine the SSL channel with the token-based call credentials
+            composite_creds = grpc.composite_channel_credentials(
+                channel_creds, call_creds)
+            self.channel = grpc.aio.secure_channel(address, composite_creds)
+
+        self.stub = eval_service_pb2_grpc.EvalServiceStub(self.channel)
+        rpc_id = "{:032x}".format(random.getrandbits(128))
+        self.metadata = grpc.aio.Metadata(
+            ("client-rpc-id", rpc_id),
+        )
+
+    async def ping(self):
+        request = eval_request_pb2.PingRequest()
+        response = await self.stub.Ping(request, metadata=self.metadata)
+        return response
+
+    async def connect(self):
+        request = eval_connect_pb2.EvalConnectRequest()
+        request.client_id = "me"
+        response = await self.stub.Connect(request, metadata=self.metadata)
+        return response
+
+    async def set_evalconfig(self, experiment: str):
+        data = None
+        with open(experiment, "rb") as f:
+            data = f.read()
+        request = eval_config_pb2.EvalConfigRequest()
+        request.yaml_config = data
+        response = await self.stub.EvalConfig(request, metadata=self.metadata)
+        return response
+
+    async def get_evalinputs(self):
+        request = eval_request_pb2.EvalInputRequest()
+        get_evalinputs_stream = self.stub.ListEvalInputs(
+            request, metadata=self.metadata
+        )
+        while True:
+            response = await get_evalinputs_stream.read()
+            if response == grpc.aio.EOF:
+                break
+            yield response
+
+    async def eval(self, evalinputs):
+        eval_call = self.stub.Eval(metadata=self.metadata)
+        for eval_input in evalinputs:
+            await eval_call.write(eval_input)
+        await eval_call.done_writing()
+        response = await eval_call
+        return response
+
+
+async def run(experiment: str, endpoint: str) -> None:
+    logger = Logger.with_default_handlers(name="evalbench-logger")
+    evalbenchclient = EvalbenchClient(endpoint)
+    response = await evalbenchclient.ping()
+    logger.info(f"ping Returned: {response.response}")
+
+    response = await evalbenchclient.connect()
+    logger.info(f"connect Returned: {response.response}")
+
+    response = await evalbenchclient.ping()
+    logger.info(f"ping Returned: {response.response}")
+
+    response = await evalbenchclient.set_evalconfig(experiment)
+    logger.info(f"get_evalinput Returned: {response.response}")
+
+    evalInputs = []
+    async for response in evalbenchclient.get_evalinputs():
+        evalInputs.append(response)
+    logger.info(f"evalInputs: {len(evalInputs)}")
+    response = await evalbenchclient.eval(evalInputs)
+    logger.info(f"eval Returned: {response.response}")
+
+
+async def main():
+    logger = Logger.with_default_handlers(name="evalbench-logger")
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--experiment", dest="experiment")
+    parser.add_argument("--endpoint", dest="endpoint")
+    known_args, _ = parser.parse_known_args()
+
+    await run(known_args.experiment, known_args.endpoint)
+    # get a set of all running tasks
+    all_tasks = asyncio.all_tasks()
+    # get the current tasks
+    current_task = asyncio.current_task()
+    # remove the current task from the list of all tasks
+    all_tasks.remove(current_task)
+    # report a message
+    print(f"Main waiting for {len(all_tasks)} tasks...")
+    # suspend until all tasks are completed
+    if len(all_tasks) > 0:
+        await asyncio.wait(all_tasks)
+    await logger.shutdown()
+
+
+if __name__ == "__main__":
+    asyncio.run(main())

evalbench/databases/__init__.py

@@ -0,0 +1,42 @@
+from .postgres import PGDB
+from .mysql import MySQLDB
+from .sqlserver import SQLServerDB
+from .sqlite import SQLiteDB
+from .db import DB
+from .bigquery import BQDB
+from .bigtable import BigtableDB
+from .alloydb import AlloyDB
+from .alloydb_omni import AlloyDBOmni
+from .spanner import SpannerDB
+from .mongodb import MongoDB
+
+
+def get_database(db_config, db_name) -> DB:
+    # if db_name is provided:
+    #   - It will override the provided default database_name
+    #   - This is useful as the default db may be "postgres" or a default only used for setup
+    if db_name:
+        suffix = db_config.get("db_name_suffix", "")
+        db_config["database_name"] = f"{db_name}{suffix}"
+
+    if db_config["db_type"] == "postgres":
+        return PGDB(db_config)
+    if db_config["db_type"] == "spanner":
+        return SpannerDB(db_config)
+    if db_config["db_type"] == "mysql":
+        return MySQLDB(db_config)
+    if db_config["db_type"] == "sqlserver":
+        return SQLServerDB(db_config)
+    if db_config["db_type"] == "sqlite":
+        return SQLiteDB(db_config)
+    if db_config["db_type"] == "bigquery":
+        return BQDB(db_config)
+    if db_config["db_type"] == "alloydb":
+        return AlloyDB(db_config)
+    if db_config["db_type"] == "alloydb_omni":
+        return AlloyDBOmni(db_config)
+    if db_config["db_type"] == "bigtable":
+        return BigtableDB(db_config)
+    if db_config["db_type"] == "mongodb":
+        return MongoDB(db_config)
+    raise ValueError("DB Type not Supported")

evalbench/databases/alloydb.py

@@ -0,0 +1,49 @@
+
+from .db import DB
+from .postgres import PGDB
+import sqlalchemy
+from sqlalchemy.pool import NullPool
+from google.cloud.alloydb.connector import Connector as AlloyDBConnector
+from google.cloud.alloydb.connector import IPTypes as AlloyDBIPTypes
+
+CONNECTOR = AlloyDBConnector()
+
+
+class AlloyDB(PGDB):
+    def __init__(self, db_config):
+        """
+        Initializes the AlloyDB connection, overriding the PGDB's
+        default Google Cloud SQL connection mechanism.
+        """
+        super().__init__(db_config)
+        self.nl_config = db_config['nl_config']
+
+        if 'api_endpoint' in db_config:
+            CONNECTOR._alloydb_api_endpoint = db_config['api_endpoint']
+
+        def get_conn_alloydb():
+            return CONNECTOR.connect(
+                self.db_path,
+                "pg8000",
+                user=self.username,
+                password=self.password,
+                db=self.db_name,
+                enable_iam_auth=self.use_adc,  # handled in PGDB
+                ip_type=AlloyDBIPTypes.PUBLIC,
+            )
+
+        def get_engine_args_alloydb():
+            common_args = {
+                "creator": get_conn_alloydb,
+                "connect_args": {"command_timeout": 60},
+            }
+            if "is_tmp_db" in db_config:
+                common_args["poolclass"] = NullPool
+            else:
+                common_args["pool_size"] = 50
+                common_args["pool_recycle"] = 300
+            return common_args
+
+        self.engine = sqlalchemy.create_engine(
+            "postgresql+pg8000://", **get_engine_args_alloydb()
+        )

evalbench/databases/alloydb_omni.py

@@ -0,0 +1,43 @@
+
+from .db import DB
+from .postgres import PGDB
+import sqlalchemy
+import pg8000
+from sqlalchemy.pool import NullPool
+
+
+class AlloyDBOmni(PGDB):
+    def __init__(self, db_config):
+        """
+        Initializes the AlloyDB connection, overriding the PGDB's
+        default Google Cloud SQL connection mechanism.
+        """
+        super().__init__(db_config)
+        self.nl_config = db_config['nl_config']
+        self.host = db_config.get("host", "localhost")
+        self.port = db_config.get("port", 5432)
+
+        def get_conn_alloydb():
+            return pg8000.connect(
+                user=self.username,
+                password=self.password,
+                host=self.host,
+                port=self.port,
+                database=self.db_name
+            )
+
+        def get_engine_args_alloydb():
+            common_args = {
+                "creator": get_conn_alloydb,
+                "connect_args": {"command_timeout": 60},
+            }
+            if "is_tmp_db" in db_config:
+                common_args["poolclass"] = NullPool
+            else:
+                common_args["pool_size"] = 50
+                common_args["pool_recycle"] = 300
+            return common_args
+
+        self.engine = sqlalchemy.create_engine(
+            "postgresql+pg8000://", **get_engine_args_alloydb()
+        )

evalbench/databases/bigquery.py

@@ -0,0 +1,297 @@
+from google.cloud import bigquery
+import logging
+import re
+from .db import DB
+from .util import with_cache_execute, DatabaseSchema
+from util.rate_limit import rate_limit, ResourceExhaustedError
+from typing import List, Optional, Tuple, Any, Dict
+import json
+import sqlparse
+from google.cloud.bigquery import QueryJobConfig, ConnectionProperty
+from util.gcp import get_gcp_project
+from google.api_core.exceptions import GoogleAPICallError
+
+
+class BQDB(DB):
+
+    #####################################################
+    #####################################################
+    # Database Connection Setup Logic
+    #####################################################
+    #####################################################
+
+    def __init__(self, db_config):
+        super().__init__(db_config)
+        self.project_id = get_gcp_project("")
+        self.location = db_config.get("location", "US")
+        self.client = bigquery.Client(project=self.project_id)
+        self.tmp_users = []
+
+    #####################################################
+    #####################################################
+    # Database Specific Execution Logic and Handling
+    #####################################################
+    #####################################################
+
+    def _execute_queries(self, query: str, job_config: Optional[bigquery.QueryJobConfig] = None) -> List:
+        result: List = []
+        for sub_query in sqlparse.split(query):
+            if sub_query:
+                resultset = self.client.query(sub_query, job_config)
+                rows = resultset.result()
+                if rows:
+                    for row in rows:
+                        result.append(dict(row))
+        return result
+
+    def batch_execute(self, commands: list[str]):
+        for command in commands:
+            self.execute(command)
+
+    def execute(
+        self, query: str, eval_query: Optional[str] = None, use_cache=False, rollback=False
+    ) -> Tuple[Any, Any, Any]:
+        if query.strip() == "":
+            return None, None, None
+        if not use_cache or not self.cache_client or eval_query:
+            return self._execute(query, eval_query, rollback)
+        return with_cache_execute(
+            query, f"{self.project_id}.{self.db_name}", self._execute, self.cache_client
+        )
+
+    def _execute(
+        self, query: str, eval_query: Optional[str] = None, rollback=False
+    ) -> Tuple[Any, Any, Any]:
+        def _run_execute(query: str, eval_query: Optional[str] = None, rollback=False):
+            result: List = []
+            eval_result: List = []
+            error = None
+            query_replaced = query.replace("{{dataset}}", self.db_name)
+            if eval_query is not None:
+                eval_query_replaced = eval_query.replace(
+                    "{{dataset}}", self.db_name)
+            try:
+                if rollback:
+                    try:
+                        initial_query = "SELECT 1;"
+                        job_config = QueryJobConfig(create_session=True)
+                        init_job = self.client.query(
+                            initial_query, job_config=job_config)
+                        init_job.result()
+                        session_id = init_job.session_info.session_id
+                        conn_props = [ConnectionProperty(
+                            key="session_id", value=session_id)]
+
+                        self.client.query(
+                            "BEGIN TRANSACTION;",
+                            job_config=QueryJobConfig(
+                                connection_properties=conn_props)
+                        ).result()
+
+                        result = self._execute_queries(
+                            query_replaced, job_config=QueryJobConfig(connection_properties=conn_props))
+
+                        if eval_query:
+                            eval_result = self._execute_queries(
+                                eval_query_replaced, job_config=QueryJobConfig(connection_properties=conn_props))
+
+                        self.client.query(
+                            "ROLLBACK TRANSACTION;",
+                            job_config=QueryJobConfig(
+                                connection_properties=conn_props)
+                        ).result()
+
+                    except Exception as e:
+                        error = str(e)
+                        print(f"Error: {error}")
+
+                    finally:
+                        if 'session_id' in locals():
+                            self.client.query(
+                                "CALL BQ.ABORT_SESSION();",
+                                job_config=QueryJobConfig(
+                                    connection_properties=conn_props)
+                            ).result()
+                if not rollback:
+                    result = self._execute_queries(query_replaced)
+
+                if eval_query and not rollback:
+                    eval_result = self._execute_queries(eval_query_replaced)
+
+            except (GoogleAPICallError, Exception) as e:
+                error = str(e)
+                if "resources exceeded" in error:
+                    raise ResourceExhaustedError(
+                        f"BigQuery resources exhausted: {e}") from e
+                elif "quota exceeded" in error:
+                    raise ResourceExhaustedError(
+                        f"BigQuery quota exceeded: {e}") from e
+                else:
+                    print(error)
+
+            return result, eval_result, error
+
+        try:
+            return rate_limit(
+                (query, eval_query, rollback),
+                _run_execute,
+                self.execs_per_minute,
+                self.semaphore,
+                self.max_attempts,
+            )
+        except ResourceExhaustedError as e:
+            logging.info(
+                "Resource Exhausted on Postgres DB. Giving up execution. Try reducing execs_per_minute."
+            )
+            return None, None, None
+
+    def get_metadata(self) -> dict:
+        metadata = {}
+        try:
+            for table in self.client.list_tables(self.db_name):
+                schema = self.client.get_table(table.reference).schema
+                metadata[table.table_id] = [
+                    {"name": f.name, "type": f.field_type} for f in schema]
+        except Exception as e:
+            print(
+                f"Error while fetching metadata for dataset '{self.db_name}': {e}")
+        return metadata
+
+    #####################################################
+    #####################################################
+    # Setup / Teardown of temporary databases
+    #####################################################
+    #####################################################
+
+    def generate_ddl(self, schema: DatabaseSchema) -> List[str]:
+        ddl_statements = []
+        try:
+            for table in schema.tables:
+                columns = ", ".join(
+                    [f"{col.name} {col.type}" for col in table.columns])
+                ddl_statements.append(
+                    f"CREATE TABLE `{self.project_id}.{self.db_name}.{table.name}` ({columns})"
+                )
+        except Exception as e:
+            print(f"Error generating DDL statements: {e}")
+        return ddl_statements
+
+    def create_tmp_database(self, database_name: str):
+        dataset_ref = bigquery.Dataset(f"{self.project_id}.{database_name}")
+        dataset_ref.location = self.location
+        self.client.create_dataset(dataset_ref, exists_ok=True)
+        self.tmp_dbs.append(database_name)
+
+    def drop_tmp_database(self, database_name: str):
+        try:
+            self.client.delete_dataset(
+                dataset=database_name,
+                delete_contents=True,
+                not_found_ok=True,
+            )
+            if database_name in self.tmp_dbs:
+                self.tmp_dbs.remove(database_name)
+        except Exception as e:
+            logging.warning(f"Failed to drop dataset {database_name}: {e}")
+
+    def drop_all_tables(self):
+        try:
+            tables = list(self.client.list_tables(self.db_name))
+
+            if tables:
+                for table in tables:
+                    full_table_id = f"{self.project_id}.{self.db_name}.{table.table_id}"
+                    self.client.delete_table(full_table_id)
+
+        except Exception as e:
+            raise RuntimeError(
+                f"Failed to drop tables in dataset {self.db_name}: {e}")
+
+    def _is_float(self, value) -> bool:
+        try:
+            float(value)
+        except ValueError:
+            return False
+        return True
+
+    def _get_column_name_to_type_mapping(self, sql_statements: List[str]) -> Dict[str, Dict[str, str]]:
+        schema_mapping = {}
+
+        for statement in sql_statements:
+            table_match = re.search(
+                r'CREATE TABLE\s+`{{dataset}}\.(\w+)`', statement)
+            if not table_match:
+                continue
+            table_name = table_match.group(1)
+            column_section_match = re.search(
+                r'\(\n(.*?)\n\)', statement, re.DOTALL)
+            if not column_section_match:
+                continue
+            columns_raw = column_section_match.group(1).split(",\n")
+            column_type_mapping = {}
+            for col in columns_raw:
+                if col.strip().startswith("PRIMARY KEY"):
+                    continue
+                col_parts = col.strip().split()
+                if len(col_parts) >= 2:
+                    column_name = col_parts[0]
+                    column_type = col_parts[1]
+                    column_type_mapping[column_name] = column_type
+
+            schema_mapping[table_name] = column_type_mapping
+
+        return schema_mapping
+
+    def insert_data(self, data: dict[str, List[str]], setup: Optional[List[str]] = None):
+        if not data:
+            return
+        schema_mapping = self._get_column_name_to_type_mapping(setup)
+        insertion_statements = []
+
+        for table_name in data:
+            column_names = list(schema_mapping[table_name].keys())
+            for row in data[table_name]:
+                formatted_values = []
+
+                for index, value in enumerate(row):
+                    col_name = column_names[index]
+                    col_type = schema_mapping[table_name][col_name].upper()
+
+                    if col_type == 'BOOL':
+                        if value == "'1'":
+                            formatted_values.append("TRUE")
+                        elif value == "'0'":
+                            formatted_values.append("FALSE")
+                        else:
+                            formatted_values.append(f"{value}")
+                    elif self._is_float(value):
+                        formatted_values.append(f"{value}")
+                    elif col_type == 'JSON':
+                        formatted_values.append(f"PARSE_JSON({value})")
+                    else:
+                        escaped_value = value.replace("''", "\\'")
+                        formatted_values.append(f"{escaped_value}")
+
+                inline_columns = ", ".join(formatted_values)
+                insertion_statements.append(
+                    f"INSERT INTO `{self.project_id}.{self.db_name}.{table_name}` VALUES ({inline_columns});"
+                )
+        try:
+            self.batch_execute(insertion_statements)
+        except RuntimeError as error:
+            raise RuntimeError(f"Could not insert data into database: {error}")
+
+    #####################################################
+    #####################################################
+    # Database User Management
+    #####################################################
+    #####################################################
+
+    def close_connections(self):
+        pass
+
+    def create_tmp_users(self, dql_user: str, dml_user: str, tmp_password: str):
+        pass
+
+    def delete_tmp_user(self, username: str):
+        pass