Merge pull request #823 from NVIDIA/feature/configurable-event-source-uid

Make `source_uid` of events configurable.

Author: sklinglernv

docs/user_guides/configuration-guide.md

@@ -556,6 +556,14 @@ This temperature will be used for the tasks that require deterministic behavior
 lowest_temperature: 0.1
 ```
 
+### Event Source ID
+
+This ID will be used as the `source_uid` for all events emitted by the Colang runtime. Setting this to something else than the default value (default value is `NeMoGuardrails-Colang-2.x`) is useful if you need to distinguish multiple Colang runtimes in your system (e.g. in a multi-agent scenario).
+
+```yaml
+event_source_uid : colang-agent-1
+```
+
 ### Custom Data
 
 If you need to pass additional configuration data to any custom component for your configuration, you can use the `custom_data` field.
@@ -785,6 +793,7 @@ To enable tracing, set the enabled flag to true under the tracing section in you
 tracing:
   enabled: true
 ```
+
 > **Note**: You must install the necessary dependencies to use tracing adapters.
 
   ```bash

nemoguardrails/colang/v2_x/runtime/statemachine.py

@@ -71,6 +71,7 @@
     InternalEvents,
     State,
 )
+from nemoguardrails.rails.llm.config import RailsConfig
 from nemoguardrails.utils import console, new_event_dict, new_readable_uuid, new_uuid
 
 log = logging.getLogger(__name__)
@@ -1829,7 +1830,7 @@ def _is_done_flow(flow_state: FlowState) -> bool:
 
 
 def _generate_umim_event(state: State, event: Event) -> Dict[str, Any]:
-    umim_event = create_umim_event(event, event.arguments)
+    umim_event = create_umim_event(event, event.arguments, state.rails_config)
     state.outgoing_events.append(umim_event)
     log.info("[bold violet]<- Action[/]: %s", event)
 
@@ -2385,10 +2386,14 @@ def create_internal_event(
     return event
 
 
-def create_umim_event(event: Event, event_args: Dict[str, Any]) -> Dict[str, Any]:
+def create_umim_event(
+    event: Event, event_args: Dict[str, Any], config: Optional[RailsConfig]
+) -> Dict[str, Any]:
     """Returns an outgoing UMIM event for the provided action data"""
     new_event_args = dict(event_args)
-    new_event_args["source_uid"] = "NeMoGuardrails-Colang-2.x"
+    new_event_args.setdefault(
+        "source_uid", config.event_source_uid if config else "NeMoGuardrails-Colang-2.x"
+    )
     if isinstance(event, ActionEvent) and event.action_uid is not None:
         if "action_uid" in new_event_args:
             event.action_uid = new_event_args["action_uid"]

nemoguardrails/rails/llm/config.py

@@ -508,6 +508,7 @@ def _join_config(dest_config: dict, additional_config: dict):
         "lowest_temperature",
         "enable_multi_step_generation",
         "colang_version",
+        "event_source_uid",
         "custom_data",
         "prompting_mode",
         "knowledge_base",
@@ -863,6 +864,11 @@ class RailsConfig(BaseModel):
         "This means it will not be altered in any way. ",
     )
 
+    event_source_uid: str = Field(
+        default="NeMoGuardrails-Colang-2.x",
+        description="The source ID of events sent by the Colang Runtime. Useful to identify the component that has sent an event.",
+    )
+
     tracing: TracingConfig = Field(
         default_factory=TracingConfig,
         description="Configuration for tracing.",

tests/utils.py

@@ -319,7 +319,7 @@ def is_data_in_events(
     return True
 
 
-def _init_state(colang_content) -> State:
+def _init_state(colang_content, yaml_content: Optional[str] = None) -> State:
     config = create_flow_configs_from_flow_list(
         parse_colang_file(
             filename="",
@@ -329,8 +329,11 @@ def _init_state(colang_content) -> State:
         )["flows"]
     )
 
+    rails_config = None
+    if yaml_content:
+        rails_config = RailsConfig.from_content(colang_content, yaml_content)
     json.dump(config, sys.stdout, indent=4, cls=EnhancedJsonEncoder)
-    state = State(flow_states=[], flow_configs=config)
+    state = State(flow_states=[], flow_configs=config, rails_config=rails_config)
     initialize_state(state)
     print("---------------------------------")
     json.dump(state.flow_configs, sys.stdout, indent=4, cls=EnhancedJsonEncoder)

tests/v2_x/test_event_mechanics.py

@@ -81,6 +81,32 @@ def test_send_umim_action_event_overwriting_default_parameters():
     )
 
 
+def test_change_umim_event_source_id():
+    """Test to send an UMIM event."""
+
+    content = """
+    flow main
+      send StartUtteranceBotAction(script="Hello world")
+    """
+
+    config = """
+    colang_version: "2.x"
+    event_source_uid : agent-1
+    """
+
+    state = run_to_completion(_init_state(content, config), start_main_flow_event)
+    assert is_data_in_events(
+        state.outgoing_events,
+        [
+            {
+                "type": "StartUtteranceBotAction",
+                "script": "Hello world",
+                "source_uid": "agent-1",
+            }
+        ],
+    )
+
+
 def test_match_umim_action_event():
     """Test to match an UMIM event."""