move activity registration

Author: marcusmotill

src/google/adk/integrations/temporal/__init__.py

@@ -139,6 +139,51 @@ async def wrapper(*args, **kw):
 
         return wrapper
 
+    async def before_model_callback(
+        self, *, callback_context: CallbackContext, llm_request: LlmRequest
+    ) -> LlmResponse | None:
+        # Construct dynamic activity name for visibility
+        agent_name = callback_context.agent_name
+        activity_name = f"{agent_name}.generate_content"
+
+        # Execute with dynamic name
+        response_dicts = await workflow.execute_activity(
+            activity_name,
+            args=[llm_request],
+            **self.activity_options
+        )
+        
+        # Rehydrate LlmResponse objects safely
+        responses = []
+        for d in response_dicts:
+            try:
+                responses.append(LlmResponse.model_validate(d))
+            except Exception as e:
+                raise RuntimeError(f"Failed to deserialized LlmResponse from activity result: {e}") from e
+
+        # Simple consolidation: return the last complete response
+        return responses[-1] if responses else None
+
+
+
+
+class AdkWorkerPlugin(SimplePlugin):
+    """A Temporal Worker Plugin configured for ADK.
+    
+    This plugin configures:
+    1. Pydantic Payload Converter (required for ADK objects).
+    2. Sandbox Passthrough for `google.adk` and `google.genai`.
+    """
+
+    def __init__(self):
+        super().__init__(
+            name="adk_worker_plugin",
+            data_converter=self._configure_data_converter,
+            workflow_runner=self._configure_workflow_runner,
+            activities=[self.dynamic_activity],
+            worker_interceptors=[AdkInterceptor()]
+        )
+
     @staticmethod
     @activity.defn(dynamic=True)
     async def dynamic_activity(args: Sequence[RawValue]) -> Any:
@@ -147,8 +192,8 @@ async def dynamic_activity(args: Sequence[RawValue]) -> Any:
 
         # Check if this is a generate_content call
         if activity_type.endswith(".generate_content") or activity_type == "google.adk.generate_content":
-             return await TemporalPlugin._handle_generate_content(args)
-             
+            return await AdkWorkerPlugin._handle_generate_content(args)
+                
         raise ValueError(f"Unknown dynamic activity: {activity_type}")
 
     @staticmethod
@@ -191,50 +236,6 @@ async def _handle_generate_content(args: List[Any]) -> list[dict[str, Any]]:
             for r in responses
         ]
 
-
-    async def before_model_callback(
-        self, *, callback_context: CallbackContext, llm_request: LlmRequest
-    ) -> LlmResponse | None:
-        # Construct dynamic activity name for visibility
-        agent_name = callback_context.agent_name
-        activity_name = f"{agent_name}.generate_content"
-
-        # Execute with dynamic name
-        response_dicts = await workflow.execute_activity(
-            activity_name,
-            args=[llm_request],
-            **self.activity_options
-        )
-        
-        # Rehydrate LlmResponse objects safely
-        responses = []
-        for d in response_dicts:
-            try:
-                responses.append(LlmResponse.model_validate(d))
-            except Exception as e:
-                raise RuntimeError(f"Failed to deserialized LlmResponse from activity result: {e}") from e
-
-        # Simple consolidation: return the last complete response
-        return responses[-1] if responses else None
-
-
-
-
-class AdkWorkerPlugin(SimplePlugin):
-    """A Temporal Worker Plugin configured for ADK.
-    
-    This plugin configures:
-    1. Pydantic Payload Converter (required for ADK objects).
-    2. Sandbox Passthrough for `google.adk` and `google.genai`.
-    """
-    def __init__(self):
-        super().__init__(
-            name="adk_worker_plugin",
-            data_converter=self._configure_data_converter,
-            workflow_runner=self._configure_workflow_runner,
-            worker_interceptors=[AdkInterceptor()]
-        )
-
     def _configure_data_converter(self, converter: DataConverter | None) -> DataConverter:
         if converter is None:
             return DataConverter(

tests/integration/manual_test_temporal_integration.py

@@ -154,7 +154,7 @@ async def run(self, topic: str) -> str:
             agent=coordinator, 
             app_name="multi_agent_app",
             session_service=session_service,
-            plugins=[TemporalPlugin()]
+            plugins=[TemporalPlugin(activity_options={'start_to_close_timeout': timedelta(minutes=2)})]
         )
 
         # 4. Run
@@ -196,12 +196,9 @@ async def test_temporal_integration():
         task_queue="adk-task-queue",
         activities=[
             get_weather, 
-            TemporalPlugin.dynamic_activity, 
-            # Note: We can add specific wrapper activities if we want distinct names in UI,
-            # but generate_content is the generic one used by TemporalPlugin.
         ],
         workflows=[WeatherAgent, MultiAgentWorkflow],
-        plugins=[AdkWorkerPlugin()] # <--- Use the class based plugin
+        plugins=[AdkWorkerPlugin()]
     ):
         print("Worker started.")
         # Test Weather Agent

tests/unittests/integrations/test_temporal.py

@@ -143,7 +143,7 @@ def test_temporal_plugin_before_model(self):
         callback_context.agent_name = "test-agent"
         callback_context.invocation_context.agent.model = "test-agent-model"
 
-        llm_request = LlmRequest(model=None, prompt="hi")
+        llm_request = LlmRequest(model="test-agent-model", prompt="hi")
 
         # Run callback
         loop = asyncio.new_event_loop()