wehub-resource-sync
482 lines
19 KiB
Python
482 lines
19 KiB
Python
## Copyright 2025 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
|
|
## https://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
|
|
|
|
import json
|
|
import logging
|
|
import os
|
|
|
|
import pandas as pd
|
|
import streamlit as st
|
|
from dotenv import load_dotenv
|
|
from google.cloud import storage
|
|
from src import vapo_lib
|
|
|
|
# Load environment variables
|
|
load_dotenv("src/.env")
|
|
|
|
# Configure logging to the console
|
|
logging.basicConfig(
|
|
level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
|
|
)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# --- Constants ---
|
|
BASE_OPTIMIZATION_PREFIX = "optimization/"
|
|
OPTIMIZATION_JOBS_SUBDIR = "optimization_jobs/"
|
|
|
|
from google.cloud import aiplatform
|
|
|
|
|
|
def list_custom_training_jobs(project_id: str, location: str):
|
|
"""Lists all custom training jobs and their statuses in a given project and location.
|
|
|
|
Args:
|
|
project_id: The Google Cloud project ID.
|
|
location: The region for the Agent Platform jobs, e.g., "us-central1".
|
|
|
|
Returns:
|
|
A list of dictionaries, where each dictionary contains details of a custom job.
|
|
"""
|
|
# Initialize the Agent Platform client
|
|
# The API endpoint is determined by the location
|
|
client_options = {"api_endpoint": f"{location}-aiplatform.googleapis.com"}
|
|
client = aiplatform.gapic.JobServiceClient(client_options=client_options)
|
|
|
|
# The parent resource path format
|
|
parent = f"projects/{project_id}/locations/{location}"
|
|
|
|
# Make the API request to list custom jobs
|
|
response = client.list_custom_jobs(parent=parent)
|
|
|
|
# Process the response and format the output
|
|
jobs_list = []
|
|
print(f"Fetching jobs from project '{project_id}' in '{location}'...")
|
|
for job in response:
|
|
job_info = {
|
|
"display_name": job.display_name,
|
|
"name": job.name,
|
|
"status": job.state.name, # .name gets the string representation of the enum
|
|
}
|
|
jobs_list.append(job_info)
|
|
|
|
print(f"Found {len(jobs_list)} jobs.")
|
|
return jobs_list
|
|
|
|
|
|
# --- Example Usage ---
|
|
if __name__ == "__main__":
|
|
# Replace with your project ID and desired location
|
|
PROJECT_ID = os.getenv("PROJECT_ID")
|
|
LOCATION = os.getenv("LOCATION")
|
|
|
|
# Ensure you have authenticated with Google Cloud CLI:
|
|
# gcloud auth application-default login
|
|
|
|
# And have the necessary permissions (e.g., "Agent Platform User" role)
|
|
|
|
try:
|
|
all_jobs = list_custom_training_jobs(project_id=PROJECT_ID, location=LOCATION)
|
|
|
|
# Print the results
|
|
if all_jobs:
|
|
print("\n--- Job Statuses ---")
|
|
for job in all_jobs:
|
|
print(f" - Name: {job['display_name']:<40} Status: {job['status']}")
|
|
print("--------------------\n")
|
|
else:
|
|
print("No custom jobs found.")
|
|
|
|
except Exception as e:
|
|
print(
|
|
"\nAn error occurred. Please ensure your project ID and location are correct,"
|
|
)
|
|
print(f"and that you have authenticated correctly. Error: {e}")
|
|
|
|
|
|
def safe_json_loads(s):
|
|
"""Safely loads a JSON string, returning the original value on failure."""
|
|
if not isinstance(s, str):
|
|
return s
|
|
try:
|
|
return json.loads(s)
|
|
except (json.JSONDecodeError, TypeError):
|
|
return s
|
|
|
|
|
|
@st.cache_data(ttl=300)
|
|
def list_gcs_directories(
|
|
bucket_name: str, prefix: str, _storage_client: storage.Client
|
|
) -> list[str]:
|
|
"""Lists 'directories' in GCS under a given prefix.
|
|
A 'directory' is inferred from the common prefixes of objects.
|
|
Caches the result for 5 minutes to improve performance.
|
|
"""
|
|
if not bucket_name:
|
|
st.warning("BUCKET environment variable is not set.")
|
|
return []
|
|
if not _storage_client:
|
|
st.warning("Storage client is not initialized.")
|
|
return []
|
|
|
|
bucket = _storage_client.bucket(bucket_name)
|
|
retrieved_prefixes = set()
|
|
try:
|
|
for page in bucket.list_blobs(prefix=prefix, delimiter="/").pages:
|
|
retrieved_prefixes.update(page.prefixes)
|
|
|
|
# The retrieved prefixes are the "subdirectories".
|
|
# e.g., for prefix 'optimization/', a retrieved prefix might be 'optimization/op_id/'.
|
|
# We want to extract just 'op_id'.
|
|
dir_names = []
|
|
for p in retrieved_prefixes:
|
|
name = p.replace(prefix, "").strip("/")
|
|
if name:
|
|
dir_names.append(name)
|
|
return sorted(set(dir_names))
|
|
except Exception as e:
|
|
st.error(
|
|
f"Error listing GCS directories under gs://{bucket_name}/{prefix}: {e}"
|
|
)
|
|
logger.error(
|
|
f"Error listing GCS directories under gs://{bucket_name}/{prefix}: {e}",
|
|
exc_info=True,
|
|
)
|
|
return []
|
|
|
|
|
|
def _display_interactive_results(results_ui: vapo_lib.ResultsUI) -> None:
|
|
"""Processes results from a VAPO run and displays them in an interactive
|
|
Streamlit UI with tabs for each prompt version.
|
|
"""
|
|
try:
|
|
if (
|
|
not hasattr(results_ui, "templates")
|
|
or not results_ui.templates
|
|
or not hasattr(results_ui, "eval_results")
|
|
):
|
|
logger.info(
|
|
"ResultsUI object does not have 'templates' or 'eval_results', or templates list is empty. Falling back."
|
|
)
|
|
st.info(
|
|
"No completed runs found yet in this directory. The evaluation might still be running or failed to produce results."
|
|
)
|
|
else:
|
|
processed_results_for_tabs = []
|
|
for i, template_summary_df in enumerate(results_ui.templates):
|
|
if (
|
|
not isinstance(template_summary_df, pd.DataFrame)
|
|
or template_summary_df.empty
|
|
):
|
|
logger.warning(
|
|
f"Template summary data at index {i} is not a non-empty DataFrame. Skipping."
|
|
)
|
|
continue
|
|
|
|
# Get the detailed results to perform the custom calculation
|
|
detailed_eval_df = pd.DataFrame()
|
|
if i < len(results_ui.eval_results) and isinstance(
|
|
results_ui.eval_results[i], pd.DataFrame
|
|
):
|
|
detailed_eval_df = results_ui.eval_results[i]
|
|
|
|
# Add a custom exact_match calculation. This is more robust than simple
|
|
# string comparison as it handles differences in JSON key order and whitespace.
|
|
if (
|
|
not detailed_eval_df.empty
|
|
and "ground_truth" in detailed_eval_df.columns
|
|
and "reference" in detailed_eval_df.columns
|
|
):
|
|
# Parse the JSON strings into Python objects before comparing.
|
|
parsed_ground_truths = detailed_eval_df["ground_truth"].apply(
|
|
safe_json_loads
|
|
)
|
|
parsed_references = detailed_eval_df["reference"].apply(
|
|
safe_json_loads
|
|
)
|
|
|
|
# Create a boolean series for the comparison
|
|
is_match = parsed_ground_truths.eq(parsed_references)
|
|
|
|
# Map boolean to 'yes'/'no' for display in the detailed table
|
|
detailed_eval_df["calculated_exact_match"] = is_match.map(
|
|
{True: "yes", False: "no"}
|
|
)
|
|
|
|
# Calculate the mean from the boolean series for the summary metric
|
|
new_exact_match_mean = is_match.mean()
|
|
template_summary_df["metrics.calculated_exact_match/mean"] = (
|
|
new_exact_match_mean
|
|
)
|
|
|
|
prompt_text = "Prompt text not found in template data."
|
|
if "prompt" in template_summary_df.columns:
|
|
prompt_text = template_summary_df["prompt"].iloc[0]
|
|
else:
|
|
logger.warning(
|
|
f"Column 'prompt' not found in template_summary_df at index {i}."
|
|
)
|
|
|
|
# Determine the primary score and build the tab name.
|
|
primary_score_label = "Score"
|
|
primary_score_value = "N/A"
|
|
if "metrics.calculated_exact_match/mean" in template_summary_df.columns:
|
|
primary_score_label = "Calculated Exact Match"
|
|
primary_score_value = template_summary_df[
|
|
"metrics.calculated_exact_match/mean"
|
|
].iloc[0]
|
|
else:
|
|
# Fallback to the first available metric
|
|
mean_metric_columns = [
|
|
col
|
|
for col in template_summary_df.columns
|
|
if col.startswith("metrics.") and "/mean" in col
|
|
]
|
|
if mean_metric_columns:
|
|
first_metric_col = mean_metric_columns[0]
|
|
primary_score_label = (
|
|
first_metric_col.replace("metrics.", "")
|
|
.replace("/mean", "")
|
|
.replace("_", " ")
|
|
.title()
|
|
)
|
|
primary_score_value = template_summary_df[
|
|
first_metric_col
|
|
].iloc[0]
|
|
|
|
# Build the tab name with all available metrics for a quick overview.
|
|
tab_name_metrics_parts = []
|
|
mean_metric_columns = [
|
|
col
|
|
for col in template_summary_df.columns
|
|
if col.startswith("metrics.") and "/mean" in col
|
|
]
|
|
for metric_col in mean_metric_columns:
|
|
metric_name_short = metric_col.replace("metrics.", "").replace(
|
|
"/mean", ""
|
|
)
|
|
metric_val = template_summary_df[metric_col].iloc[0]
|
|
if metric_name_short == "calculated_exact_match" and isinstance(
|
|
metric_val, float
|
|
):
|
|
tab_name_metrics_parts.append(
|
|
f"{metric_name_short}: {metric_val:.1%}"
|
|
)
|
|
else:
|
|
tab_name_metrics_parts.append(
|
|
f"{metric_name_short}: {metric_val:.3f}"
|
|
if isinstance(metric_val, float)
|
|
else f"{metric_name_short}: {metric_val}"
|
|
)
|
|
|
|
tab_name = f"Template {i}"
|
|
if tab_name_metrics_parts:
|
|
tab_name += f" ({', '.join(tab_name_metrics_parts)})"
|
|
|
|
current_summary_df_display = template_summary_df.copy()
|
|
if "prompt" in current_summary_df_display.columns:
|
|
current_summary_df_display = current_summary_df_display.drop(
|
|
columns=["prompt"]
|
|
)
|
|
|
|
processed_results_for_tabs.append(
|
|
{
|
|
"name": tab_name,
|
|
"template_text": prompt_text,
|
|
"primary_score_label": primary_score_label,
|
|
"primary_score_value": primary_score_value,
|
|
"summary_metrics_df": current_summary_df_display,
|
|
"detailed_eval_df": detailed_eval_df,
|
|
}
|
|
)
|
|
|
|
if (
|
|
processed_results_for_tabs
|
|
): # If we successfully processed data, show the new UI
|
|
st.write("### Interactive Prompt Versions")
|
|
tab_titles = [res["name"] for res in processed_results_for_tabs]
|
|
tabs = st.tabs(tab_titles)
|
|
|
|
for i, tab_content in enumerate(tabs):
|
|
with tab_content:
|
|
result_data = processed_results_for_tabs[i]
|
|
|
|
st.subheader("Prompt Template")
|
|
# Sanitize tab name for key
|
|
clean_key_name = "".join(
|
|
filter(str.isalnum, result_data["name"])
|
|
)
|
|
st.text_area(
|
|
"Template",
|
|
value=result_data["template_text"],
|
|
height=200,
|
|
disabled=True,
|
|
key=f"template_view_{clean_key_name}_{i}",
|
|
)
|
|
|
|
st.subheader("Primary Score")
|
|
score_val = result_data["primary_score_value"]
|
|
score_label = result_data["primary_score_label"]
|
|
if score_label == "Calculated Exact Match" and isinstance(
|
|
score_val, float
|
|
):
|
|
st.metric(label=score_label, value=f"{score_val:.2%}")
|
|
else:
|
|
st.metric(
|
|
label=score_label,
|
|
value=f"{score_val:.4f}"
|
|
if isinstance(score_val, float)
|
|
else str(score_val),
|
|
)
|
|
|
|
if not result_data["summary_metrics_df"].empty:
|
|
st.subheader("Summary Metrics (from templates.json)")
|
|
st.dataframe(result_data["summary_metrics_df"])
|
|
|
|
if not result_data["detailed_eval_df"].empty:
|
|
st.subheader(
|
|
"Detailed Evaluation Results (from eval_results.json)"
|
|
)
|
|
st.dataframe(result_data["detailed_eval_df"])
|
|
else:
|
|
st.caption(
|
|
"No detailed evaluation results available for this template."
|
|
)
|
|
else:
|
|
st.warning("No valid results could be processed for display.")
|
|
|
|
except Exception as e:
|
|
st.error(f"An error occurred while trying to display results: {e}")
|
|
logger.error(f"Error in results display section: {e}", exc_info=True)
|
|
st.markdown(
|
|
"For now, you can access the results directly at the GCS path shown above."
|
|
)
|
|
|
|
|
|
def main() -> None:
|
|
"""Renders the Streamlit page for viewing Prompt Optimization Results."""
|
|
st.set_page_config(
|
|
layout="wide",
|
|
page_title="Prompt Optimization Results",
|
|
page_icon="assets/favicon.ico",
|
|
)
|
|
st.header("Prompt Optimization Results Browser")
|
|
|
|
if "storage_client" not in st.session_state:
|
|
try:
|
|
st.session_state["storage_client"] = storage.Client()
|
|
logger.info("Storage client initialized.")
|
|
except Exception as e:
|
|
st.error(f"Failed to initialize Google Cloud Storage client: {e}")
|
|
logger.error(
|
|
f"Failed to initialize Google Cloud Storage client: {e}", exc_info=True
|
|
)
|
|
st.session_state["storage_client"] = None
|
|
return
|
|
|
|
bucket_name = os.getenv("BUCKET")
|
|
if not bucket_name:
|
|
st.error("BUCKET environment variable is not set. Please configure it in .env.")
|
|
return
|
|
|
|
# --- Step 1: Select Operation ID ---
|
|
op_ids = list_gcs_directories(
|
|
bucket_name, BASE_OPTIMIZATION_PREFIX, st.session_state.storage_client
|
|
)
|
|
if not op_ids:
|
|
st.info(
|
|
f"No optimization operation IDs found under gs://{bucket_name}/{BASE_OPTIMIZATION_PREFIX}"
|
|
)
|
|
return
|
|
|
|
if "op_id" in st.session_state and st.session_state.op_id:
|
|
st.caption(
|
|
f"Hint: The last optimization run you initiated had the ID: `{st.session_state.op_id}`."
|
|
)
|
|
|
|
selected_op_id = st.selectbox(
|
|
"Select an Operation ID:", options=[None, *op_ids], key="selected_op_id_results"
|
|
)
|
|
if not selected_op_id:
|
|
st.write("Please select an Operation ID to see its optimization job runs.")
|
|
return
|
|
|
|
st.divider()
|
|
|
|
# --- Step 2: Select Experiment Run ---
|
|
st.subheader(f"Optimization Job Runs for Operation ID: {selected_op_id}")
|
|
optimization_jobs_prefix = (
|
|
f"{BASE_OPTIMIZATION_PREFIX}{selected_op_id}/{OPTIMIZATION_JOBS_SUBDIR}"
|
|
)
|
|
experiment_runs = list_gcs_directories(
|
|
bucket_name, optimization_jobs_prefix, st.session_state.storage_client
|
|
)
|
|
|
|
if not experiment_runs:
|
|
st.info(
|
|
f"No completed optimization job runs found under gs://{bucket_name}/{optimization_jobs_prefix}"
|
|
)
|
|
return
|
|
|
|
selected_run = st.selectbox(
|
|
"Select an Optimization Job Run:",
|
|
options=[None, *experiment_runs],
|
|
key="selected_experiment_run",
|
|
)
|
|
if not selected_run:
|
|
st.write("Please select an optimization job run to view its results.")
|
|
return
|
|
|
|
st.divider()
|
|
|
|
# --- Step 3: Check Job Status and Display Results ---
|
|
st.subheader(f"Results for: {selected_run}")
|
|
|
|
project_id = os.getenv("PROJECT_ID")
|
|
location = os.getenv("LOCATION")
|
|
|
|
if not project_id or not location:
|
|
st.error("PROJECT_ID or REGION environment variables are not set.")
|
|
return
|
|
|
|
try:
|
|
jobs = list_custom_training_jobs(project_id=project_id, location=location)
|
|
job_status = "Not Found"
|
|
for job in jobs:
|
|
if job["display_name"] == selected_run:
|
|
job_status = job["status"]
|
|
break
|
|
|
|
st.info(f"Status for job '{selected_run}': **{job_status}**")
|
|
|
|
if job_status == "JOB_STATE_FAILED":
|
|
st.error(
|
|
"This optimization job has failed. Please check the logs in the Agent Platform console for more details."
|
|
)
|
|
return
|
|
if job_status not in ["JOB_STATE_SUCCEEDED", "JOB_STATE_CANCELLED"]:
|
|
st.warning(
|
|
f"Job is currently in status: {job_status}. Results may be incomplete."
|
|
)
|
|
|
|
except Exception as e:
|
|
st.error(f"Could not retrieve job status. Error: {e}")
|
|
logger.error(
|
|
f"Failed to retrieve job status for {selected_run}: {e}", exc_info=True
|
|
)
|
|
|
|
run_uri = f"gs://{bucket_name}/{optimization_jobs_prefix}{selected_run}"
|
|
st.info(f"Loading results from: {run_uri}")
|
|
results_ui = vapo_lib.ResultsUI(run_uri)
|
|
_display_interactive_results(results_ui)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|