From bd0eeacb567c476ac37883e9bd8b77a3046868c7 Mon Sep 17 00:00:00 2001 From: Erin Limbo Date: Wed, 26 Aug 2026 14:36:43 -0700 Subject: [PATCH 1/2] feat(cortado): add multi-turn intermediate SQL execution, Set Match diffing, and tool telemetry --- evalbench/dataset/cortadoinput.py | 13 +- evalbench/eval_service.py | 9 +- evalbench/evaluator/cortadoevaluator.py | 354 ++++++++++++++++++--- evalbench/evaluator/cortadoorchestrator.py | 2 +- evalbench/generators/models/grpc_proxy.py | 10 + evalbench/test/cortadoorchestrator_test.py | 175 ++++++++++ evalbench/work/agentscorework.py | 75 ++++- 7 files changed, 589 insertions(+), 49 deletions(-) diff --git a/evalbench/dataset/cortadoinput.py b/evalbench/dataset/cortadoinput.py index 0eec1e3d..36c08929 100644 --- a/evalbench/dataset/cortadoinput.py +++ b/evalbench/dataset/cortadoinput.py @@ -27,6 +27,7 @@ def __init__(self, raw_dict: dict, job_id: str = "", trace_id: str = ""): self.agent_results = [] self.scoring_results = [] + self.other = {} @classmethod def init_from_proto(cls, proto): @@ -39,23 +40,31 @@ def init_from_proto(cls, proto): raw_dict["id"] = str(getattr(proto, "id", "-1")) - return cls( + obj = cls( raw_dict=raw_dict, job_id=getattr(proto, "job_id", ""), trace_id=getattr(proto, "trace_id", ""), ) + if hasattr(proto, "other"): + for k, v in proto.other.items(): + obj.other[k] = v + return obj def to_proto(self): """Packs the object into the Protobuf to send to Google3.""" # Note: You must import eval_request_pb2 here to prevent circular dependencies from evalproto import eval_request_pb2 - return eval_request_pb2.EvalInputRequest( + proto_req = eval_request_pb2.EvalInputRequest( id=int(self.id) if self.id.isdigit() else 0, payload=self.payload_str, # We map starting_prompt to nl_prompt for backwards compatibility nl_prompt=self.nl_prompt ) + if hasattr(self, "other") and isinstance(self.other, dict): + for k, v in self.other.items(): + proto_req.other[k] = str(v) + return proto_req def copy(self): return copy.deepcopy(self) diff --git a/evalbench/eval_service.py b/evalbench/eval_service.py index bf315503..70afff64 100644 --- a/evalbench/eval_service.py +++ b/evalbench/eval_service.py @@ -411,7 +411,12 @@ async def read_from_client(): logging.debug("Read task cancelled as expected.") # Process final scoring and reporting - job_id, run_time, results_tf, scores_tf = orchestrator.process() + process_res = orchestrator.process() + if len(process_res) == 5: + job_id, run_time, results_tf, scores_tf, multi_trial_scores_tf = process_res + else: + job_id, run_time, results_tf, scores_tf = process_res + multi_trial_scores_tf = None reporters = get_reporters(config.get( "reporting") or {}, job_id, run_time) @@ -427,7 +432,7 @@ async def read_from_client(): run_time, results_tf, scores_tf, - None, # Added None for multi_trial_scores_tf + multi_trial_scores_tf, config, model_config, db_configs, diff --git a/evalbench/evaluator/cortadoevaluator.py b/evalbench/evaluator/cortadoevaluator.py index 64ceb803..2c26c6f1 100644 --- a/evalbench/evaluator/cortadoevaluator.py +++ b/evalbench/evaluator/cortadoevaluator.py @@ -1,11 +1,14 @@ # cortadoevaluator.py -from typing import Any, List, Dict +from typing import Any, List, Dict, Tuple, Optional import datetime import concurrent.futures import logging import json +import re +import threading +import databases from dataset.cortadoinput import EvalCortadoRequest from generators.models.grpc_proxy import GrpcProxyModel from util.config import load_yaml_config @@ -13,11 +16,144 @@ from work.agentgenwork import AgentGenWork from evaluator.simulateduser import SimulatedUser from work.agentscorework import AgentScoreWork +from scorers.setmatcher import SetMatcher + + +def extract_golden_sql_for_turn(scenario: Dict[str, Any], turn: int, dialect: str = "") -> str: + """Extracts expected golden SQL for a specific turn from scenario plan or golden_sql.""" + # 1. Check structured turns + turns = scenario.get("turns") + if isinstance(turns, list) and turn < len(turns) and isinstance(turns[turn], dict): + turn_sql = turns[turn].get("golden_sql") or turns[turn].get("sql") + if turn_sql: + if isinstance(turn_sql, list) and len(turn_sql) > 0: + return str(turn_sql[0]).strip() + return str(turn_sql).strip() + + # 2. Parse from conversation_plan if present + conversation_plan = scenario.get("conversation_plan", "") + plan_text = "" + if isinstance(conversation_plan, list): + if turn < len(conversation_plan): + item = conversation_plan[turn] + if isinstance(item, dict): + sql_val = item.get("golden_sql") or item.get("sql") or item.get("expected_sql") + if sql_val: + return str(sql_val).strip() + plan_text = str(item) + else: + plan_text = "\n".join(str(p) for p in conversation_plan) + elif isinstance(conversation_plan, str): + plan_text = conversation_plan + + if plan_text: + # Check for turn-specific segment in plan (e.g. "Turn 1: ... Turn 2: ...") + turn_num = turn + 1 + turn_pattern = re.compile( + rf'(?:Turn|Step)\s*{turn_num}\b[:\.\-]?\s*(.*?)(?=(?:Turn|Step)\s*\d+\b[:\.\-]|\Z)', + re.DOTALL | re.IGNORECASE, + ) + turn_match = turn_pattern.search(plan_text) + target_text = turn_match.group(1) if turn_match else plan_text + + # Look for SQL patterns within the targeted section + sql_patterns = [ + r"(?:Agent should execute SQL|execute SQL|golden SQL|expected SQL|run SQL|SQL)\s*[:=]\s*[`'\"]([^`'\"]+)[`'\"]", + r"```(?:sql)?\s*([\s\S]*?)\s*```", + r"(?:Agent should execute SQL|execute SQL|golden SQL|expected SQL|run SQL|SQL)\s*[:=]\s*(SELECT\b[\s\S]*?)(?:[\.;\n]|$)", + ] + for pattern in sql_patterns: + match = re.search(pattern, target_text, re.IGNORECASE) + if match: + extracted = match.group(1).strip() + if extracted: + return extracted + + # 3. Fallback to scenario-level golden_sql + golden_sql = scenario.get("golden_sql") + if golden_sql: + if isinstance(golden_sql, dict): + sqls = golden_sql.get(dialect, golden_sql.get("googlesql", [])) + if not sqls and len(golden_sql) > 0: + sqls = list(golden_sql.values())[0] + if isinstance(sqls, list) and len(sqls) > 0: + if turn < len(sqls): + return str(sqls[turn]).strip() + return str(sqls[0]).strip() + elif isinstance(sqls, str): + return sqls.strip() + elif isinstance(golden_sql, list) and len(golden_sql) > 0: + if turn < len(golden_sql): + return str(golden_sql[turn]).strip() + return str(golden_sql[0]).strip() + elif isinstance(golden_sql, str): + return golden_sql.strip() + + return "" + + +def extract_tools_and_skills_from_turn( + eval_result: Any, agent_text: str = "", sql_reply: str = "" +) -> Tuple[List[str], List[str]]: + """Extracts tool calls and skills from eval_result other metadata, agent text, and SQL.""" + tools: List[str] = [] + skills: List[str] = [] + + def _extract_from_obj(obj: Any): + if isinstance(obj, str): + try: + parsed = json.loads(obj) + _extract_from_obj(parsed) + return + except Exception: + pass + for tool_name in ["dataplex_search", "query_data_tool", "execute_sql_tool", "query_data"]: + if tool_name in obj and tool_name not in tools: + tools.append(tool_name) + elif isinstance(obj, dict): + for k, v in obj.items(): + if k in ("actionName", "action_name", "tool_name", "tool", "name") and isinstance(v, str): + if v and v not in tools: + tools.append(v) + elif k in ("tools", "tool_calls", "actions") and isinstance(v, list): + for item in v: + _extract_from_obj(item) + elif k in ("skills", "skill_calls") and isinstance(v, list): + for item in v: + if isinstance(item, str) and item not in skills: + skills.append(item) + else: + _extract_from_obj(v) + elif isinstance(obj, list): + for item in obj: + _extract_from_obj(item) + + other = {} + if hasattr(eval_result, "other") and isinstance(eval_result.other, dict): + other = eval_result.other + elif isinstance(eval_result, dict) and "other" in eval_result and isinstance(eval_result["other"], dict): + other = eval_result["other"] + + for k, v in other.items(): + if "debug" in k.lower() or "tool" in k.lower() or "telemetry" in k.lower() or k == "macchiato_debug_info": + _extract_from_obj(v) + + if sql_reply and sql_reply.strip(): + if not any(t in tools for t in ("query_data_tool", "execute_sql_tool", "query_data")): + tools.append("query_data_tool") + + return tools, skills class CortadoEvaluator: - def __init__(self, config): + def __init__(self, config: Dict[str, Any], db_configs: Optional[Dict[str, Any]] = None): self.config = config + self.db_configs = db_configs or config.get("db_configs", {}) or {} + self._db_lock = threading.Lock() + self._db_cache: Dict[Tuple[str, str], Any] = {} + + # Initialize SetMatcher scorer + self.set_matcher = SetMatcher(self.config.get("scorers", {}).get("set_match", {})) # Load model config model_config = config @@ -31,12 +167,61 @@ def __init__(self, config): self.generator = GrpcProxyModel(model_config) else: raise ValueError( - f"CortadoEvaluator requires 'grpc_proxy' generator, got {generator_type}") + f"CortadoEvaluator requires 'grpc_proxy' generator, got {generator_type}" + ) runner_config = self.config.get("runners", {}) self.agent_runners = runner_config.get("agent_runners", 10) self.agentrunner = mprunner.MPRunner(self.agent_runners) + def _get_db(self, database_name: str, dialect: str = "bigquery") -> Any: + if not database_name: + return None + cache_key = (dialect, database_name) + with self._db_lock: + if cache_key in self._db_cache: + return self._db_cache[cache_key] + + db_cfg = None + if self.db_configs: + if isinstance(self.db_configs, dict): + dialect_configs = self.db_configs.get( + dialect, self.db_configs.get("bigquery", []) + ) + if isinstance(dialect_configs, list) and dialect_configs: + db_cfg = dialect_configs[0] + elif isinstance(dialect_configs, dict): + db_cfg = dialect_configs + elif isinstance(self.db_configs, list) and self.db_configs: + db_cfg = self.db_configs[0] + + if not db_cfg: + db_cfg = self.config.get("database_config") or {"db_type": "bigquery"} + + db_cfg_copy = db_cfg.copy() if isinstance(db_cfg, dict) else {"db_type": "bigquery"} + if "db_type" not in db_cfg_copy: + db_cfg_copy["db_type"] = dialect or "bigquery" + + try: + db = databases.get_database(db_cfg_copy, database_name) + self._db_cache[cache_key] = db + return db + except Exception as e: + logging.warning( + f"Could not initialize database '{database_name}' for dialect '{dialect}': {e}" + ) + return None + + def _execute_sql(self, db: Any, sql_query: str) -> Tuple[Any, Any]: + """Executes SQL against db, returning (result_rows, error_str).""" + if not db or not sql_query or not sql_query.strip(): + return None, None + try: + res, _, err = db.execute(sql_query, use_cache=True, rollback=True) + return res, err + except Exception as e: + return None, str(e) + def evaluate(self, dataset: List[EvalCortadoRequest], job_id: str, run_time: datetime.datetime): eval_outputs: List[Any] = [] scoring_results: List[Any] = [] @@ -58,7 +243,7 @@ def evaluate(self, dataset: List[EvalCortadoRequest], job_id: str, run_time: dat eval_result=item, job_id=job_id, metadata=metadata, - simulated_user=simulated_user + simulated_user=simulated_user, ) self.agentrunner.execute_work(work) @@ -71,58 +256,128 @@ def evaluate(self, dataset: List[EvalCortadoRequest], job_id: str, run_time: dat if hasattr(modified_item, "scoring_results"): scoring_results.extend(modified_item.scoring_results) except Exception as e: - logging.error( - f"Error getting result from future: {e}", exc_info=True) + logging.error(f"Error getting result from future: {e}", exc_info=True) return eval_outputs, scoring_results def process_scenario( - self, scenario: Dict[str, Any], eval_result: Any, job_id: str, - metadata: Dict[str, Any], simulated_user: Any = None + self, + scenario: Dict[str, Any], + eval_result: Any, + job_id: str, + metadata: Dict[str, Any], + simulated_user: Any = None, ) -> Any: """Communication between Cortado and the Simulated User.""" current_prompt = scenario.get("starting_prompt", "") max_turns = scenario.get("max_turns", 1) conversation_plan = scenario.get("conversation_plan", []) - conversation_history = [] + conversation_history: List[Dict[str, str]] = [] + turn_history: List[Dict[str, Any]] = [] last_agent_text = "" last_sql_reply = "" + last_golden_sql = "" + last_gen_res = None + last_golden_res = None + last_gen_err = None + last_golden_err = None + + accumulated_tools: List[str] = [] + accumulated_skills: List[str] = [] - # Parity tracking lists - accumulated_tools = [] - accumulated_skills = [] + database_name = scenario.get("database") or metadata.get("database", "") + dialects = scenario.get("dialects") or metadata.get("dialects", ["bigquery"]) + dialect = dialects[0] if isinstance(dialects, list) and dialects else (dialects if isinstance(dialects, str) else "bigquery") + db = self._get_db(database_name, dialect) for turn in range(max_turns): - logging.info( - f"Turn {turn + 1}/{max_turns} - Prompt: {current_prompt}") + logging.info(f"Turn {turn + 1}/{max_turns} - Prompt: {current_prompt}") # Inject the current prompt into the object eval_result.nl_prompt = current_prompt # Hand it to the gRPC Proxy (blocks until client replies) agent_text = "" + sql_reply = "" try: self.generator.generate(eval_result) nl_reply = getattr(eval_result, "generated_nl_response", "") sql_reply = getattr(eval_result, "generated_sql", "") - last_sql_reply = sql_reply agent_text = nl_reply except Exception as e: - logging.error(f'gRPC generation failed: {e}', exc_info=True) + logging.error(f"gRPC generation failed: {e}", exc_info=True) agent_text = f"Error: {e}" - last_sql_reply = "" + sql_reply = "" last_agent_text = agent_text - logging.info( - f"Turn {turn + 1}/{max_turns} - Agent Reply to Simulated User: {agent_text}") + logging.info(f"Turn {turn + 1}/{max_turns} - Agent Reply to Simulated User: {agent_text}") + + # Extract tools & skills for this turn + turn_tools, turn_skills = extract_tools_and_skills_from_turn( + eval_result, agent_text, sql_reply + ) + accumulated_tools.extend(turn_tools) + accumulated_skills.extend(turn_skills) + + # Golden SQL extraction for this turn + turn_golden_sql = extract_golden_sql_for_turn(scenario, turn, dialect) + + # Execute SQLs against DB if present + golden_res, golden_err = self._execute_sql(db, turn_golden_sql) + gen_res, gen_err = self._execute_sql(db, sql_reply) + + # Calculate Set Match score for this turn + turn_set_match_score = 0.0 + if turn_golden_sql or sql_reply: + try: + score, _ = self.set_matcher.compare( + nl_prompt=current_prompt, + golden_query=turn_golden_sql or "", + query_type="dql", + golden_execution_result=golden_res if golden_res is not None else [], + golden_eval_result="", + golden_error=str(golden_err) if golden_err else "", + generated_query=sql_reply or "", + generated_execution_result=gen_res if gen_res is not None else [], + generated_eval_result="", + generated_error=str(gen_err) if gen_err else "", + ) + turn_set_match_score = float(score) + except Exception as e: + logging.warning(f"SetMatcher error on turn {turn + 1}: {e}") + turn_set_match_score = 0.0 + + # Record turn in turn_history + turn_record = { + "turn": turn + 1, + "user_prompt": current_prompt, + "agent_response": agent_text, + "generated_sql": sql_reply, + "golden_sql": turn_golden_sql, + "generated_execution_result": gen_res, + "generated_error": str(gen_err) if gen_err else None, + "golden_execution_result": golden_res, + "golden_error": str(golden_err) if golden_err else None, + "set_match": turn_set_match_score, + "tools": turn_tools, + } + turn_history.append(turn_record) + + if sql_reply or not last_sql_reply: + last_sql_reply = sql_reply + last_golden_sql = turn_golden_sql + last_gen_res = gen_res + last_golden_res = golden_res + last_gen_err = gen_err + last_golden_err = golden_err # Log history conversation_history.append({ "user": current_prompt, - "agent": agent_text + "agent": agent_text, }) # Invoke Simulated User to check plan and generate next turn @@ -131,8 +386,7 @@ def process_scenario( conversation_plan, conversation_history, agent_text ) if "TERMINATE" in next_response: - logging.info( - "Simulated user met the goal and terminated the conversation.") + logging.info("Simulated user met the goal and terminated the conversation.") break current_prompt = next_response else: @@ -140,19 +394,41 @@ def process_scenario( # Finalize and Score self._finalize_scenario( - scenario, last_agent_text, conversation_history, - accumulated_tools, accumulated_skills, - eval_result, job_id, metadata, - last_sql_reply + scenario=scenario, + last_response=last_agent_text, + conversation_history=conversation_history, + accumulated_tools=accumulated_tools, + accumulated_skills=accumulated_skills, + eval_result=eval_result, + job_id=job_id, + metadata=metadata, + last_sql=last_sql_reply, + turn_history=turn_history, + last_golden_sql=last_golden_sql, + last_gen_res=last_gen_res, + last_golden_res=last_golden_res, + last_gen_err=last_gen_err, + last_golden_err=last_golden_err, ) return eval_result def _finalize_scenario( - self, scenario: Dict[str, Any], last_response: str, + self, + scenario: Dict[str, Any], + last_response: str, conversation_history: List[Dict[str, str]], - accumulated_tools: List[str], accumulated_skills: List[str], - eval_result: Any, job_id: str, metadata: Dict[str, Any], - last_sql: str + accumulated_tools: List[str], + accumulated_skills: List[str], + eval_result: Any, + job_id: str, + metadata: Dict[str, Any], + last_sql: str, + turn_history: Optional[List[Dict[str, Any]]] = None, + last_golden_sql: str = "", + last_gen_res: Any = None, + last_golden_res: Any = None, + last_gen_err: Any = None, + last_golden_err: Any = None, ): """Packages the conversation and sends it to the scoring engine.""" @@ -162,23 +438,27 @@ def _finalize_scenario( "stderr": "", "returncode": 0 if not last_response.startswith("Error") else 1, "prompt_generator_error": None, - "generated_error": None, + "generated_error": str(last_gen_err) if last_gen_err else None, "sql_generator_error": None, - "golden_error": None, - "generated_sql": last_sql, + "golden_error": str(last_golden_err) if last_golden_err else None, + "generated_sql": last_sql if last_sql else "skipped", + "golden_sql": last_golden_sql, + "generated_result": last_gen_res if last_gen_res is not None else accumulated_tools, + "golden_result": last_golden_res if last_golden_res is not None else scenario.get("expected_trajectory", []), "prompt": scenario["starting_prompt"], "conversation_history": json.dumps(conversation_history, indent=2), + "turn_history": turn_history or [], "scenario": scenario, - "accumulated_tools": accumulated_tools, # Passes empty list for now - "accumulated_skills": accumulated_skills, # Passes empty list for now + "accumulated_tools": accumulated_tools, + "accumulated_skills": accumulated_skills, "job_id": job_id, - "metadata": metadata + "metadata": metadata, } score_work = AgentScoreWork( config=self.config, eval_output=eval_output_data, - scoring_results=eval_result.scoring_results + scoring_results=eval_result.scoring_results, ) score_work.run() eval_result.agent_results.append(eval_output_data) diff --git a/evalbench/evaluator/cortadoorchestrator.py b/evalbench/evaluator/cortadoorchestrator.py index d4fa72b6..8fcc9ae4 100644 --- a/evalbench/evaluator/cortadoorchestrator.py +++ b/evalbench/evaluator/cortadoorchestrator.py @@ -18,7 +18,7 @@ def __init__(self, config, db_configs, setup_config, report_progress=False): self.total_scoring_results = [] def evaluate(self, dataset: list[EvalCortadoRequest]): - evaluator = CortadoEvaluator(self.config) + evaluator = CortadoEvaluator(self.config, db_configs=self.db_configs) eval_outputs, scoring_results = evaluator.evaluate( dataset, self.job_id, self.run_time ) diff --git a/evalbench/generators/models/grpc_proxy.py b/evalbench/generators/models/grpc_proxy.py index 78e20b55..4e4581d9 100644 --- a/evalbench/generators/models/grpc_proxy.py +++ b/evalbench/generators/models/grpc_proxy.py @@ -98,14 +98,24 @@ def get_val(obj, *keys, default=None): nl_response = getattr( inbound_response, "generated_nl_response", "") sql_response = getattr(inbound_response, "generated_sql", "") + other_dict = {} + if hasattr(inbound_response, "other"): + for k, v in inbound_response.other.items(): + other_dict[k] = v # Update the eval_output object with the results from the client. if hasattr(eval_output, "__setitem__"): eval_output["generated_sql"] = sql_response eval_output["generated_nl_response"] = nl_response + if "other" not in eval_output or not isinstance(eval_output["other"], dict): + eval_output["other"] = {} + eval_output["other"].update(other_dict) else: setattr(eval_output, "generated_sql", sql_response) setattr(eval_output, "generated_nl_response", nl_response) + if not hasattr(eval_output, "other") or not isinstance(eval_output.other, dict): + setattr(eval_output, "other", {}) + eval_output.other.update(other_dict) return eval_output diff --git a/evalbench/test/cortadoorchestrator_test.py b/evalbench/test/cortadoorchestrator_test.py index 1864b4e4..9f2a8c87 100644 --- a/evalbench/test/cortadoorchestrator_test.py +++ b/evalbench/test/cortadoorchestrator_test.py @@ -49,6 +49,181 @@ def test_evaluate_and_process_returns_5_tuple(self, mock_evaluator_class): self.assertEqual(len(scores_data), 1) self.assertEqual(scores_data[0]["score"], 100) + def test_extract_golden_sql_structured_turns(self): + from evaluator.cortadoevaluator import extract_golden_sql_for_turn + scenario = { + "turns": [ + {"turn": 1, "golden_sql": "SELECT col1 FROM tbl1;"}, + {"turn": 2, "golden_sql": ["SELECT col2 FROM tbl2;"]}, + ] + } + self.assertEqual( + extract_golden_sql_for_turn(scenario, 0), + "SELECT col1 FROM tbl1;", + ) + self.assertEqual( + extract_golden_sql_for_turn(scenario, 1), + "SELECT col2 FROM tbl2;", + ) + + def test_extract_golden_sql_conversation_plan(self): + from evaluator.cortadoevaluator import extract_golden_sql_for_turn + scenario = { + "conversation_plan": ( + "Turn 1: The user asks for spending. Agent should execute SQL: 'SELECT id, sum(spending) FROM schools;'\n" + "Turn 2: The user filters by district. Agent should execute SQL: 'SELECT id, sum(spending) FROM schools WHERE district = 12;'" + ) + } + self.assertEqual( + extract_golden_sql_for_turn(scenario, 0), + "SELECT id, sum(spending) FROM schools;", + ) + self.assertEqual( + extract_golden_sql_for_turn(scenario, 1), + "SELECT id, sum(spending) FROM schools WHERE district = 12;", + ) + + def test_extract_tools_and_skills_telemetry(self): + from evaluator.cortadoevaluator import extract_tools_and_skills_from_turn + eval_result = MagicMock() + eval_result.other = { + "macchiato_debug_info": json.dumps({ + "actions": [ + {"actionName": "dataplex_search"}, + {"actionName": "query_data_tool"} + ] + }) + } + tools, skills = extract_tools_and_skills_from_turn( + eval_result, + agent_text="Here are the results", + sql_reply="SELECT 1;" + ) + self.assertIn("dataplex_search", tools) + self.assertIn("query_data_tool", tools) + + def test_eval_cortado_request_other_proto_roundtrip(self): + from dataset.cortadoinput import EvalCortadoRequest + req = EvalCortadoRequest( + raw_dict={"id": "123", "starting_prompt": "test prompt"} + ) + req.other = {"macchiato_debug_info": "sample_debug_info"} + + proto = req.to_proto() + self.assertEqual(proto.other["macchiato_debug_info"], "sample_debug_info") + + restored = EvalCortadoRequest.init_from_proto(proto) + self.assertEqual(restored.other.get("macchiato_debug_info"), "sample_debug_info") + + @patch("evaluator.cortadoevaluator.databases.get_database") + @patch("evaluator.cortadoevaluator.GrpcProxyModel") + def test_cortado_evaluator_multi_turn_execution(self, mock_grpc_model, mock_get_database): + from evaluator.cortadoevaluator import CortadoEvaluator + from dataset.cortadoinput import EvalCortadoRequest + + mock_db = MagicMock() + mock_db.execute.side_effect = [ + ([{"col": 1}], None, None), # Turn 1 golden + ([{"col": 1}], None, None), # Turn 1 generated + ([{"col": 2}], None, None), # Turn 2 golden + ([{"col": 2}], None, None), # Turn 2 generated + ] + mock_get_database.return_value = mock_db + + mock_generator = MagicMock() + def mock_generate(eval_result): + if "district" in eval_result.nl_prompt: + eval_result.generated_nl_response = "Filtered results" + eval_result.generated_sql = "SELECT col FROM table WHERE district = 1;" + else: + eval_result.generated_nl_response = "Initial results" + eval_result.generated_sql = "SELECT col FROM table;" + mock_generator.generate.side_effect = mock_generate + mock_grpc_model.return_value = mock_generator + + config = { + "model_config": {"generator": "grpc_proxy"}, + "scorers": {"set_match": {}}, + "runners": {"agent_runners": 1} + } + evaluator = CortadoEvaluator(config=config, db_configs={"bigquery": [{"db_type": "bigquery"}]}) + + scenario = { + "id": "scenario_01", + "starting_prompt": "Find all schools", + "max_turns": 2, + "conversation_plan": ( + "Turn 1: Find all schools. Agent should execute SQL: 'SELECT col FROM table;'\n" + "Turn 2: Filter by district. Agent should execute SQL: 'SELECT col FROM table WHERE district = 1;'" + ), + "database": "test_db", + "dialects": ["bigquery"] + } + eval_result = EvalCortadoRequest(raw_dict=scenario) + + simulated_user = MagicMock() + simulated_user.get_next_response.return_value = "Now filter by district 1" + + evaluator.process_scenario( + scenario=scenario, + eval_result=eval_result, + job_id="test_job", + metadata={"dialects": ["bigquery"], "database": "test_db"}, + simulated_user=simulated_user + ) + + self.assertEqual(len(eval_result.agent_results), 1) + final_output = eval_result.agent_results[0] + self.assertIn("turn_history", final_output) + self.assertEqual(len(final_output["turn_history"]), 2) + + turn_1 = final_output["turn_history"][0] + self.assertEqual(turn_1["turn"], 1) + self.assertEqual(turn_1["set_match"], 100.0) + self.assertEqual(turn_1["generated_sql"], "SELECT col FROM table;") + + turn_2 = final_output["turn_history"][1] + self.assertEqual(turn_2["turn"], 2) + self.assertEqual(turn_2["set_match"], 100.0) + self.assertEqual(turn_2["generated_sql"], "SELECT col FROM table WHERE district = 1;") + + @patch("work.agentscorework.scorer.compare") + def test_agent_score_work_populates_turn_data(self, mock_scorer_compare): + from work.agentscorework import AgentScoreWork + eval_output = { + "eval_id": "test_eval", + "scenario": {"starting_prompt": "Test prompt"}, + "metadata": {"dialects": ["bigquery"], "database": "test_db"}, + "turn_history": [ + {"turn": 1, "golden_sql": "SELECT 1;", "generated_sql": "SELECT 1;"}, + {"turn": 2, "golden_sql": "SELECT 2;", "generated_sql": "SELECT 2;"}, + ], + "generated_sql": "SELECT 2;", + "golden_sql": "SELECT 2;", + "generated_result": [{"val": 2}], + "golden_result": [{"val": 2}], + "accumulated_tools": ["dataplex_search", "query_data_tool"], + } + scoring_results = [] + work = AgentScoreWork(config={}, eval_output=eval_output, scoring_results=scoring_results) + work.run() + + mock_scorer_compare.assert_called_once() + call_kwargs = mock_scorer_compare.call_args.kwargs + scoring_item = call_kwargs["eval_output_item"] + + self.assertEqual(scoring_item["golden_sql"], "SELECT 2;") + self.assertEqual(scoring_item["generated_sql"], "SELECT 2;") + self.assertEqual(scoring_item["generated_result"], [{"val": 2}]) + self.assertEqual(scoring_item["accumulated_tools"], ["dataplex_search", "query_data_tool"]) + + # Check multi-turn dual rollup metrics in scoring_results + comparators = {r["comparator"]: r["score"] for r in scoring_results} + self.assertIn("set_match_all_turns", comparators) + self.assertIn("set_match_mean", comparators) + self.assertIn("set_match_turn_1", comparators) + self.assertIn("set_match_turn_2", comparators) + if __name__ == "__main__": unittest.main() diff --git a/evalbench/work/agentscorework.py b/evalbench/work/agentscorework.py index d3ec0d44..335a58ab 100644 --- a/evalbench/work/agentscorework.py +++ b/evalbench/work/agentscorework.py @@ -39,22 +39,44 @@ def run(self, work_config: Any = None) -> Any: """ scenario = self.eval_output.get("scenario", {}) metadata = self.eval_output.get("metadata", {}) + golden_sql = self.eval_output.get("golden_sql", "") + generated_sql = self.eval_output.get("generated_sql", "") + turn_history = self.eval_output.get("turn_history", []) + if not golden_sql and turn_history: + for turn_item in turn_history: + if turn_item.get("golden_sql"): + golden_sql = turn_item["golden_sql"] + break + if not generated_sql and turn_history: + for turn_item in reversed(turn_history): + if turn_item.get("generated_sql"): + generated_sql = turn_item["generated_sql"] + break + golden_result = self.eval_output.get("golden_result") + if golden_result is None: + golden_result = scenario.get("expected_trajectory", []) + generated_result = self.eval_output.get("generated_result") + if generated_result is None: + generated_result = self.eval_output.get("accumulated_tools", []) scoring_item = { "id": self.eval_output.get("eval_id"), "nl_prompt": scenario.get("starting_prompt", ""), - "golden_sql": "", - "query_type": "", - "golden_result": scenario.get("expected_trajectory", []), + "golden_sql": golden_sql, + "query_type": "dql", + "golden_result": golden_result, "golden_eval_results": "", - "golden_error": "", - "generated_sql": "skipped", - "generated_result": self.eval_output.get("accumulated_tools", []), + "golden_error": self.eval_output.get("golden_error", ""), + "generated_sql": generated_sql if generated_sql else "skipped", + "generated_result": generated_result, "eval_results": self.eval_output, - "generated_error": None, + "generated_error": self.eval_output.get("generated_error"), "dialects": metadata.get("dialects", []), "database": metadata.get("database", "unknown"), "job_id": self.eval_output.get("job_id"), + "turn_history": turn_history, + "accumulated_tools": self.eval_output.get("accumulated_tools", []), + "accumulated_skills": self.eval_output.get("accumulated_skills", []), } scorer.compare( @@ -64,4 +86,43 @@ def run(self, work_config: Any = None) -> Any: global_models=self.global_models ) + # Multi-turn rollup metrics calculation + if turn_history: + sql_turns = [t for t in turn_history if t.get("golden_sql") or t.get("generated_sql")] + if sql_turns: + set_match_scores = [t.get("set_match", 0.0) for t in sql_turns] + all_turns_score = 100.0 if all(s == 100.0 for s in set_match_scores) else 0.0 + mean_score = sum(set_match_scores) / len(set_match_scores) + + base_item = { + "id": self.eval_output.get("eval_id"), + "generated_sql": generated_sql if generated_sql else "skipped", + "generated_error": self.eval_output.get("generated_error"), + "dialects": metadata.get("dialects", []), + "database": metadata.get("database", "unknown"), + "job_id": self.eval_output.get("job_id"), + "comparison_logs": None, + "comparison_error": None, + } + + # Record multi-turn aggregate metrics + self.scoring_results.append({ + **base_item, + "comparator": "set_match_all_turns", + "score": all_turns_score, + }) + self.scoring_results.append({ + **base_item, + "comparator": "set_match_mean", + "score": mean_score, + }) + + for t_idx, t in enumerate(turn_history): + if "set_match" in t: + self.scoring_results.append({ + **base_item, + "comparator": f"set_match_turn_{t_idx + 1}", + "score": float(t["set_match"]), + }) + return self.eval_output From 38b90d8a3b36059c32fd4ee8e5962185c1e4e2d6 Mon Sep 17 00:00:00 2001 From: Erin <139916914+erinlimbogoogle@users.noreply.github.com> Date: Wed, 26 Aug 2026 15:49:43 -0700 Subject: [PATCH 2/2] pycheck pycheck --- evalbench/test/cortadoorchestrator_test.py | 1 + 1 file changed, 1 insertion(+) diff --git a/evalbench/test/cortadoorchestrator_test.py b/evalbench/test/cortadoorchestrator_test.py index 9f2a8c87..e7e7c217 100644 --- a/evalbench/test/cortadoorchestrator_test.py +++ b/evalbench/test/cortadoorchestrator_test.py @@ -131,6 +131,7 @@ def test_cortado_evaluator_multi_turn_execution(self, mock_grpc_model, mock_get_ mock_get_database.return_value = mock_db mock_generator = MagicMock() + def mock_generate(eval_result): if "district" in eval_result.nl_prompt: eval_result.generated_nl_response = "Filtered results"