Spaces:
Running
Running
| """MongoDB store for syncing and querying experiment results.""" | |
| import logging | |
| from datetime import datetime, timezone | |
| from typing import Any | |
| from pymongo import MongoClient | |
| from pymongo.collection import Collection | |
| logger = logging.getLogger(__name__) | |
| class MongoDBStore: | |
| """Read/write store for solar-eval experiment data in MongoDB. | |
| Collections: | |
| - experiments: One doc per (project, run_id) — metadata + aggregate scores. | |
| - sample_results: One doc per (experiment_id, sample_idx) — per-sample inference + eval. | |
| - diff_details: One doc per correction diff — category-level detail for qualitative view. | |
| """ | |
| def __init__(self, connection_uri: str, database: str = "solar_eval"): | |
| self.client = MongoClient(connection_uri) | |
| self.db = self.client[database] | |
| self.experiments: Collection = self.db["experiments"] | |
| self.sample_results: Collection = self.db["sample_results"] | |
| self.diff_details: Collection = self.db["diff_details"] | |
| def ensure_indexes(self) -> None: | |
| """Create indexes for efficient querying.""" | |
| # experiments | |
| self.experiments.create_index([("project", 1), ("dataset_group", 1)]) | |
| self.experiments.create_index("prompt_id") | |
| self.experiments.create_index("model") | |
| self.experiments.create_index([("project", 1), ("run_id", 1)], unique=True) | |
| # sample_results | |
| self.sample_results.create_index( | |
| [("experiment_id", 1), ("sample_idx", 1)], unique=True | |
| ) | |
| # diff_details | |
| self.diff_details.create_index([("experiment_id", 1), ("sample_idx", 1)]) | |
| self.diff_details.create_index("category") | |
| def upsert_experiment(self, experiment: dict[str, Any]) -> str: | |
| """Insert or update an experiment document. Returns the _id.""" | |
| key = {"project": experiment["project"], "run_id": experiment["run_id"]} | |
| result = self.experiments.update_one( | |
| key, | |
| {"$set": {**experiment, "synced_at": datetime.now(timezone.utc)}}, | |
| upsert=True, | |
| ) | |
| if result.upserted_id: | |
| return str(result.upserted_id) | |
| doc = self.experiments.find_one(key) | |
| return str(doc["_id"]) if doc else "" | |
| def bulk_insert_samples(self, samples: list[dict[str, Any]]) -> int: | |
| """Bulk upsert sample results. Returns count of upserted/modified docs.""" | |
| if not samples: | |
| return 0 | |
| from pymongo import UpdateOne | |
| ops = [ | |
| UpdateOne( | |
| {"experiment_id": s["experiment_id"], "sample_idx": s["sample_idx"]}, | |
| {"$set": s}, | |
| upsert=True, | |
| ) | |
| for s in samples | |
| ] | |
| result = self.db["sample_results"].bulk_write(ops) | |
| return result.upserted_count + result.modified_count | |
| def bulk_insert_diffs(self, diffs: list[dict[str, Any]]) -> int: | |
| """Bulk insert diff details. Returns count inserted.""" | |
| if not diffs: | |
| return 0 | |
| # Clear existing diffs for this experiment first | |
| if diffs: | |
| exp_id = diffs[0].get("experiment_id") | |
| if exp_id: | |
| self.diff_details.delete_many({"experiment_id": exp_id}) | |
| result = self.diff_details.insert_many(diffs) | |
| return len(result.inserted_ids) | |
| def get_experiments(self, project: str, **filters: Any) -> list[dict[str, Any]]: | |
| """Query experiments with optional filters (dataset_group, model, etc.).""" | |
| query: dict[str, Any] = {"project": project} | |
| query.update(filters) | |
| return list(self.experiments.find(query).sort("created_at", -1)) | |
| def get_samples(self, experiment_id: str, **filters: Any) -> list[dict[str, Any]]: | |
| """Get per-sample results for an experiment.""" | |
| query: dict[str, Any] = {"experiment_id": experiment_id} | |
| query.update(filters) | |
| return list(self.sample_results.find(query).sort("sample_idx", 1)) | |
| def get_diffs(self, experiment_id: str, sample_idx: int | None = None) -> list[dict[str, Any]]: | |
| """Get diff details for an experiment, optionally filtered by sample.""" | |
| query: dict[str, Any] = {"experiment_id": experiment_id} | |
| if sample_idx is not None: | |
| query["sample_idx"] = sample_idx | |
| return list(self.diff_details.find(query).sort("sample_idx", 1)) | |
| def close(self) -> None: | |
| """Close the MongoDB connection.""" | |
| self.client.close() | |