refactor: unify data preparations and research input snapshots

backend/app/research/experiments.py

@@ -11,7 +11,8 @@ from ..backtests.contracts import Candidate, DraftInput, PreviewInput, Simulatio
 from ..backtests.service import Backtests, uid
 from ..catalog.research_metadata import ResearchMetadata
 from ..catalog.service import Catalog
-from ..models import Alpha, BacktestRun, CatalogResource, ResearchExperiment, TemplateInput
+from ..models import Alpha, BacktestRun, CatalogResource, ResearchExperiment
+from ..preparations.service import Preparations
 from .assets import Assets
 from .expressions import GROUPS, analyze, expand
 from .serialization import encode_snapshot as jsonable_encoder
@@ -86,7 +87,7 @@ class Experiments:
                     "candidates": experiment["candidates"],
                     "hypothesis": experiment["hypothesis"],
                     "input_references": [
-                        {k: entry[k] for k in ("id", "collection_version", "scope", "dataset_id")}
+                        {k: entry[k] for k in ("id", "preparation_id", "preparation_version", "scope", "dataset_ids")}
                         for entry in experiment["inputs"]
                     ],
                     "template_reference": {
@@ -133,6 +134,7 @@ class Experiments:
         return validation
 
     async def create(self, body, kind="template", extra_evidence=None, *, parent_snapshots=None):
+        await Preparations(self.db).bind(body)
         asset = await self.assets.get(body.asset_id, body.version, "template") if body.asset_id else None
         template = TemplateSpec.model_validate(asset["content"]) if asset else body.template
         scope = scope_of(body.settings)
@@ -318,7 +320,8 @@ class Experiments:
                         kind=source_kind or experiment["kind"],
                         reference=reference or experiment_id,
                         research_id=experiment_id,
-                        template_input_id=inputs[0]["id"] if len(inputs) == 1 else None,
+                        input_snapshot_ids=[i["id"] for i in inputs],
+                        input_snapshot_id=inputs[0]["id"] if len(inputs) == 1 else None,
                         hypothesis=experiment["hypothesis"][:2000],
                     ),
                     candidates=[
@@ -342,6 +345,7 @@ class Experiments:
         original = parents[0]
         base = seed_settings(original["settings"])
         expression = original["expression"]
+        await Preparations(self.db).bind(body)
         snapshots, _ = await self.inputs(body.input_ids)
         groups = defaultdict(list)
         for snapshot in snapshots:
@@ -405,6 +409,7 @@ class Experiments:
         )
 
     async def generation_context(self, body):
+        await Preparations(self.db).bind(body)
         snapshots, fields = await self.inputs(body.input_ids)
         parents = await self.parents(body.parent_alpha_ids, body.parent_experiment_ids)
         metadata = await ResearchMetadata(self.db).operators(limit=100)
@@ -414,7 +419,7 @@ class Experiments:
             "hypothesis": body.hypothesis,
             "method": body.method,
             "inputs": [
-                {"id": item["id"], "scope": item["scope"], "dataset_id": item["dataset_id"]}
+                {"id": item["id"], "scope": item["scope"], "dataset_ids": item["dataset_ids"], "name": item["name"], "fields": item["fields"][:100]}
                 for item in snapshots
             ],
             "fields": dict(list(fields.items())[:300]),
@@ -424,9 +429,3 @@ class Experiments:
             ],
             "parents": [{**p, "candidates": p.get("candidates", [])[:10]} for p in parents],
         }
-
-    async def available_inputs(self, limit=100):
-        rows = await self.db.scalars(
-            select(TemplateInput).order_by(TemplateInput.created_at.desc()).limit(limit)
-        )
-        return {"items": [await self.catalog.input(row.id) for row in rows]}