From ba36b21118b5062cb50484b24dcbd676cea433dc Mon Sep 17 00:00:00 2001
From: Baudouin Raoult <baudouin.raoult@ecmwf.int>
Date: Wed, 16 Aug 2023 15:32:10 +0100
Subject: [PATCH] Add retrieve

---
 ai_models/inputs/__init__.py | 1 +
 ai_models/model.py           | 3 +++
 2 files changed, 4 insertions(+)

diff --git a/ai_models/inputs/__init__.py b/ai_models/inputs/__init__.py
index 6a44fbd..4ccbf4a 100644
--- a/ai_models/inputs/__init__.py
+++ b/ai_models/inputs/__init__.py
@@ -29,6 +29,7 @@ def fields_sfc(self):
                     param=self.owner.param_sfc,
                     grid=self.owner.grid,
                     area=self.owner.area,
+                    **self.owner.retrieve,
                 )
                 for date, time in self.owner.datetimes()
             ],
diff --git a/ai_models/model.py b/ai_models/model.py
index 3601be7..6f2acd9 100644
--- a/ai_models/model.py
+++ b/ai_models/model.py
@@ -52,6 +52,7 @@ def add(self, field):
 class Model:
     lagged = False
     assets_extra_dir = None
+    retrieve = {}  # Extra parameters for retrieve
 
     def __init__(self, input, output, download_assets, **kwargs):
         self.input = get_input(input, self, **kwargs)
@@ -252,6 +253,8 @@ def print_requests(self):
             grid=self.grid,
             area=self.area,
         )
+        first.update(self.retrieve)
+
         for date, time in self.datetimes():  # noqa F402
             param, level = self.param_level_pl