This CL enhances asyncio event loop management and test isolation.
- **BigQuery Analytics Plugin:** Ensure the asyncio event loop is consistently closed within the BigQuery analytics plugin. This prevents potential resource leaks. Add checks to handle potential deadlocks in Python 3.13+ when creating loops during interpreter shutdown.
- **Test Thread Pool Cleanup:** Introduce a pytest fixture (`cleanup_thread_pools`) to automatically shut down and clear all tool-related thread pools after each test run in `test_functions_thread_pool.py`. This improves test isolation and prevents order-dependent test failures.
- **Streaming Test Loop Restoration:** Refactor event loop handling in `test_streaming.py`. A new `_run_with_loop` method is introduced in the custom test runners to create a temporary event loop for each test execution, run the coroutine, and crucially, restore the original event loop afterwards. This prevents tests from interfering with each other's loop state.
- **Resource Closure:** Ensure services are closed properly in tests by adding `await service.close()` in `test_service_factory.py` and using `async with session_service` in `test_session_service.py`.
PiperOrigin-RevId: 863305565
Improves test reliability by guaranteeing the BigQueryAgentAnalyticsPlugin is shut down after each test execution. This is achieved by wrapping test logic in try/finally blocks, using a yield fixture, or employing a new asynccontextmanager to ensure the plugin.shutdown() method is always called, preventing potential resource leaks between tests.
PiperOrigin-RevId: 862951816
This change introduces a `flush` method to the `BigQueryAgentAnalyticsPlugin`. This ensures that all pending log events are written to BigQuery before the agent's run completes.
Key changes:
- Added `flush()` method to `BigQueryAgentAnalyticsPlugin` to force write of pending events.
PiperOrigin-RevId: 859263853
This refactors the BigQueryAgentAnalyticsPlugin to use the standard OpenTelemetry API for trace and span ID generation and propagation, replacing the custom ContextVar implementation.
Key changes:
- Utilizes `opentelemetry.trace` for starting/ending spans.
- Correctly uses `opentelemetry.context` for context attachment and detachment.
- Span information is now derived from the OpenTelemetry context when available.
- Added a fallback mechanism to ensure span_id and parent_span_id are still populated if the OpenTelemetry SDK is not initialized.
To get standard OpenTelemetry trace information in BigQuery logs, users should install `opentelemetry-sdk` and initialize a global `TracerProvider` in their application *before* initializing ADK components.
Example minimal initialization:
```python
# Install: pip install opentelemetry-sdk
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
trace.set_tracer_provider(TracerProvider())
```
PiperOrigin-RevId: 858965562
* **Async Safety:** Improved TraceManager context variable handling to ensure correct context isolation in concurrent asynchronous operations. This was achieved by using immutable tuples for the span stack and making copies of context dictionaries before modification.
* **Enhanced Logging:** The BigQueryAgentAnalyticsPlugin now captures richer metadata, including:
* Root agent name (via a new context variable).
* LLM model name and version.
* Usage metadata from LLM requests and responses.
* **Serialization Fix:** Updated BigQueryAgentAnalyticsPlugin to prevent JSON serialization errors when logging custom objects (e.g., Dataclasses). These are now automatically converted to dictionaries or string representations to ensure successful insertion into BigQuery.
PiperOrigin-RevId: 855415320
Fix discovery engine search tool, bigquery agent analytics plugin, and application integration tool to correctly handle the ADC quota project override -- the x-goog-user-project should be set based on the ADC quota project, per gcloud auth team's requirements.
Co-authored-by: Kathy Wu <wukathy@google.com>
PiperOrigin-RevId: 853841124
Performance: Switched to BigQuery Storage Write API with async batching, reducing agent latency.
Multimodal: Native support for GCS offloading (ObjectRef) for images, video, and large text.
Reliability: Added connection pooling, retries, and a "rescue flush" for safe shutdown on Cloud Run.
Observability: Fixed distributed tracing hierarchy with ContextVars support.
PiperOrigin-RevId: 843062561
This update enhances the BigQuery agent analytics plugin:
* **Schema Field Descriptions:** The recommended BigQuery table schema now includes descriptions for each field, improving data understandability.
* **Optimized Table Structure:** The plugin now creates the table with daily partitioning on `timestamp` and clustering on `event_type`, `agent`, and `user_id` by default.
* **Truncation Flag:** A new boolean field `is_truncated` is added to the schema to show if the `content` was truncated.
PiperOrigin-RevId: 832436799
This update improves the `BigQueryAgentAnalyticsPlugin` in several ways:
* Corrects the PyArrow schema generation to accurately reflect BigQuery field nullability based on the `mode` attribute.
* Introduces a configurable `shutdown_timeout` in `BigQueryLoggerConfig` to manage how long the plugin waits for pending logs to flush during shutdown.
* Adds more robust error handling within the `shutdown` method and background write tasks, particularly for event loop closure issues.
* Improves internal logging to provide better diagnostics.
* Ensures consistent use of safe content formatting.
PiperOrigin-RevId: 831225837
This change introduces a shutdown lifecycle hook for plugins. The `PluginManager` now has an `async def shutdown()` method that will call `await plugin.shutdown()` on any registered plugins that implement the method. This is called from `Runner.close()`, allowing plugins to perform cleanup tasks like flushing logs or closing connections when the runner instance is being closed. This improves the reliability of plugins that perform background operations.
PiperOrigin-RevId: 831037737
This change updates the `BigQueryAgentAnalyticsPlugin` (formerly in `bigquery_logging_plugin.py`) to perform BigQuery writes asynchronously in background tasks, preventing blocking of the main agent execution flow. Initialization of BigQuery clients and table creation is also made asynchronous. The content logged in various callbacks has been streamlined and simplified. The plugin file has been renamed to `bigquery_agent_analytics_plugin.py` to match the class name.
PiperOrigin-RevId: 829603260
