autoformat.sh

Author: BowedMars2

src/google/adk/runners.py

@@ -67,19 +67,22 @@
 
 logger = logging.getLogger('google_adk.' + __name__)
 
+
 class StopSignal:
+
   def __init__(self):
     self.stopped = False
-  
+
   def stop(self):
     self.stopped = True
 
   def reset(self):
     self.stopped = False
-  
+
   def is_set(self) -> bool:
     return self.stopped
 
+
 def _is_tool_call_or_response(event: Event) -> bool:
   return bool(event.get_function_calls() or event.get_function_responses())
 
@@ -625,7 +628,7 @@ async def execute(ctx: InvocationContext) -> AsyncGenerator[Event]:
                 session=session,
                 execute_fn=execute,
                 is_live_call=False,
-                stop_signal=stop_signal
+                stop_signal=stop_signal,
             )
         ) as agen:
           async for event in agen:
@@ -842,7 +845,7 @@ async def _exec_with_plugin(
       execute_fn: Callable[[InvocationContext], AsyncGenerator[Event, None]],
       is_live_call: bool = False,
       *,
-      stop_signal: Optional[StopSignal] = None
+      stop_signal: Optional[StopSignal] = None,
   ) -> AsyncGenerator[Event, None]:
     """Wraps execution with plugin callbacks.
 
@@ -904,10 +907,7 @@ async def _exec_with_plugin(
                 invocation_id=invocation_context.invocation_id,
                 author='model',
                 interrupted=True,
-                content=types.Content(
-                  role='model',
-                  parts=[]
-                )
+                content=types.Content(role='model', parts=[]),
             )
             yield interrupted_event
 

tests/unittests/runners/test_cancel_async.py

@@ -1,95 +1,96 @@
 import asyncio
-import pytest
-from google.adk.runners import InMemoryRunner
+
+from google.adk.agents import BaseAgent
 from google.adk.agents.run_config import RunConfig
 from google.adk.agents.run_config import StreamingMode
-from google.adk.runners import StopSignal
-from google.genai.types import Part, Content
-from google.adk.agents import BaseAgent
 from google.adk.events import Event
+from google.adk.runners import InMemoryRunner
+from google.adk.runners import StopSignal
+from google.genai.types import Content
+from google.genai.types import Part
+import pytest
+
 
 # Used to tell the main loop when to interrupt the run_async
 class TriggerSignal:
+
   def __init__(self):
     self.triggered = False
-  
+
   def set(self):
     self.triggered = True
 
   def reset(self):
     self.triggered = False
-  
+
   def is_set(self) -> bool:
     return self.triggered
 
+
 # Mocks an LLM, outputting events every 1s
 class MockAgent(BaseAgent):
-    name: str = "MockAgent"
-    
-    async def run_async(self, ctx: Any) -> AsyncGenerator[Event, None]:
-        cycle = 0
-
-        while True:
-            await asyncio.sleep(1)
-            cycle += 1
-            yield Event(
-                author='model',
-                content=Content(
-                    role='model',
-                    parts=[Part(text=f"{cycle}")]
-                )
-            )
+  name: str = "MockAgent"
+
+  async def run_async(self, ctx: Any) -> AsyncGenerator[Event, None]:
+    cycle = 0
+
+    while True:
+      await asyncio.sleep(1)
+      cycle += 1
+      yield Event(
+          author="model",
+          content=Content(role="model", parts=[Part(text=f"{cycle}")]),
+      )
+
 
 APP_NAME = "adk_test_app"
 USER_ID = "adk_test_user"
 STOP_SIGNAL = StopSignal()
 TRIGGER_SIGNAL = TriggerSignal()
 
+
 @pytest.mark.asyncio
 async def test_stop_signal():
-    async def consume_stream(runner, trigger_signal, **kwargs):
-        cycle = 0
-
-        async for event in runner.run_async(**kwargs):
-            cycle += 1
-            content = event.content.parts[0].text
-
-            if not content:
-                # The fourth cycle should be interrupted
-                # Only interrupted cycles would not have content objects given this MockAgent
-                assert cycle == 4
-                assert event.interrupted
-            if content and content == "3":
-                # Tell main loop to interrupt the fourth cycle
-                trigger_signal.set()         
-
-    client = MockAgent()
-    runner = InMemoryRunner(agent=client, app_name=APP_NAME)
-    
-    session = await runner.session_service.create_session(
-        app_name=APP_NAME, user_id=USER_ID
-    )
-
-    message = Content(
-        role="user",
-        parts=[Part(text="Necessary message")]
-    )
-    
-    task = asyncio.create_task(consume_stream(
-        runner=runner,
-        trigger_signal=TRIGGER_SIGNAL,
-        session_id=session.id, 
-        user_id=USER_ID, 
-        new_message=message, 
-        stop_signal=STOP_SIGNAL,
-        run_config=RunConfig(
-            streaming_mode=StreamingMode.SSE
-        )
-    ))
-
-    # Wait for 3 events to pass
-    while not TRIGGER_SIGNAL.is_set():
-        await asyncio.sleep(0.1)
-    
-    # Interrupt the run_async
-    STOP_SIGNAL.stop()
+  async def consume_stream(runner, trigger_signal, **kwargs):
+    cycle = 0
+
+    async for event in runner.run_async(**kwargs):
+      cycle += 1
+      content = event.content.parts[0].text
+
+      if not content:
+        # The fourth cycle should be interrupted
+        # Only interrupted cycles would not have content objects given this MockAgent
+        assert cycle == 4
+        assert event.interrupted
+      if content and content == "3":
+        # Tell main loop to interrupt the fourth cycle
+        trigger_signal.set()
+
+  client = MockAgent()
+  runner = InMemoryRunner(agent=client, app_name=APP_NAME)
+
+  session = await runner.session_service.create_session(
+      app_name=APP_NAME, user_id=USER_ID
+  )
+
+  message = Content(role="user", parts=[Part(text="Necessary message")])
+
+  task = asyncio.create_task(
+      consume_stream(
+          runner=runner,
+          trigger_signal=TRIGGER_SIGNAL,
+          session_id=session.id,
+          user_id=USER_ID,
+          new_message=message,
+          stop_signal=STOP_SIGNAL,
+          run_config=RunConfig(streaming_mode=StreamingMode.SSE),
+      )
+  )
+
+  # Wait for 3 events to pass
+  while not TRIGGER_SIGNAL.is_set():
+    await asyncio.sleep(0.1)
+
+  # Interrupt the run_async
+  STOP_SIGNAL.stop()