chore: force flush OTel logs at the end of each request

Similar to what is already done for spans. Logs and spans are flushed in concurrently.

PiperOrigin-RevId: 826620266

Author: vertex-sdk-bot

vertexai/agent_engines/templates/adk.py

@@ -25,6 +25,7 @@
 )
 
 import asyncio
+from collections.abc import Awaitable
 import queue
 import threading
 import warnings
@@ -231,9 +232,10 @@ def _warn(msg: str):
     _warn._LOGGER.warning(msg)  # pyright: ignore[reportFunctionMemberAccess]
 
 
-def _force_flush_traces():
+async def _force_flush_otel(tracing_enabled: bool, logging_enabled: bool):
     try:
         import opentelemetry.trace
+        import opentelemetry._logs
     except (ImportError, AttributeError):
         _warn(
             "Could not force flush traces. opentelemetry-api is not installed. Please call  'pip install google-cloud-aiplatform[agent_engines]'."
@@ -242,15 +244,26 @@ def _force_flush_traces():
 
     try:
         import opentelemetry.sdk.trace
+        import opentelemetry.sdk._logs
     except (ImportError, AttributeError):
         _warn(
             "Could not force flush traces. opentelemetry-sdk is not installed. Please call  'pip install google-cloud-aiplatform[agent_engines]'."
         )
         return None
 
-    provider = opentelemetry.trace.get_tracer_provider()
-    if isinstance(provider, opentelemetry.sdk.trace.TracerProvider):
-        _ = provider.force_flush()
+    coros: List[Awaitable[bool]] = []
+
+    if tracing_enabled:
+        tracer_provider = opentelemetry.trace.get_tracer_provider()
+        if isinstance(tracer_provider, opentelemetry.sdk.trace.TracerProvider):
+            coros.append(asyncio.to_thread(tracer_provider.force_flush))
+
+    if logging_enabled:
+        logger_provider = opentelemetry._logs.get_logger_provider()
+        if isinstance(logger_provider, opentelemetry.sdk._logs.LoggerProvider):
+            coros.append(asyncio.to_thread(logger_provider.force_flush))
+
+    await asyncio.gather(*coros, return_exceptions=True)
 
 
 def _default_instrumentor_builder(
@@ -900,9 +913,11 @@ async def async_stream_query(
                 # Yield the event data as a dictionary
                 yield _utils.dump_event_for_json(event)
         finally:
-            # Avoid trace data loss having to do with CPU throttling on instance turndown
-            if self._tracing_enabled():
-                _ = await asyncio.to_thread(_force_flush_traces)
+            # Avoid telemetry data loss having to do with CPU throttling on instance turndown
+            _ = await _force_flush_otel(
+                tracing_enabled=self._tracing_enabled(),
+                logging_enabled=self._telemetry_enabled(),
+            )
 
     def stream_query(
         self,
@@ -1054,9 +1069,11 @@ async def streaming_agent_run_with_events(self, request_json: str):
                     user_id=request.user_id,
                     session_id=session.id,
                 )
-            # Avoid trace data loss having to do with CPU throttling on instance turndown
-            if self._tracing_enabled():
-                _ = await asyncio.to_thread(_force_flush_traces)
+            # Avoid telemetry data loss having to do with CPU throttling on instance turndown
+            _ = await _force_flush_otel(
+                tracing_enabled=self._tracing_enabled(),
+                logging_enabled=self._telemetry_enabled(),
+            )
 
     async def async_get_session(
         self,