Merge branch 'main' into feature/ollama-llm

Author: ayman3000

src/google/adk/plugins/bigquery_agent_analytics_plugin.py

@@ -660,6 +660,9 @@ def get_first_token_time(span_id: str) -> Optional[float]:
 # ==============================================================================
 # HELPER: BATCH PROCESSOR
 # ==============================================================================
+_SHUTDOWN_SENTINEL = object()
+
+
 class BatchProcessor:
   """Handles asynchronous batching and writing of events to BigQuery."""
 
@@ -809,11 +812,18 @@ async def _batch_writer(self) -> None:
               self._queue.get(), timeout=self.flush_interval
           )
 
+        if first_item is _SHUTDOWN_SENTINEL:
+          self._queue.task_done()
+          continue
+
         batch.append(first_item)
 
         while len(batch) < self.batch_size:
           try:
             item = self._queue.get_nowait()
+            if item is _SHUTDOWN_SENTINEL:
+              self._queue.task_done()
+              continue
             batch.append(item)
           except asyncio.QueueEmpty:
             break
@@ -831,6 +841,13 @@ async def _batch_writer(self) -> None:
       except asyncio.CancelledError:
         logger.info("Batch writer task cancelled.")
         break
+      except RuntimeError as e:
+        if "Event loop is closed" in str(e):
+          logger.info("Batch writer loop closed: %s", e)
+          break
+        # Re-raise other RuntimeErrors (or log them below)
+        logger.error("RuntimeError in batch writer loop: %s", e, exc_info=True)
+        await asyncio.sleep(1)
       except Exception as e:
         logger.error("Error in batch writer loop: %s", e, exc_info=True)
         await asyncio.sleep(1)
@@ -939,12 +956,24 @@ async def shutdown(self, timeout: float = 5.0) -> None:
     """
     self._shutdown = True
     logger.info("BatchProcessor shutting down, draining queue...")
+
+    # Signal the writer to wake up and check shutdown status
+    try:
+      self._queue.put_nowait(_SHUTDOWN_SENTINEL)
+    except asyncio.QueueFull:
+      # If queue is full, the writer is active and will check _shutdown soon
+      pass
+
     if self._batch_processor_task:
       try:
         await asyncio.wait_for(self._batch_processor_task, timeout=timeout)
       except asyncio.TimeoutError:
         logger.warning("BatchProcessor shutdown timed out, cancelling worker.")
         self._batch_processor_task.cancel()
+        try:
+          await self._batch_processor_task
+        except asyncio.CancelledError:
+          pass
       except Exception as e:
         logger.error("Error during BatchProcessor shutdown: %s", e)
 
@@ -1626,51 +1655,55 @@ def get_credentials():
   def _atexit_cleanup(batch_processor: "BatchProcessor") -> None:
     """Clean up batch processor on script exit."""
     # Check if the batch_processor object is still alive
-    if batch_processor and not batch_processor._shutdown:
-      # Emergency Flush: Rescue any logs remaining in the queue
-      remaining_items = []
-      try:
-        while True:
-          remaining_items.append(batch_processor._queue.get_nowait())
-      except (asyncio.QueueEmpty, AttributeError):
-        pass
-
-      if remaining_items:
-        # We need a new loop and client to flush these
-        async def rescue_flush():
-          try:
-            # Create a short-lived client just for this flush
+    try:
+      if batch_processor and not batch_processor._shutdown:
+        # Emergency Flush: Rescue any logs remaining in the queue
+        remaining_items = []
+        try:
+          while True:
+            remaining_items.append(batch_processor._queue.get_nowait())
+        except (asyncio.QueueEmpty, AttributeError):
+          pass
+
+        if remaining_items:
+          # We need a new loop and client to flush these
+          async def rescue_flush():
             try:
-              # Note: This relies on google.auth.default() working in this context.
-              # pylint: disable=g-import-not-at-top
-              from google.cloud.bigquery_storage_v1.services.big_query_write.async_client import BigQueryWriteAsyncClient
-
-              # pylint: enable=g-import-not-at-top
-              client = BigQueryWriteAsyncClient()
+              # Create a short-lived client just for this flush
+              try:
+                # Note: This relies on google.auth.default() working in this context.
+                # pylint: disable=g-import-not-at-top
+                from google.cloud.bigquery_storage_v1.services.big_query_write.async_client import BigQueryWriteAsyncClient
+
+                # pylint: enable=g-import-not-at-top
+                client = BigQueryWriteAsyncClient()
+              except Exception as e:
+                logger.warning("Could not create rescue client: %s", e)
+                return
+
+              # Patch batch_processor.write_client temporarily
+              old_client = batch_processor.write_client
+              batch_processor.write_client = client
+              try:
+                # Force a write
+                await batch_processor._write_rows_with_retry(remaining_items)
+                logger.info("Rescued logs flushed successfully.")
+              except Exception as e:
+                logger.error("Failed to flush rescued logs: %s", e)
+              finally:
+                batch_processor.write_client = old_client
             except Exception as e:
-              logger.warning("Could not create rescue client: %s", e)
-              return
+              logger.error("Rescue flush failed: %s", e)
 
-            # Patch batch_processor.write_client temporarily
-            old_client = batch_processor.write_client
-            batch_processor.write_client = client
-            try:
-              # Force a write
-              await batch_processor._write_rows_with_retry(remaining_items)
-              logger.info("Rescued logs flushed successfully.")
-            except Exception as e:
-              logger.error("Failed to flush rescued logs: %s", e)
-            finally:
-              batch_processor.write_client = old_client
+          try:
+            loop = asyncio.new_event_loop()
+            loop.run_until_complete(rescue_flush())
+            loop.close()
           except Exception as e:
-            logger.error("Rescue flush failed: %s", e)
-
-        try:
-          loop = asyncio.new_event_loop()
-          loop.run_until_complete(rescue_flush())
-          loop.close()
-        except Exception as e:
-          logger.error("Failed to run rescue loop: %s", e)
+            logger.error("Failed to run rescue loop: %s", e)
+    except ReferenceError:
+      # batch_processor already GC'd, nothing to do
+      pass
 
   def _ensure_schema_exists(self) -> None:
     """Ensures the BigQuery table exists with the correct schema."""

src/google/adk/tools/openapi_tool/openapi_spec_parser/rest_api_tool.py

@@ -14,6 +14,7 @@
 
 from __future__ import annotations
 
+import logging
 import ssl
 from typing import Any
 from typing import Callable
@@ -48,6 +49,8 @@
 from .operation_parser import OperationParser
 from .tool_auth_handler import ToolAuthHandler
 
+logger = logging.getLogger("google_adk." + __name__)
+
 
 def snake_to_lower_camel(snake_case_string: str):
   """Converts a snake_case string to a lower_camel_case string.
@@ -158,6 +161,7 @@ def __init__(
     self._default_headers: Dict[str, str] = {}
     self._ssl_verify = ssl_verify
     self._header_provider = header_provider
+    self._logger = logger
     if should_parse_operation:
       self._operation_parser = OperationParser(self.operation)
 
@@ -493,14 +497,40 @@ async def call(
       if provider_headers:
         request_params.setdefault("headers", {}).update(provider_headers)
 
+    # Log the API request
+    self._logger.debug(
+        "API Request: %s %s",
+        request_params.get("method", "").upper(),
+        request_params.get("url", ""),
+    )
+    self._logger.debug("API Request params: %s", request_params.get("params"))
+    if "json" in request_params:
+      self._logger.debug("API Request body: %s", request_params.get("json"))
+
     response = requests.request(**request_params)
 
+    # Log the API response
+    self._logger.debug(
+        "API Response: %s %s - Status: %d",
+        request_params.get("method", "").upper(),
+        request_params.get("url", ""),
+        response.status_code,
+    )
+
     # Parse API response
     try:
       response.raise_for_status()  # Raise HTTPError for bad responses
-      return response.json()  # Try to decode JSON
+      result = response.json()  # Try to decode JSON
+      self._logger.debug("API Response body: %s", result)
+      return result
     except requests.exceptions.HTTPError:
       error_details = response.content.decode("utf-8")
+      self._logger.warning(
+          "API call failed for tool %s: Status %d - %s",
+          self.name,
+          response.status_code,
+          error_details,
+      )
       return {
           "error": (
               f"Tool {self.name} execution failed. Analyze this execution error"
@@ -510,6 +540,7 @@ async def call(
           )
       }
     except ValueError:
+      self._logger.debug("API Response (non-JSON): %s", response.text)
       return {"text": response.text}  # Return text if not JSON
 
   def __str__(self):

src/google/adk/tools/vertex_ai_search_tool.py

@@ -21,6 +21,7 @@
 from google.genai import types
 from typing_extensions import override
 
+from ..agents.readonly_context import ReadonlyContext
 from ..utils.model_name_utils import is_gemini_1_model
 from ..utils.model_name_utils import is_gemini_model
 from .base_tool import BaseTool
@@ -38,6 +39,25 @@ class VertexAiSearchTool(BaseTool):
   Attributes:
     data_store_id: The Vertex AI search data store resource ID.
     search_engine_id: The Vertex AI search engine resource ID.
+
+  To dynamically customize the search configuration at runtime (e.g., set
+  filter based on user context), subclass this tool and override the
+  `_build_vertex_ai_search_config` method.
+
+  Example:
+    ```python
+    class DynamicFilterSearchTool(VertexAiSearchTool):
+      def _build_vertex_ai_search_config(
+          self, ctx: ReadonlyContext
+      ) -> types.VertexAISearch:
+        user_id = ctx.state.get('user_id')
+        return types.VertexAISearch(
+            datastore=self.data_store_id,
+            engine=self.search_engine_id,
+            filter=f"user_id = '{user_id}'",
+            max_results=self.max_results,
+        )
+    ```
   """
 
   def __init__(
@@ -90,6 +110,30 @@ def __init__(
     self.max_results = max_results
     self.bypass_multi_tools_limit = bypass_multi_tools_limit
 
+  def _build_vertex_ai_search_config(
+      self, readonly_context: ReadonlyContext
+  ) -> types.VertexAISearch:
+    """Builds the VertexAISearch configuration.
+
+    Override this method in a subclass to dynamically customize the search
+    configuration based on the context (e.g., set filter based on session
+    state).
+
+    Args:
+      readonly_context: The readonly context with access to state and session
+        info.
+
+    Returns:
+      The VertexAISearch configuration to use for this request.
+    """
+    return types.VertexAISearch(
+        datastore=self.data_store_id,
+        data_store_specs=self.data_store_specs,
+        engine=self.search_engine_id,
+        filter=self.filter,
+        max_results=self.max_results,
+    )
+
   @override
   async def process_llm_request(
       self,
@@ -106,14 +150,20 @@ async def process_llm_request(
       llm_request.config = llm_request.config or types.GenerateContentConfig()
       llm_request.config.tools = llm_request.config.tools or []
 
+      # Build the search config (can be overridden by subclasses)
+      vertex_ai_search_config = self._build_vertex_ai_search_config(
+          tool_context
+      )
+
       # Format data_store_specs concisely for logging
-      if self.data_store_specs:
+      if vertex_ai_search_config.data_store_specs:
         spec_ids = [
             spec.data_store.split('/')[-1] if spec.data_store else 'unnamed'
-            for spec in self.data_store_specs
+            for spec in vertex_ai_search_config.data_store_specs
         ]
         specs_info = (
-            f'{len(self.data_store_specs)} spec(s): [{", ".join(spec_ids)}]'
+            f'{len(vertex_ai_search_config.data_store_specs)} spec(s):'
+            f' [{", ".join(spec_ids)}]'
         )
       else:
         specs_info = None
@@ -122,23 +172,17 @@ async def process_llm_request(
           'Adding Vertex AI Search tool config to LLM request: '
           'datastore=%s, engine=%s, filter=%s, max_results=%s, '
           'data_store_specs=%s',
-          self.data_store_id,
-          self.search_engine_id,
-          self.filter,
-          self.max_results,
+          vertex_ai_search_config.datastore,
+          vertex_ai_search_config.engine,
+          vertex_ai_search_config.filter,
+          vertex_ai_search_config.max_results,
           specs_info,
       )
 
       llm_request.config.tools.append(
           types.Tool(
               retrieval=types.Retrieval(
-                  vertex_ai_search=types.VertexAISearch(
-                      datastore=self.data_store_id,
-                      data_store_specs=self.data_store_specs,
-                      engine=self.search_engine_id,
-                      filter=self.filter,
-                      max_results=self.max_results,
-                  )
+                  vertex_ai_search=vertex_ai_search_config
               )
           )
       )