Fix final few type errors and add colang directory to pyright pre-commits

Author: tgasser-nv

nemoguardrails/colang/v1_0/runtime/runtime.py

@@ -457,16 +457,18 @@ async def _run_input_rails_in_parallel(
         self, flows: List[str], events: List[dict]
     ) -> ActionResult:
         """Run the input rails in parallel."""
+
         pre_events = [
-            (await create_event({"_type": "StartInputRail", "flow_id": flow})).events[0]
+            (await create_event({"_type": "StartInputRail", "flow_id": flow})).events
             for flow in flows
         ]
+        pre_events = [event[0] for event in pre_events if event]
+
         post_events = [
-            (
-                await create_event({"_type": "InputRailFinished", "flow_id": flow})
-            ).events[0]
+            (await create_event({"_type": "InputRailFinished", "flow_id": flow})).events
             for flow in flows
         ]
+        post_events = [event[0] for event in post_events if event]
 
         return await self._run_flows_in_parallel(
             flows=flows, events=events, pre_events=pre_events, post_events=post_events
@@ -476,18 +478,20 @@ async def _run_output_rails_in_parallel(
         self, flows: List[str], events: List[dict]
     ) -> ActionResult:
         """Run the output rails in parallel."""
+
         pre_events = [
-            (await create_event({"_type": "StartOutputRail", "flow_id": flow})).events[
-                0
-            ]
+            (await create_event({"_type": "StartOutputRail", "flow_id": flow})).events
             for flow in flows
         ]
+        pre_events = [event[0] for event in pre_events if event]
+
         post_events = [
             (
                 await create_event({"_type": "OutputRailFinished", "flow_id": flow})
-            ).events[0]
+            ).events
             for flow in flows
         ]
+        post_events = [event[0] for event in post_events if event]
 
         return await self._run_flows_in_parallel(
             flows=flows, events=events, pre_events=pre_events, post_events=post_events

nemoguardrails/colang/v2_x/runtime/runtime.py

@@ -294,10 +294,10 @@ async def _process_start_action(
 
     async def _get_action_resp(
         self, action_meta: Dict[str, Any], action_name: str, kwargs: Dict[str, Any]
-    ) -> Tuple[Union[str, Dict[str, Any]], str]:
+    ) -> Tuple[Union[Optional[str], Dict[str, Any]], str]:
         """Interact with actions and get response from action-server and system actions."""
         # default response
-        result: Union[str, Dict[str, Any]] = {}
+        result: Union[Optional[str], Dict[str, Any]] = {}
         status: str = "failed"
         try:
             # Call the Actions Server if it is available.
@@ -309,6 +309,7 @@ async def _get_action_resp(
                 result, status = await self.action_dispatcher.execute_action(
                     action_name, kwargs
                 )
+
             else:
                 url = urljoin(
                     self.config.actions_server_url, "/v1/actions/run"

pyproject.toml

@@ -157,6 +157,7 @@ pyright = "^1.1.405"
 include = [
   "nemoguardrails/rails/**",
   "nemoguardrails/actions/**",
+  "nemoguardrails/colang/**",
   "nemoguardrails/cli/**",
   "nemoguardrails/kb/**",
   "nemoguardrails/logging/**",