Storage

DynamoDB for experiment rows, S3 for full result JSON. Schema, experiment resolution, and configuration: Persistence and Configuration.

uaef.storage.dynamodb_s3

DynamoDB + S3 storage layer for UAEF experiments and evaluation results.

DynamoDB: One row per experiment with average scores and metadata. S3: Full evaluation results as JSON per experiment.

DynamoS3Storage

Storage layer using DynamoDB for experiment metadata and S3 for full results.

DynamoDB schema (one row per experiment): - experiment_id (PK): str - experiment_name: str - experiment_objective: str | None - created_at: str (ISO 8601) - updated_at: str (ISO 8601) - evaluation_count: int - average_scores: dict (metric_name -> average score) - overall_average_score: Decimal - result_path: str (s3://bucket/prefix/experiment_id.json) - metadata: dict

S3 object (one JSON file per experiment): { "experiment_id": "...", "experiment_name": "...", "evaluations": [ ... full EvaluationResult dicts ... ] }

Source code in uaef/storage/dynamodb_s3.py

class DynamoS3Storage:
    """
    Storage layer using DynamoDB for experiment metadata and S3 for full results.

    DynamoDB schema (one row per experiment):
        - experiment_id (PK): str
        - experiment_name: str
        - experiment_objective: str | None
        - created_at: str (ISO 8601)
        - updated_at: str (ISO 8601)
        - evaluation_count: int
        - average_scores: dict  (metric_name -> average score)
        - overall_average_score: Decimal
        - result_path: str (s3://bucket/prefix/experiment_id.json)
        - metadata: dict

    S3 object (one JSON file per experiment):
        {
            "experiment_id": "...",
            "experiment_name": "...",
            "evaluations": [ ... full EvaluationResult dicts ... ]
        }
    """

    def __init__(self, config: Optional[StorageConfig] = None):
        """Initialize DynamoDB + S3 storage.

        Args:
            config: Storage configuration. Uses global config if None.
        """
        self.config = config or get_config().storage
        region = self.config.region or get_config().aws.region

        self._dynamodb = boto3.resource("dynamodb", region_name=region)
        self._s3 = boto3.client("s3", region_name=region)
        self._table_name = self.config.dynamodb_table_name
        self._region = region
        self._prefix = self.config.s3_prefix
        self._table = self._dynamodb.Table(self._table_name)

        logger.info(
            f"DynamoS3Storage initialized: table={self._table_name}, "
            f"bucket={self.config.s3_bucket or '<account default>'}, prefix={self._prefix}"
        )

    @cached_property
    def account_id(self) -> str:
        """Caller's AWS account, resolved on first S3 use (STS GetCallerIdentity).

        Every S3 call passes it as ExpectedBucketOwner so results are never
        read from or written to a same-named bucket in someone else's account.
        """
        sts = boto3.client("sts", region_name=self._region)
        return sts.get_caller_identity()["Account"]

    @cached_property
    def bucket(self) -> str:
        """Results bucket: the configured name, else uaef-results-<account-id>-<region>.

        Bucket names are global, so a fixed default could resolve to a bucket
        owned by another account.
        """
        return self.config.s3_bucket or f"uaef-results-{self._account_id}-{self._region}"

    @property
    def _bucket(self) -> str:
        return self.bucket

    @property
    def _account_id(self) -> str:
        return self.account_id

    # ------------------------------------------------------------------
    # Table provisioning
    # ------------------------------------------------------------------

    def create_table_if_not_exists(self) -> None:
        """Create the DynamoDB table if it does not already exist."""
        try:
            self._dynamodb.create_table(
                TableName=self._table_name,
                KeySchema=[
                    {"AttributeName": "experiment_id", "KeyType": "HASH"},
                ],
                AttributeDefinitions=[
                    {"AttributeName": "experiment_id", "AttributeType": "S"},
                ],
                BillingMode="PAY_PER_REQUEST",
            )
            self._table.wait_until_exists()
            logger.info(f"Created DynamoDB table '{self._table_name}'")
        except ClientError as e:
            if e.response["Error"]["Code"] == "ResourceInUseException":
                logger.debug(f"Table '{self._table_name}' already exists")
            else:
                raise

    def create_bucket_if_not_exists(self) -> None:
        """Create the S3 bucket if it does not already exist."""
        try:
            self._s3.head_bucket(Bucket=self._bucket, ExpectedBucketOwner=self._account_id)
            logger.debug(f"S3 bucket '{self._bucket}' already exists")
        except ClientError as e:
            error_code = e.response["Error"]["Code"]
            if error_code in ("404", "NoSuchBucket"):
                region = self.config.region or get_config().aws.region
                create_kwargs = {"Bucket": self._bucket}
                # us-east-1 doesn't need LocationConstraint
                if region != "us-east-1":
                    create_kwargs["CreateBucketConfiguration"] = {
                        "LocationConstraint": region
                    }
                self._s3.create_bucket(**create_kwargs)
                logger.info(f"Created S3 bucket '{self._bucket}'")
            elif error_code in ("403", "AccessDenied"):
                # HeadBucket returns 403 both for a bucket owned by another
                # account (ExpectedBucketOwner mismatch) and for an own bucket
                # the caller may not s3:ListBucket.
                raise PermissionError(
                    f"Access denied to S3 bucket '{self._bucket}': either it is owned by "
                    f"another account than {self._account_id}, or this identity lacks "
                    f"s3:ListBucket on it. Grant access, or set UAEF_S3_BUCKET to a "
                    f"bucket this account owns."
                ) from e
            else:
                raise

    # ------------------------------------------------------------------
    # S3 helpers
    # ------------------------------------------------------------------

    def _s3_key(self, experiment_id: str) -> str:
        return f"{self._prefix}{experiment_id}.json"

    def _result_path(self, experiment_id: str) -> str:
        return f"s3://{self._bucket}/{self._s3_key(experiment_id)}"

    def _read_s3_results(
        self, experiment_id: str, result_path: Optional[str] = None
    ) -> Dict[str, Any]:
        """Read the full results JSON from S3.

        ``result_path`` (``s3://bucket/key``, as stored on the experiment) is
        preferred so experiments saved under an earlier bucket/prefix
        configuration stay readable; ownership is still enforced.
        """
        bucket, key = self._bucket, self._s3_key(experiment_id)
        if result_path and result_path.startswith("s3://") and "/" in result_path[5:]:
            bucket, key = result_path[5:].split("/", 1)
        try:
            resp = self._s3.get_object(
                Bucket=bucket,
                ExpectedBucketOwner=self._account_id,
                Key=key,
            )
            return json.loads(resp["Body"].read().decode("utf-8"))
        except ClientError as e:
            if e.response["Error"]["Code"] == "NoSuchKey":
                return {
                    "experiment_id": experiment_id,
                    "evaluations": [],
                }
            raise

    def _write_s3_results(self, experiment_id: str, data: Dict[str, Any]) -> None:
        """Write the full results JSON to S3."""
        self._s3.put_object(
            Bucket=self._bucket,
            ExpectedBucketOwner=self._account_id,
            Key=self._s3_key(experiment_id),
            Body=json.dumps(data, default=_decimal_default, indent=2).encode("utf-8"),
            ContentType="application/json",
        )

    # ------------------------------------------------------------------
    # Core operations
    # ------------------------------------------------------------------

    def save_experiment(
        self,
        experiment_id: str,
        experiment_name: str,
        evaluation_results: List[Dict[str, Any]],
        experiment_objective: Optional[str] = None,
        metadata: Optional[Dict[str, Any]] = None,
        created_by: Optional[str] = None,
    ) -> Dict[str, Any]:
        """
        Save or update an experiment with its evaluation results.

        Computes average scores from all evaluation results, writes the full
        results to S3, and upserts the DynamoDB row.

        Args:
            experiment_id: Unique experiment identifier.
            experiment_name: Human-readable name.
            evaluation_results: List of EvaluationResult dicts (serialized).
            experiment_objective: Optional description/objective.
            metadata: Optional extra metadata.
            created_by: Security review H-02: the identity (e.g. Cognito
                ``sub``) of the caller who created this experiment, used by
                the UAEF Service to scope ``list_experiments``/read access to
                the owner. Optional — library-only callers with no per-caller
                identity concept (no service in front of them) may omit it,
                in which case the experiment has no owner attribute and is
                only reachable directly by ``experiment_id``, never via an
                owner-scoped listing. Preserved across updates (an existing
                experiment's ``created_by`` is never overwritten by a later
                save with a different or missing value), matching how
                ``created_at`` is already preserved below.

        Returns:
            The DynamoDB item dict that was written.
        """
        now = datetime.utcnow().isoformat()

        # --- Compute average scores ---
        metric_totals: Dict[str, List[float]] = {}
        overall_scores: List[float] = []

        for er in evaluation_results:
            if "overall_score" in er and er["overall_score"] is not None:
                overall_scores.append(float(er["overall_score"]))
            for dim in er.get("dimension_results", []):
                for ms in dim.get("metric_scores", []):
                    name = ms.get("metric_name", "")
                    score = ms.get("score")
                    if name and score is not None:
                        metric_totals.setdefault(name, []).append(float(score))

        average_scores = {
            name: round(sum(vals) / len(vals), 4)
            for name, vals in metric_totals.items()
        }
        overall_avg = (
            round(sum(overall_scores) / len(overall_scores), 4)
            if overall_scores
            else 0.0
        )

        # --- Write full results to S3 ---
        s3_data = {
            "experiment_id": experiment_id,
            "experiment_name": experiment_name,
            "experiment_objective": experiment_objective,
            "evaluations": evaluation_results,
        }
        self._write_s3_results(experiment_id, s3_data)

        # --- Upsert DynamoDB row ---
        item = {
            "experiment_id": experiment_id,
            "experiment_name": experiment_name,
            "experiment_objective": experiment_objective or "",
            "created_at": now,
            "updated_at": now,
            "evaluation_count": len(evaluation_results),
            "average_scores": _to_dynamo(average_scores),
            "overall_average_score": _to_dynamo(overall_avg),
            "result_path": self._result_path(experiment_id),
            "metadata": _to_dynamo(metadata or {}),
        }
        if created_by:
            item["created_by"] = created_by

        # Preserve original created_at (and created_by — security review
        # H-02: an experiment's owner must not change on a later save that
        # omits it or names someone else) if the experiment already exists.
        existing = self.get_experiment(experiment_id)
        if existing and "created_at" in existing:
            item["created_at"] = existing["created_at"]
        if existing and existing.get("created_by"):
            item["created_by"] = existing["created_by"]

        self._table.put_item(Item=item)
        logger.info(
            f"Saved experiment {experiment_id}: "
            f"{len(evaluation_results)} evaluations, "
            f"overall_avg={overall_avg}"
        )
        return item

    def get_experiment(self, experiment_id: str) -> Optional[Dict[str, Any]]:
        """
        Get experiment metadata from DynamoDB.

        Args:
            experiment_id: Experiment ID.

        Returns:
            DynamoDB item dict or None if not found.
        """
        try:
            # Consistent read: callers commonly read an experiment straight
            # after save_experiment(), which an eventual read can miss.
            resp = self._table.get_item(
                Key={"experiment_id": experiment_id}, ConsistentRead=True
            )
            return resp.get("Item")
        except ClientError as e:
            logger.warning(f"get_experiment({experiment_id}) failed: {e}")
            return None

    def get_experiment_by_name(self, name: str) -> Optional[Dict[str, Any]]:
        """
        Find an experiment by name via DynamoDB scan.

        Args:
            name: Experiment name to search for.

        Returns:
            DynamoDB item dict or None if not found.
        """
        resp = self._table.scan(
            FilterExpression="experiment_name = :n",
            ExpressionAttributeValues={":n": name},
            Limit=1,
        )
        items = resp.get("Items", [])
        return items[0] if items else None

    def get_full_results(self, experiment_id: str) -> Dict[str, Any]:
        """
        Get the full evaluation results from S3.

        Args:
            experiment_id: Experiment ID.

        Returns:
            Full results dict including all evaluations.
        """
        experiment = self.get_experiment(experiment_id) or {}
        return self._read_s3_results(experiment_id, experiment.get("result_path"))

    #: Name of the GSI on ``created_by`` used to scope ``list_experiments`` to
    #: a single owner (security review H-02). Must match the index name the
    #: CDK-provisioned table declares (see ``uaef-service/infra/storage_stack.py``).
    CREATED_BY_INDEX_NAME = "created_by-index"

    def list_experiments(self, created_by: Optional[str] = None) -> List[Dict[str, Any]]:
        """
        List experiments from DynamoDB.

        Args:
            created_by: Security review H-02: when provided, scopes the
                listing to experiments owned by this identity (e.g. a
                Cognito ``sub``), via the ``created_by`` GSI — never a
                client-side filter over a full scan, so an experiment
                belonging to another caller is never even read out of the
                table. When ``None`` (the default — unchanged library-only
                behavior), lists every experiment via a full table scan, as
                before; library callers with no per-caller identity concept
                have no owner boundary to enforce. The UAEF Service (which
                does have caller identity) always passes ``created_by``.

        Returns:
            List of experiment item dicts.
        """
        items: List[Dict[str, Any]] = []
        if created_by:
            kwargs: Dict[str, Any] = {
                "IndexName": self.CREATED_BY_INDEX_NAME,
                "KeyConditionExpression": "created_by = :cb",
                "ExpressionAttributeValues": {":cb": created_by},
            }
            resp = self._table.query(**kwargs)
            items.extend(resp.get("Items", []))
            while "LastEvaluatedKey" in resp:
                resp = self._table.query(
                    **kwargs, ExclusiveStartKey=resp["LastEvaluatedKey"]
                )
                items.extend(resp.get("Items", []))
            return items

        resp = self._table.scan()
        items.extend(resp.get("Items", []))
        while "LastEvaluatedKey" in resp:
            resp = self._table.scan(ExclusiveStartKey=resp["LastEvaluatedKey"])
            items.extend(resp.get("Items", []))
        return items

    def delete_experiment(self, experiment_id: str) -> None:
        """
        Delete an experiment from DynamoDB and S3.

        Args:
            experiment_id: Experiment ID to delete.
        """
        self._table.delete_item(Key={"experiment_id": experiment_id})
        try:
            self._s3.delete_object(
                Bucket=self._bucket,
                ExpectedBucketOwner=self._account_id,
                Key=self._s3_key(experiment_id),
            )
        except ClientError:
            pass
        logger.info(f"Deleted experiment {experiment_id}")

account_id cached property

Caller's AWS account, resolved on first S3 use (STS GetCallerIdentity).

Every S3 call passes it as ExpectedBucketOwner so results are never read from or written to a same-named bucket in someone else's account.

bucket cached property

Results bucket: the configured name, else uaef-results--.

Bucket names are global, so a fixed default could resolve to a bucket owned by another account.

create_bucket_if_not_exists()

Create the S3 bucket if it does not already exist.

Source code in uaef/storage/dynamodb_s3.py

def create_bucket_if_not_exists(self) -> None:
    """Create the S3 bucket if it does not already exist."""
    try:
        self._s3.head_bucket(Bucket=self._bucket, ExpectedBucketOwner=self._account_id)
        logger.debug(f"S3 bucket '{self._bucket}' already exists")
    except ClientError as e:
        error_code = e.response["Error"]["Code"]
        if error_code in ("404", "NoSuchBucket"):
            region = self.config.region or get_config().aws.region
            create_kwargs = {"Bucket": self._bucket}
            # us-east-1 doesn't need LocationConstraint
            if region != "us-east-1":
                create_kwargs["CreateBucketConfiguration"] = {
                    "LocationConstraint": region
                }
            self._s3.create_bucket(**create_kwargs)
            logger.info(f"Created S3 bucket '{self._bucket}'")
        elif error_code in ("403", "AccessDenied"):
            # HeadBucket returns 403 both for a bucket owned by another
            # account (ExpectedBucketOwner mismatch) and for an own bucket
            # the caller may not s3:ListBucket.
            raise PermissionError(
                f"Access denied to S3 bucket '{self._bucket}': either it is owned by "
                f"another account than {self._account_id}, or this identity lacks "
                f"s3:ListBucket on it. Grant access, or set UAEF_S3_BUCKET to a "
                f"bucket this account owns."
            ) from e
        else:
            raise

create_table_if_not_exists()

Create the DynamoDB table if it does not already exist.

Source code in uaef/storage/dynamodb_s3.py

def create_table_if_not_exists(self) -> None:
    """Create the DynamoDB table if it does not already exist."""
    try:
        self._dynamodb.create_table(
            TableName=self._table_name,
            KeySchema=[
                {"AttributeName": "experiment_id", "KeyType": "HASH"},
            ],
            AttributeDefinitions=[
                {"AttributeName": "experiment_id", "AttributeType": "S"},
            ],
            BillingMode="PAY_PER_REQUEST",
        )
        self._table.wait_until_exists()
        logger.info(f"Created DynamoDB table '{self._table_name}'")
    except ClientError as e:
        if e.response["Error"]["Code"] == "ResourceInUseException":
            logger.debug(f"Table '{self._table_name}' already exists")
        else:
            raise

delete_experiment(experiment_id)

Delete an experiment from DynamoDB and S3.

Parameters:

Name	Type	Description	Default
experiment_id	str	Experiment ID to delete.	required

Source code in uaef/storage/dynamodb_s3.py

def delete_experiment(self, experiment_id: str) -> None:
    """
    Delete an experiment from DynamoDB and S3.

    Args:
        experiment_id: Experiment ID to delete.
    """
    self._table.delete_item(Key={"experiment_id": experiment_id})
    try:
        self._s3.delete_object(
            Bucket=self._bucket,
            ExpectedBucketOwner=self._account_id,
            Key=self._s3_key(experiment_id),
        )
    except ClientError:
        pass
    logger.info(f"Deleted experiment {experiment_id}")

get_experiment(experiment_id)

Get experiment metadata from DynamoDB.

Parameters:

Name	Type	Description	Default
experiment_id	str	Experiment ID.	required

Returns:

Type	Description
Optional[Dict[str, Any]]	DynamoDB item dict or None if not found.

Source code in uaef/storage/dynamodb_s3.py

def get_experiment(self, experiment_id: str) -> Optional[Dict[str, Any]]:
    """
    Get experiment metadata from DynamoDB.

    Args:
        experiment_id: Experiment ID.

    Returns:
        DynamoDB item dict or None if not found.
    """
    try:
        # Consistent read: callers commonly read an experiment straight
        # after save_experiment(), which an eventual read can miss.
        resp = self._table.get_item(
            Key={"experiment_id": experiment_id}, ConsistentRead=True
        )
        return resp.get("Item")
    except ClientError as e:
        logger.warning(f"get_experiment({experiment_id}) failed: {e}")
        return None

get_experiment_by_name(name)

Find an experiment by name via DynamoDB scan.

Parameters:

Name	Type	Description	Default
name	str	Experiment name to search for.	required

Returns:

Type	Description
Optional[Dict[str, Any]]	DynamoDB item dict or None if not found.

Source code in uaef/storage/dynamodb_s3.py

def get_experiment_by_name(self, name: str) -> Optional[Dict[str, Any]]:
    """
    Find an experiment by name via DynamoDB scan.

    Args:
        name: Experiment name to search for.

    Returns:
        DynamoDB item dict or None if not found.
    """
    resp = self._table.scan(
        FilterExpression="experiment_name = :n",
        ExpressionAttributeValues={":n": name},
        Limit=1,
    )
    items = resp.get("Items", [])
    return items[0] if items else None

get_full_results(experiment_id)

Get the full evaluation results from S3.

Parameters:

Name	Type	Description	Default
experiment_id	str	Experiment ID.	required

Returns:

Type	Description
Dict[str, Any]	Full results dict including all evaluations.

Source code in uaef/storage/dynamodb_s3.py

def get_full_results(self, experiment_id: str) -> Dict[str, Any]:
    """
    Get the full evaluation results from S3.

    Args:
        experiment_id: Experiment ID.

    Returns:
        Full results dict including all evaluations.
    """
    experiment = self.get_experiment(experiment_id) or {}
    return self._read_s3_results(experiment_id, experiment.get("result_path"))

list_experiments(created_by=None)

List experiments from DynamoDB.

Parameters:

Name	Type	Description	Default
created_by	Optional[str]	Security review H-02: when provided, scopes the listing to experiments owned by this identity (e.g. a Cognito sub), via the created_by GSI — never a client-side filter over a full scan, so an experiment belonging to another caller is never even read out of the table. When None (the default — unchanged library-only behavior), lists every experiment via a full table scan, as before; library callers with no per-caller identity concept have no owner boundary to enforce. The UAEF Service (which does have caller identity) always passes created_by.	None

Returns:

Type	Description
List[Dict[str, Any]]	List of experiment item dicts.

Source code in uaef/storage/dynamodb_s3.py

def list_experiments(self, created_by: Optional[str] = None) -> List[Dict[str, Any]]:
    """
    List experiments from DynamoDB.

    Args:
        created_by: Security review H-02: when provided, scopes the
            listing to experiments owned by this identity (e.g. a
            Cognito ``sub``), via the ``created_by`` GSI — never a
            client-side filter over a full scan, so an experiment
            belonging to another caller is never even read out of the
            table. When ``None`` (the default — unchanged library-only
            behavior), lists every experiment via a full table scan, as
            before; library callers with no per-caller identity concept
            have no owner boundary to enforce. The UAEF Service (which
            does have caller identity) always passes ``created_by``.

    Returns:
        List of experiment item dicts.
    """
    items: List[Dict[str, Any]] = []
    if created_by:
        kwargs: Dict[str, Any] = {
            "IndexName": self.CREATED_BY_INDEX_NAME,
            "KeyConditionExpression": "created_by = :cb",
            "ExpressionAttributeValues": {":cb": created_by},
        }
        resp = self._table.query(**kwargs)
        items.extend(resp.get("Items", []))
        while "LastEvaluatedKey" in resp:
            resp = self._table.query(
                **kwargs, ExclusiveStartKey=resp["LastEvaluatedKey"]
            )
            items.extend(resp.get("Items", []))
        return items

    resp = self._table.scan()
    items.extend(resp.get("Items", []))
    while "LastEvaluatedKey" in resp:
        resp = self._table.scan(ExclusiveStartKey=resp["LastEvaluatedKey"])
        items.extend(resp.get("Items", []))
    return items

save_experiment(experiment_id, experiment_name, evaluation_results, experiment_objective=None, metadata=None, created_by=None)

Save or update an experiment with its evaluation results.

Computes average scores from all evaluation results, writes the full results to S3, and upserts the DynamoDB row.

Parameters:

Name	Type	Description	Default
experiment_id	str	Unique experiment identifier.	required
experiment_name	str	Human-readable name.	required
evaluation_results	List[Dict[str, Any]]	List of EvaluationResult dicts (serialized).	required
experiment_objective	Optional[str]	Optional description/objective.	None
metadata	Optional[Dict[str, Any]]	Optional extra metadata.	None
created_by	Optional[str]	Security review H-02: the identity (e.g. Cognito sub) of the caller who created this experiment, used by the UAEF Service to scope list_experiments/read access to the owner. Optional — library-only callers with no per-caller identity concept (no service in front of them) may omit it, in which case the experiment has no owner attribute and is only reachable directly by experiment_id, never via an owner-scoped listing. Preserved across updates (an existing experiment's created_by is never overwritten by a later save with a different or missing value), matching how created_at is already preserved below.	None

Returns:

Type	Description
Dict[str, Any]	The DynamoDB item dict that was written.

Source code in uaef/storage/dynamodb_s3.py

def save_experiment(
    self,
    experiment_id: str,
    experiment_name: str,
    evaluation_results: List[Dict[str, Any]],
    experiment_objective: Optional[str] = None,
    metadata: Optional[Dict[str, Any]] = None,
    created_by: Optional[str] = None,
) -> Dict[str, Any]:
    """
    Save or update an experiment with its evaluation results.

    Computes average scores from all evaluation results, writes the full
    results to S3, and upserts the DynamoDB row.

    Args:
        experiment_id: Unique experiment identifier.
        experiment_name: Human-readable name.
        evaluation_results: List of EvaluationResult dicts (serialized).
        experiment_objective: Optional description/objective.
        metadata: Optional extra metadata.
        created_by: Security review H-02: the identity (e.g. Cognito
            ``sub``) of the caller who created this experiment, used by
            the UAEF Service to scope ``list_experiments``/read access to
            the owner. Optional — library-only callers with no per-caller
            identity concept (no service in front of them) may omit it,
            in which case the experiment has no owner attribute and is
            only reachable directly by ``experiment_id``, never via an
            owner-scoped listing. Preserved across updates (an existing
            experiment's ``created_by`` is never overwritten by a later
            save with a different or missing value), matching how
            ``created_at`` is already preserved below.

    Returns:
        The DynamoDB item dict that was written.
    """
    now = datetime.utcnow().isoformat()

    # --- Compute average scores ---
    metric_totals: Dict[str, List[float]] = {}
    overall_scores: List[float] = []

    for er in evaluation_results:
        if "overall_score" in er and er["overall_score"] is not None:
            overall_scores.append(float(er["overall_score"]))
        for dim in er.get("dimension_results", []):
            for ms in dim.get("metric_scores", []):
                name = ms.get("metric_name", "")
                score = ms.get("score")
                if name and score is not None:
                    metric_totals.setdefault(name, []).append(float(score))

    average_scores = {
        name: round(sum(vals) / len(vals), 4)
        for name, vals in metric_totals.items()
    }
    overall_avg = (
        round(sum(overall_scores) / len(overall_scores), 4)
        if overall_scores
        else 0.0
    )

    # --- Write full results to S3 ---
    s3_data = {
        "experiment_id": experiment_id,
        "experiment_name": experiment_name,
        "experiment_objective": experiment_objective,
        "evaluations": evaluation_results,
    }
    self._write_s3_results(experiment_id, s3_data)

    # --- Upsert DynamoDB row ---
    item = {
        "experiment_id": experiment_id,
        "experiment_name": experiment_name,
        "experiment_objective": experiment_objective or "",
        "created_at": now,
        "updated_at": now,
        "evaluation_count": len(evaluation_results),
        "average_scores": _to_dynamo(average_scores),
        "overall_average_score": _to_dynamo(overall_avg),
        "result_path": self._result_path(experiment_id),
        "metadata": _to_dynamo(metadata or {}),
    }
    if created_by:
        item["created_by"] = created_by

    # Preserve original created_at (and created_by — security review
    # H-02: an experiment's owner must not change on a later save that
    # omits it or names someone else) if the experiment already exists.
    existing = self.get_experiment(experiment_id)
    if existing and "created_at" in existing:
        item["created_at"] = existing["created_at"]
    if existing and existing.get("created_by"):
        item["created_by"] = existing["created_by"]

    self._table.put_item(Item=item)
    logger.info(
        f"Saved experiment {experiment_id}: "
        f"{len(evaluation_results)} evaluations, "
        f"overall_avg={overall_avg}"
    )
    return item

get_storage(config=None)

Get or create the global DynamoS3Storage instance.

Source code in uaef/storage/dynamodb_s3.py

def get_storage(config: Optional[StorageConfig] = None) -> DynamoS3Storage:
    """Get or create the global DynamoS3Storage instance."""
    global _storage
    if _storage is None:
        _storage = DynamoS3Storage(config=config)
    return _storage

reset_storage()

Reset the global storage instance.

Source code in uaef/storage/dynamodb_s3.py

def reset_storage() -> None:
    """Reset the global storage instance."""
    global _storage
    _storage = None

uaef.data.ground_truth

Ground truth parsing and validation utilities.

Shared logic for loading, parsing, and validating ground truth datasets from various formats (CSV, Excel, JSON). Used by the demo backend, notebooks, and any future CLI or API endpoints.

ValidationResult

Bases: BaseModel

Result of ground truth validation.

Source code in uaef/data/ground_truth.py

class ValidationResult(BaseModel):
    """Result of ground truth validation."""

    valid: bool
    filename: str = ""
    rows: int = 0
    columns: List[str] = Field(default_factory=list)
    question_column: Optional[str] = None
    answer_column: Optional[str] = None
    warnings: List[str] = Field(default_factory=list)
    errors: List[str] = Field(default_factory=list)
    preview: List[Dict[str, Any]] = Field(default_factory=list)

build_ground_truth(row, events=None)

Build a (query, GroundTruth) pair from a single ground truth row.

If the row has a context column, that value is used. Otherwise, if stream events are provided, tool result contents are extracted as context documents across supported agent frameworks (LangGraph, Strands, etc.).

Parameters:

Name	Type	Description	Default
row	Dict[str, Any]	Dictionary representing a single ground truth row.	required
events	Optional[List[Dict[str, Any]]]	Optional list of agent stream events (for ToolMessage fallback).	None

Returns:

Type	Description
Tuple[str, GroundTruth]	Tuple of (query, GroundTruth).

Source code in uaef/data/ground_truth.py

def build_ground_truth(
    row: Dict[str, Any],
    events: Optional[List[Dict[str, Any]]] = None,
) -> Tuple[str, GroundTruth]:
    """
    Build a (query, GroundTruth) pair from a single ground truth row.

    If the row has a context column, that value is used. Otherwise, if
    stream events are provided, tool result contents are extracted as
    context documents across supported agent frameworks (LangGraph, Strands, etc.).

    Args:
        row: Dictionary representing a single ground truth row.
        events: Optional list of agent stream events (for ToolMessage fallback).

    Returns:
        Tuple of (query, GroundTruth).
    """
    query, expected, context, expected_tools = parse_ground_truth_row(row)

    if context:
        context_documents = [context]
    elif events:
        context_documents = _extract_tool_message_context(events)
    else:
        context_documents = []

    gt = GroundTruth(
        expected_output=expected,
        expected_tool_calls=expected_tools,
        context_documents=context_documents,
    )
    return query, gt

build_multi_turn_ground_truth(turns, events=None)

Build a GroundTruth from multiple conversation turns.

Aggregates expected tool calls across all turns. The top-level expected_output is the combined per-turn expected answers (with turn labels), used by full-trace metrics like ConversationCompletenessMetric. Per-turn expected outputs are also stored in expected_arguments["per_turn_expected"] as a list so that per-turn evaluation can create isolated ground truths per turn.

Context is sourced from the context column of any turn that has it; falls back to tool result contents from the provided events across supported agent frameworks (LangGraph, Strands, etc.).

Parameters:

Name	Type	Description	Default
turns	List[Dict[str, Any]]	List of row dicts, one per conversation turn.	required
events	Optional[List[Dict[str, Any]]]	Optional list of agent stream events (for ToolMessage fallback).	None

Returns:

Type	Description
GroundTruth	Aggregated GroundTruth for the multi-turn session.

Source code in uaef/data/ground_truth.py

def build_multi_turn_ground_truth(
    turns: List[Dict[str, Any]],
    events: Optional[List[Dict[str, Any]]] = None,
) -> GroundTruth:
    """
    Build a GroundTruth from multiple conversation turns.

    Aggregates expected tool calls across all turns. The top-level
    `expected_output` is the combined per-turn expected answers (with
    turn labels), used by full-trace metrics like
    ConversationCompletenessMetric. Per-turn expected outputs are also
    stored in `expected_arguments["per_turn_expected"]` as a list so
    that per-turn evaluation can create isolated ground truths per turn.

    Context is sourced from the context column of any turn that has it;
    falls back to tool result contents from the provided events across
    supported agent frameworks (LangGraph, Strands, etc.).

    Args:
        turns: List of row dicts, one per conversation turn.
        events: Optional list of agent stream events (for ToolMessage fallback).

    Returns:
        Aggregated GroundTruth for the multi-turn session.
    """
    all_tools: List[ToolCall] = []
    context_parts: List[str] = []
    per_turn_expected: List[str] = []
    per_turn_tools: List[List[ToolCall]] = []
    expected_parts: List[str] = []

    for i, turn in enumerate(turns):
        _, expected, context, tools = parse_ground_truth_row(turn)
        all_tools.extend(tools)
        per_turn_expected.append(expected)
        per_turn_tools.append(tools)
        if expected:
            expected_parts.append(f"[Turn {i + 1}] {expected}")
        if context:
            context_parts.append(context)

    if not context_parts and events:
        context_parts = _extract_tool_message_context(events)

    combined_expected = "\n".join(expected_parts) if expected_parts else ""

    return GroundTruth(
        expected_output=combined_expected,
        expected_tool_calls=all_tools,
        context_documents=context_parts,
        expected_arguments={
            "per_turn_expected": per_turn_expected,
            "per_turn_tools": per_turn_tools,
        },
    )

extract_context_from_trace(trace)

Extract context documents from a canonical trace's tool output.

Framework-agnostic counterpart to :func:_extract_tool_message_context, which needs raw provider events and so needs a per-framework extractor. Working from the canonical trace instead means one implementation serves every adapter, and callers that already hold a trace (the service worker, for instance) don't have to keep the raw payload around to get context.

Two sources, in order:

1. ToolCall.result on the trace's tool calls — the intended home for tool output.
2. Message entries with role tool — used by adapters that record tool output as conversation turns rather than on the call.

MultiAgentTrace is walked per sub-agent, because a supervisor's own trace usually holds no tool calls: the work happens in the sub-agents.

Returns an empty list when the trace carries no tool output. That is a real outcome, not a failure: some adapters record tool calls without their results, in which case there is nothing to ground a hallucination check on and the metric should report itself unavailable rather than invent context.

Source code in uaef/data/ground_truth.py

def extract_context_from_trace(trace: Any) -> List[str]:
    """Extract context documents from a canonical trace's tool output.

    Framework-agnostic counterpart to :func:`_extract_tool_message_context`,
    which needs raw provider events and so needs a per-framework extractor.
    Working from the canonical trace instead means one implementation serves
    every adapter, and callers that already hold a trace (the service worker,
    for instance) don't have to keep the raw payload around to get context.

    Two sources, in order:

    1. ``ToolCall.result`` on the trace's tool calls — the intended home for
       tool output.
    2. ``Message`` entries with role ``tool`` — used by adapters that record
       tool output as conversation turns rather than on the call.

    ``MultiAgentTrace`` is walked per sub-agent, because a supervisor's own
    trace usually holds no tool calls: the work happens in the sub-agents.

    Returns an empty list when the trace carries no tool output. That is a real
    outcome, not a failure: some adapters record tool *calls* without their
    results, in which case there is nothing to ground a hallucination check on
    and the metric should report itself unavailable rather than invent context.
    """
    # MultiAgentTrace: recurse into each sub-agent, preserving agent order.
    agent_traces = getattr(trace, "agent_traces", None)
    if isinstance(agent_traces, dict):
        documents: List[str] = []
        for sub_trace in agent_traces.values():
            documents.extend(extract_context_from_trace(sub_trace))
        return documents

    documents = []
    for tool_call in getattr(trace, "tool_calls", None) or []:
        result = getattr(tool_call, "result", None)
        if result is None:
            continue
        text = result if isinstance(result, str) else str(result)
        if text.strip():
            name = getattr(tool_call, "name", "") or "tool"
            documents.append(f"[{name}] {text}")

    for message in getattr(trace, "messages", None) or []:
        role = getattr(message, "role", None)
        # MessageRole is a str enum, so compare on value to accept both forms.
        if str(getattr(role, "value", role)).lower() != "tool":
            continue
        content = getattr(message, "content", None)
        if content is None:
            continue
        text = content if isinstance(content, str) else str(content)
        if text.strip():
            documents.append(text)

    return documents

load_ground_truth_file(path=None, content=None, filename=None)

Load ground truth data from a file path or raw bytes.

Supports CSV and Excel (.xlsx, .xls) formats.

Parameters:

Name	Type	Description	Default
path	Optional[str]	File path to load from. Mutually exclusive with content.	None
content	Optional[bytes]	Raw file bytes. Requires filename for format detection.	None
filename	Optional[str]	Original filename (used for format detection with content).	None

Returns:

Type	Description
DataFrame	DataFrame with ground truth data.

Raises:

Type	Description
ValueError	If neither path nor content is provided, or format is unsupported.

Source code in uaef/data/ground_truth.py

def load_ground_truth_file(
    path: Optional[str] = None,
    content: Optional[bytes] = None,
    filename: Optional[str] = None,
) -> pd.DataFrame:
    """
    Load ground truth data from a file path or raw bytes.

    Supports CSV and Excel (.xlsx, .xls) formats.

    Args:
        path: File path to load from. Mutually exclusive with content.
        content: Raw file bytes. Requires filename for format detection.
        filename: Original filename (used for format detection with content).

    Returns:
        DataFrame with ground truth data.

    Raises:
        ValueError: If neither path nor content is provided, or format is unsupported.
    """
    if path:
        p = Path(path)
        if not p.exists():
            raise ValueError(f"File not found: {path}")
        filename = p.name
        content = p.read_bytes()
    elif content is None:
        raise ValueError("Either path or content must be provided")

    if not filename:
        raise ValueError("filename is required when loading from content")

    fname_lower = filename.lower()
    if fname_lower.endswith(".xlsx") or fname_lower.endswith(".xls"):
        return pd.read_excel(io.BytesIO(content))
    elif fname_lower.endswith(".csv"):
        return pd.read_csv(io.BytesIO(content))
    elif fname_lower.endswith(".json"):
        return pd.DataFrame(json.loads(content))
    else:
        # Default to CSV
        return pd.read_csv(io.BytesIO(content))

parse_ground_truth_dataframe(df)

Parse an entire DataFrame into (query, GroundTruth) pairs.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame with ground truth data.	required

Returns:

Type	Description
List[Tuple[str, GroundTruth]]	List of (query_string, GroundTruth) tuples.

Source code in uaef/data/ground_truth.py

def parse_ground_truth_dataframe(
    df: pd.DataFrame,
) -> List[Tuple[str, GroundTruth]]:
    """
    Parse an entire DataFrame into (query, GroundTruth) pairs.

    Args:
        df: DataFrame with ground truth data.

    Returns:
        List of (query_string, GroundTruth) tuples.
    """
    results = []
    for _, row in df.iterrows():
        query, gt = build_ground_truth(row.to_dict())
        results.append((query, gt))
    return results

parse_ground_truth_row(row, query_only=False)

Extract query, expected answer, context, and expected tool calls from a ground truth row.

Handles case-insensitive column names and multiple naming conventions.

Parameters:

Name	Type	Description	Default
row	Dict[str, Any]	Dictionary representing a single ground truth row.	required
query_only	bool	If True, only extracts and returns the query string. Useful when you need the query before running the agent and will build the GroundTruth separately.	False

Returns:

Type	Description
Union[str, Tuple[str, str, str, List[ToolCall]]]	When query_only is False: Tuple of (query, expected_answer, context, expected_tool_calls).
Union[str, Tuple[str, str, str, List[ToolCall]]]	When query_only is True: query string only.

Source code in uaef/data/ground_truth.py

def parse_ground_truth_row(
    row: Dict[str, Any],
    query_only: bool = False,
) -> Union[str, Tuple[str, str, str, List[ToolCall]]]:
    """
    Extract query, expected answer, context, and expected tool calls from a
    ground truth row.

    Handles case-insensitive column names and multiple naming conventions.

    Args:
        row: Dictionary representing a single ground truth row.
        query_only: If True, only extracts and returns the query string.
                    Useful when you need the query before running the agent
                    and will build the GroundTruth separately.

    Returns:
        When query_only is False: Tuple of (query, expected_answer, context, expected_tool_calls).
        When query_only is True: query string only.
    """
    row_lower = {k.lower().strip(): v for k, v in row.items()}

    query = str(
        row_lower.get("question",
        row_lower.get("query",
        row_lower.get("input", "")))
    )

    if query_only:
        return query

    expected = str(
        row_lower.get("answer",
        row_lower.get("expected_answer",
        row_lower.get("expected_output",
        row_lower.get("expected", ""))))
    )

    raw_context = row_lower.get("context")
    context = str(raw_context) if raw_context is not None and pd.notna(raw_context) else ""

    expected_tools: List[ToolCall] = []
    raw_tools = row_lower.get("expected_tool_calls", row_lower.get("tools", None))
    # `pd.notna` returns an element-wise boolean array for list/Series inputs,
    # so guard against non-scalar values before calling it.
    is_scalar = not isinstance(raw_tools, (list, tuple, dict))
    if raw_tools and (not is_scalar or pd.notna(raw_tools)):
        try:
            parsed = json.loads(str(raw_tools)) if isinstance(raw_tools, str) else raw_tools
            if isinstance(parsed, list):
                for t in parsed:
                    expected_tools.append(ToolCall(
                        name=t.get("name", t.get("tool_name", "")),
                        arguments=t.get("arguments", t.get("parameters", {})),
                        timestamp=datetime.now(timezone.utc),
                    ))
        except (json.JSONDecodeError, TypeError):
            pass

    return query, expected, context, expected_tools

validate_ground_truth(df, filename='', preview_rows=5)

Validate a ground truth DataFrame and return structured results.

Checks for required columns, detects naming conventions, and produces a preview of the data.

Parameters:

Name	Type	Description	Default
df	DataFrame	DataFrame to validate.	required
filename	str	Original filename for reporting.	''
preview_rows	int	Number of rows to include in the preview.	5

Returns:

Type	Description
ValidationResult	ValidationResult with validation status, column mapping, warnings, and preview.

Source code in uaef/data/ground_truth.py

def validate_ground_truth(
    df: pd.DataFrame,
    filename: str = "",
    preview_rows: int = 5,
) -> ValidationResult:
    """
    Validate a ground truth DataFrame and return structured results.

    Checks for required columns, detects naming conventions, and produces
    a preview of the data.

    Args:
        df: DataFrame to validate.
        filename: Original filename for reporting.
        preview_rows: Number of rows to include in the preview.

    Returns:
        ValidationResult with validation status, column mapping, warnings, and preview.
    """
    if df.empty:
        return ValidationResult(valid=False, filename=filename, errors=["File is empty"])

    columns = list(df.columns)
    row_count = len(df)

    col_lower = {c.lower().strip(): c for c in columns}
    question_col = _resolve_column(col_lower, _QUESTION_ALIASES)
    answer_col = _resolve_column(col_lower, _ANSWER_ALIASES)

    warnings = []
    if not question_col:
        warnings.append("No 'Question' column found. The first column will be used as the query.")
        question_col = columns[0]
    if not answer_col:
        warnings.append("No 'Answer' column found. Evaluation will run without ground truth comparison.")

    preview = []
    for i, row in df.head(preview_rows).iterrows():
        preview.append({
            "index": int(i) + 1,
            "question": str(row[question_col])[:200] if question_col else "",
            "answer": str(row[answer_col])[:200] if answer_col else "",
        })

    return ValidationResult(
        valid=True,
        filename=filename,
        rows=row_count,
        columns=columns,
        question_column=question_col,
        answer_column=answer_col,
        warnings=warnings,
        preview=preview,
    )

uaef.utils.serialization

Shared serialization for EvaluationResult objects.

This is the single source of truth for converting EvaluationResult objects to/from JSON-compatible dicts. Used by both: - Local file persistence (save_metric_results / load_metric_results) - S3 persistence (DynamoS3Storage.save_experiment)

deserialize_results(data)

Deserialize a list of dicts back into EvaluationResult objects.

Parameters:

Name	Type	Description	Default
data	List[Dict[str, Any]]	List of dicts (as produced by serialize_results)	required

Returns:

Type	Description
List[EvaluationResult]	List of EvaluationResult objects

Source code in uaef/utils/serialization.py

def deserialize_results(data: List[Dict[str, Any]]) -> List[EvaluationResult]:
    """
    Deserialize a list of dicts back into EvaluationResult objects.

    Args:
        data: List of dicts (as produced by serialize_results)

    Returns:
        List of EvaluationResult objects
    """
    return [EvaluationResult(**item) for item in data]

serialize_results(results)

Serialize a list of EvaluationResult objects to JSON-compatible dicts.

This is the canonical serialization used everywhere results are persisted (local JSON files, S3 storage, API responses).

Parameters:

Name	Type	Description	Default
results	List[EvaluationResult]	List of EvaluationResult objects	required

Returns:

Type	Description
List[Dict[str, Any]]	List of JSON-serializable dicts

Source code in uaef/utils/serialization.py

def serialize_results(results: List[EvaluationResult]) -> List[Dict[str, Any]]:
    """
    Serialize a list of EvaluationResult objects to JSON-compatible dicts.

    This is the canonical serialization used everywhere results are persisted
    (local JSON files, S3 storage, API responses).

    Args:
        results: List of EvaluationResult objects

    Returns:
        List of JSON-serializable dicts
    """
    return [r.model_dump(mode="json") for r in results]

uaef.utils.save_metrics

Utility for saving and loading evaluation metric results.

Uses the same serialization format as the S3 persistence layer (EvaluationResult.model_dump(mode="json")), ensuring that results saved locally can be loaded identically to those retrieved from S3.

The JSON format is

{ "experiment_name": "...", "evaluations": [ {evaluation_result_dict}, ... ] }

This matches the structure written by DynamoS3Storage._write_s3_results().

load_metric_results(path)

Load EvaluationResult objects from a saved JSON file.

Handles both formats: - Local format: {"evaluations": [...], ...} (same as S3) - Legacy format: [...] (flat list of evaluation dicts)

Parameters:

Name	Type	Description	Default
path	str	Path to the JSON file produced by save_metric_results() or downloaded from S3 via get_full_results()	required

Returns:

Type	Description
list	List of EvaluationResult objects

Examples:

>>> from uaef.utils import load_metric_results
>>> results = load_metric_results("output/evaluation-results/my_run.json")

Source code in uaef/utils/save_metrics.py

def load_metric_results(path: str) -> list:
    """
    Load EvaluationResult objects from a saved JSON file.

    Handles both formats:
    - Local format: {"evaluations": [...], ...}  (same as S3)
    - Legacy format: [...] (flat list of evaluation dicts)

    Args:
        path: Path to the JSON file produced by save_metric_results()
              or downloaded from S3 via get_full_results()

    Returns:
        List of EvaluationResult objects

    Examples:
        >>> from uaef.utils import load_metric_results
        >>> results = load_metric_results("output/evaluation-results/my_run.json")
    """

    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)

    # Handle both formats
    if isinstance(data, dict) and "evaluations" in data:
        eval_list = data["evaluations"]
    elif isinstance(data, list):
        eval_list = data
    else:
        raise ValueError(
            f"Unrecognized format in {path}. "
            "Expected a dict with 'evaluations' key or a list of evaluation dicts."
        )

    from uaef.utils.serialization import deserialize_results
    return deserialize_results(eval_list)

save_metric_results(results, test_cases=None, *, prefix='batch_results', output_dir=None, query_key='query', expected_key='expected_output', experiment_name=None)

Save evaluation metric results to timestamped JSON and CSV files.

The JSON uses the same format as DynamoS3Storage (with an "evaluations" wrapper), ensuring local saves and S3 downloads are interchangeable. The CSV is a flat, human-readable companion artifact.

Parameters:

Name	Type	Description	Default
results		List of EvaluationResult objects from evaluate() or batch_evaluate()	required
test_cases	Optional[List[Dict]]	Optional list of test case dicts for the CSV (not required for JSON)	None
prefix	str	Filename prefix (default: "batch_results")	'batch_results'
output_dir	Optional[str]	Output directory (default: BATCH_EVAL_OUTPUT_DIR)	None
query_key	str	Key in test_cases for the query text (CSV only)	'query'
expected_key	str	Key in test_cases for the expected output text (CSV only)	'expected_output'
experiment_name	Optional[str]	Optional name stored in the JSON metadata	None

Returns:

Type	Description
str	Path to the saved JSON file (use this path with load_metric_results)

Examples:

>>> from uaef.utils import save_metric_results
>>> filepath = save_metric_results(results, test_cases, prefix="strands_weather")
>>> # Or without test_cases (JSON-only, no CSV):
>>> filepath = save_metric_results(results, prefix="my_run")

Source code in uaef/utils/save_metrics.py

def save_metric_results(
    results,
    test_cases: Optional[List[Dict]] = None,
    *,
    prefix: str = "batch_results",
    output_dir: Optional[str] = None,
    query_key: str = "query",
    expected_key: str = "expected_output",
    experiment_name: Optional[str] = None,
) -> str:
    """
    Save evaluation metric results to timestamped JSON and CSV files.

    The JSON uses the same format as DynamoS3Storage (with an "evaluations"
    wrapper), ensuring local saves and S3 downloads are interchangeable.
    The CSV is a flat, human-readable companion artifact.

    Args:
        results: List of EvaluationResult objects from evaluate() or batch_evaluate()
        test_cases: Optional list of test case dicts for the CSV (not required for JSON)
        prefix: Filename prefix (default: "batch_results")
        output_dir: Output directory (default: BATCH_EVAL_OUTPUT_DIR)
        query_key: Key in test_cases for the query text (CSV only)
        expected_key: Key in test_cases for the expected output text (CSV only)
        experiment_name: Optional name stored in the JSON metadata

    Returns:
        Path to the saved JSON file (use this path with load_metric_results)

    Examples:
        >>> from uaef.utils import save_metric_results
        >>> filepath = save_metric_results(results, test_cases, prefix="strands_weather")
        >>> # Or without test_cases (JSON-only, no CSV):
        >>> filepath = save_metric_results(results, prefix="my_run")
    """
    output_dir = output_dir or BATCH_EVAL_OUTPUT_DIR
    os.makedirs(output_dir, exist_ok=True)

    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    base_name = f"{prefix}_{timestamp}"

    # --- Save JSON (source of truth, same format as S3) ---
    json_path = os.path.join(output_dir, f"{base_name}.json")
    from uaef.utils.serialization import serialize_results
    json_data = {
        "experiment_name": experiment_name or prefix,
        "saved_at": datetime.now().isoformat(),
        "evaluation_count": len(results),
        "evaluations": serialize_results(results),
    }
    with open(json_path, "w", encoding="utf-8") as f:
        json.dump(json_data, f, indent=2, default=str)

    print(f"\u2713 Results saved to {json_path}")

    # --- Save CSV (human-readable companion) ---
    if test_cases is not None:
        csv_path = os.path.join(output_dir, f"{base_name}.csv")
        rows = []
        for i, r in enumerate(results):
            row = {
                "query_id": i + 1,
                "query": str(test_cases[i].get(query_key, "")) if i < len(test_cases) else "",
                "expected_output": str(test_cases[i].get(expected_key, "")) if i < len(test_cases) else "",
                "overall_score": r.overall_score,
                "passed": r.passed,
            }
            for dim in r.dimension_results:
                row[f"dim_{dim.dimension_name}"] = dim.aggregate_score
                for metric in dim.metric_scores:
                    row[metric.metric_name] = metric.score
            rows.append(row)

        df = pd.DataFrame(rows)
        df.to_csv(csv_path, index=False)
        print(f"\u2713 CSV saved to {csv_path}")

    return json_path