fix rag metrics (#2221)

nuclia · Jun 5, 2024 · 12ad092 · 12ad092
1 parent d7a3531
commit 12ad092
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/nucliadb/src/nucliadb/search/search/find.py b/nucliadb/src/nucliadb/search/search/find.py
@@ -149,8 +149,8 @@ async def find(
                 "query": item.model_dump_json(),
                 "time": search_time,
                 "nodes": debug_nodes_info(queried_nodes),
-                # Include step times in the log
-                **{step: metrics.elapsed(step) for step in metrics.steps()},
+                # Include metrics in the log
+                **metrics.steps(),
             },
         )
 

diff --git a/nucliadb/src/nucliadb/search/search/metrics.py b/nucliadb/src/nucliadb/search/search/metrics.py
@@ -75,8 +75,8 @@ def time(self, step: str):
         finally:
             self._end(step)
 
-    def steps(self):
-        return list(self._start_times.keys())
+    def steps(self) -> dict[str, float]:
+        return {step: self.elapsed(step) for step in self._end_times.keys()}
 
     def elapsed(self, step: str) -> float:
         return self._end_times[step] - self._start_times[step]