diff --git a/lavague-core/lavague/core/navigation.py b/lavague-core/lavague/core/navigation.py index dade236f..69c58e2c 100644 --- a/lavague-core/lavague/core/navigation.py +++ b/lavague-core/lavague/core/navigation.py @@ -158,10 +158,19 @@ def get_nodes(self, query: str) -> List[str]: html = self.driver.get_html() - with time_profiler("Retriever Inference", html_size=len(html)): + retrieved_nodes_size = {"size": 0} + + with time_profiler( + "Retriever Inference", + html_size=len(html), + retrieved_nodes_size=retrieved_nodes_size, + ): source_nodes = self.retriever.retrieve( QueryBundle(query_str=query), [html], viewport_only ) + + retrieved_nodes_size["size"] = sum(len(node) for node in source_nodes) + return source_nodes def add_knowledge(self, knowledge: str): diff --git a/lavague-core/lavague/core/utilities/profiling.py b/lavague-core/lavague/core/utilities/profiling.py index f3d84d8e..6d154df0 100644 --- a/lavague-core/lavague/core/utilities/profiling.py +++ b/lavague-core/lavague/core/utilities/profiling.py @@ -28,7 +28,11 @@ def clear_profiling_data(): @contextmanager def time_profiler( - event_name, prompt_size=None, html_size=None, full_step_profiling=False + event_name, + prompt_size=None, + html_size=None, + retrieved_nodes_size=None, + full_step_profiling=False, ): """ A context manager to profile the execution time of code blocks. @@ -53,6 +57,11 @@ def time_profiler( "duration": duration, **({"prompt_size": prompt_size} if prompt_size is not None else {}), **({"html_size": html_size} if html_size is not None else {}), + **( + {"retrieved_nodes_size": retrieved_nodes_size["size"]} + if retrieved_nodes_size is not None + else {} + ), } # append the record to the appropriate list