Updated metric and event retrieval

Author: kzscisoft

simvue/api/objects/base.py

@@ -311,15 +311,15 @@ def ids(
         _class_instance = cls(_read_only=True, _local=True)
         _count: int = 0
         for response in cls._get_all_objects(offset):
-            if count and _count > count:
-                return
             if (_data := response.get("data")) is None:
                 raise RuntimeError(
                     f"Expected key 'data' for retrieval of {_class_instance.__class__.__name__.lower()}s"
                 )
             for entry in _data:
                 yield entry["id"]
                 _count += 1
+                if count and _count > count:
+                    return
 
     @classmethod
     @pydantic.validate_call

simvue/api/objects/events.py

@@ -44,17 +44,19 @@ def get(
         **kwargs,
     ) -> typing.Generator[EventSet, None, None]:
         _class_instance = cls(_read_only=True, _local=True)
-        if (
-            _data := cls._get_all_objects(count, offset, run=run_id, **kwargs).get(
-                "data"
-            )
-        ) is None:
-            raise RuntimeError(
-                f"Expected key 'data' for retrieval of {_class_instance.__class__.__name__.lower()}s"
-            )
-
-        for _entry in _data:
-            yield EventSet(**_entry)
+        _count: int = 0
+
+        for response in cls._get_all_objects(offset, run=run_id, **kwargs):
+            if (_data := response.get("data")) is None:
+                raise RuntimeError(
+                    f"Expected key 'data' for retrieval of {_class_instance.__class__.__name__.lower()}s"
+                )
+
+            for _entry in _data:
+                yield EventSet(**_entry)
+                _count += 1
+                if _count > count:
+                    return
 
     @classmethod
     @pydantic.validate_call

simvue/api/objects/metrics.py

@@ -56,18 +56,17 @@ def get(
         count: pydantic.PositiveInt | None = None,
         offset: pydantic.PositiveInt | None = None,
         **kwargs,
-    ) -> typing.Generator[MetricSet, None, None]:
+    ) -> typing.Generator[dict[str, dict[str, list[dict[str, float]]]], None, None]:
+        """Yields the values for the given metrics for each run."""
         _class_instance = cls(_read_only=True, _local=True)
-        _data = cls._get_all_objects(
-            count,
+        yield from cls._get_all_objects(
             offset,
             metrics=json.dumps(metrics),
             runs=json.dumps(runs),
             xaxis=xaxis,
+            count=count,
             **kwargs,
         )
-        # TODO: Temp fix, just return the dictionary. Not sure what format we really want this in...
-        return _data
 
     @pydantic.validate_call
     def span(self, run_ids: list[str]) -> dict[str, int | float]:

simvue/api/request.py

@@ -320,3 +320,5 @@ def get_paginated(
     ):
         yield _response
         _offset += MAX_ENTRIES_PER_PAGE
+
+    yield _response