Make code async

Author: pratduv

src/galileo/decorator.py

@@ -930,6 +930,24 @@ def start_session(
             name=name, previous_session_id=previous_session_id, external_id=external_id
         )
 
+    async def async_start_session(
+        self, name: Optional[str] = None, previous_session_id: Optional[str] = None, external_id: Optional[str] = None
+    ) -> str:
+        """
+        Async start a session in the active context logger instance.
+
+        Args:
+            name: The name of the session. If not provided, a session name will be generated automatically.
+            previous_session_id: The id of the previous session. Defaults to None.
+            external_id: The external id of the session. Defaults to None.
+
+        Returns:
+            str: The id of the newly created session.
+        """
+        return await self.get_logger_instance().async_start_session(
+            name=name, previous_session_id=previous_session_id, external_id=external_id
+        )
+
     def clear_session(self) -> None:
         """Clear the session in the active context logger instance."""
         self.get_logger_instance().clear_session()

src/galileo/logger/logger.py

@@ -1016,22 +1016,8 @@ def conclude(
 
         return current_parent
 
-    @nop_sync
-    def flush(self) -> list[Trace]:
-        """
-        Upload all traces to Galileo.
-
-        Returns:
-        -------
-            List[Trace]: The list of uploaded traces.
-        """
-        if self.mode == "batch":
-            return self._flush_batch()
-        else:
-            self._logger.warning("Flushing in streaming mode is not supported.")
-            return list()
-
-    def _flush_batch(self):
+    async def _flush_batch(self, is_async: bool = False) -> list[Trace]:
+        # import pdb; pdb.set_trace()
         if not self.traces:
             self._logger.info("No traces to flush.")
             return list()
@@ -1044,7 +1030,7 @@ def _flush_batch(self):
 
         if self.local_metrics:
             self._logger.info("Computing local metrics...")
-            # TODO: parallelize, possibly with ThreadPoolExecutor
+            # TODO: parallelize, possibly with ThreadPoolExecutor/asyncio
             for trace in self.traces:
                 populate_local_metrics(trace, self.local_metrics)
 
@@ -1053,7 +1039,12 @@ def _flush_batch(self):
         traces_ingest_request = TracesIngestRequest(
             traces=self.traces, experiment_id=self.experiment_id, session_id=self.session_id
         )
-        self._client.ingest_traces_sync(traces_ingest_request)
+
+        if is_async:
+            await self._client.ingest_traces(traces_ingest_request)
+        else:
+            self._client.ingest_traces_sync(traces_ingest_request)
+
         logged_traces = self.traces
 
         self._logger.info("Successfully flushed %d traces.", len(logged_traces))
@@ -1072,39 +1063,39 @@ async def async_flush(self) -> list[Trace]:
             List[Trace]: The list of uploaded workflows.
         """
         if self.mode == "batch":
-            return await self._async_flush_batch()
+            return await self._flush_batch(is_async=True)
         else:
             self._logger.warning("Flushing in streaming mode is not supported.")
             return list()
 
-    async def _async_flush_batch(self) -> list[Trace]:
-        if not self.traces:
-            self._logger.info("No traces to flush.")
-            return list()
-
-        current_parent = self.current_parent()
-        if current_parent is not None:
-            self._logger.info("Concluding the active trace...")
-            last_output = get_last_output(current_parent)
-            self.conclude(output=last_output, conclude_all=True)
-
-        if self.local_metrics:
-            self._logger.info("Computing metrics for local scorers...")
-            # TODO: parallelize, possibly with asyncio to_thread/gather
-            for trace in self.traces:
-                populate_local_metrics(trace, self.local_metrics)
-
-        self._logger.info("Flushing %d traces...", len(self.traces))
-
-        traces_ingest_request = TracesIngestRequest(traces=self.traces, session_id=self.session_id)
-        await self._client.ingest_traces(traces_ingest_request)
-        logged_traces = self.traces
-
-        self._logger.info("Successfully flushed %d traces.", len(logged_traces))
+    @nop_sync
+    def flush(self) -> list[Trace]:
+        """
+        Upload all traces to Galileo.
 
-        self.traces = list()
-        self._parent_stack = deque()
-        return logged_traces
+        Returns:
+        -------
+            List[Trace]: The list of uploaded traces.
+        """
+        if self.mode == "batch":
+            # This is bad because asyncio.run() fails in environments with existing event loops
+            # (e.g. jupyter notebooks, FastAPI, etc. would fail with "cannot be called from a running event loop"
+            # Even though flush() is sync, it can be called from async contexts like:
+            # - Jupyter notebooks (which have their own event loop)
+            # - pytest-asyncio tests (where @mark.asyncio creates an event loop)
+            # - FastAPI/Django async views (where the web framework has an event loop)
+            # - Any async function that calls sync code
+            # The EventLoopThreadPool approach works in ALL environments by using dedicated threads
+            # return asyncio.run(self._flush_batch(is_async=False))
+
+            # This is good because async_run() uses EventLoopThreadPool which works in all environments
+            # by running async code in dedicated threads with their own event loops
+            from galileo_core.helpers.execution import async_run
+
+            return async_run(self._flush_batch(is_async=False))
+        else:
+            self._logger.warning("Flushing in streaming mode is not supported.")
+            return list()
 
     @nop_sync
     def terminate(self) -> None:

tests/test_openai_agents.py

@@ -127,7 +127,47 @@ async def test_simple_agent(
         assert span.metrics.duration_ns
         assert span.metrics.duration_ns > 0
 
+    await galileo_logger.async_flush()
+    payload = mock_core_api_instance.ingest_traces.call_args[0][0]
+    assert len(payload.traces) == 1
+    assert len(payload.traces[0].spans) == 1
+
+
+@vcr.use_cassette(
+    "tests/fixtures/openai_agents.yaml",
+    filter_headers=["authorization"],
+    decode_compressed_response=True,
+    record_mode=vcr.mode.NEW_EPISODES,
+)
+@patch("galileo.logger.logger.LogStreams")
+@patch("galileo.logger.logger.Projects")
+@patch("galileo.logger.logger.GalileoCoreApiClient")
+def test_simple_agent_sync_flush(
+    mock_core_api_client: Mock, mock_projects_client: Mock, mock_logstreams_client: Mock, monkeypatch: MonkeyPatch
+) -> None:
+    """Test sync flush() method - this test is NOT async"""
+    mock_core_api_instance = setup_mock_core_api_client(mock_core_api_client)
+    setup_mock_projects_client(mock_projects_client)
+    setup_mock_logstreams_client(mock_logstreams_client)
+
+    galileo_logger = GalileoLogger(project="test", log_stream="test")
+
+    # Add a simple trace manually (no async Runner.run)
+    galileo_logger.start_trace(input="Test input")
+    galileo_logger.add_llm_span(
+        input="Test input",
+        output="Test output",
+        model="gpt-4o",
+        num_input_tokens=5,
+        num_output_tokens=5,
+        total_tokens=10,
+    )
+    galileo_logger.conclude(output="Test output")
+
+    # This should work since we're NOT in an async context
     galileo_logger.flush()
+
+    # Check that sync method was called
     payload = mock_core_api_instance.ingest_traces_sync.call_args[0][0]
     assert len(payload.traces) == 1
     assert len(payload.traces[0].spans) == 1