Add prompt task name to logging messages

Author: schuellc-nvidia

nemoguardrails/actions/v2_x/generation.py

@@ -48,11 +48,13 @@
     get_element_from_head,
     get_event_from_element,
 )
+from nemoguardrails.context import llm_call_info_var
 from nemoguardrails.embeddings.index import EmbeddingsIndex, IndexItem
 from nemoguardrails.llm.filters import colang
 from nemoguardrails.llm.params import llm_params
 from nemoguardrails.llm.types import Task
 from nemoguardrails.logging import verbose
+from nemoguardrails.logging.explain import LLMCallInfo
 from nemoguardrails.utils import console, new_uuid
 
 log = logging.getLogger(__name__)
@@ -265,6 +267,10 @@ async def generate_user_intent(
         if is_embedding_only:
             return f"{potential_user_intents[0]}"
 
+        llm_call_info_var.set(
+            LLMCallInfo(task=Task.GENERATE_USER_INTENT_FROM_USER_ACTION.value)
+        )
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_USER_INTENT_FROM_USER_ACTION,
             events=events,
@@ -335,6 +341,12 @@ async def generate_user_intent_and_bot_action(
             state, user_action, max_example_flows
         )
 
+        llm_call_info_var.set(
+            LLMCallInfo(
+                task=Task.GENERATE_USER_INTENT_AND_BOT_ACTION_FROM_USER_ACTION.value
+            )
+        )
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_USER_INTENT_AND_BOT_ACTION_FROM_USER_ACTION,
             events=events,
@@ -448,6 +460,10 @@ async def generate_flow_from_instructions(
         flow_id = new_uuid()[0:4]
         flow_name = f"dynamic_{flow_id}"
 
+        llm_call_info_var.set(
+            LLMCallInfo(task=Task.GENERATE_FLOW_FROM_INSTRUCTIONS.value)
+        )
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_FLOW_FROM_INSTRUCTIONS,
             events=events,
@@ -511,6 +527,8 @@ async def generate_flow_from_name(
         for result in reversed(results):
             examples += f"{result.meta['flow']}\n"
 
+        llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_FLOW_FROM_NAME.value))
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_FLOW_FROM_NAME,
             events=events,
@@ -572,6 +590,8 @@ async def generate_flow_continuation(
 
         # TODO: add examples from the actual running flows
 
+        llm_call_info_var.set(LLMCallInfo(task=Task.GENERATE_FLOW_CONTINUATION.value))
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_FLOW_CONTINUATION,
             events=events,
@@ -687,6 +707,10 @@ async def generate_value(
                 if "GenerateValueAction" not in result.text:
                     examples += f"{result.text}\n\n"
 
+        llm_call_info_var.set(
+            LLMCallInfo(task=Task.GENERATE_VALUE_FROM_INSTRUCTION.value)
+        )
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_VALUE_FROM_INSTRUCTION,
             events=events,
@@ -793,6 +817,10 @@ async def generate_flow(
             textwrap.dedent(docstring), context=render_context, events=events
         )
 
+        llm_call_info_var.set(
+            LLMCallInfo(task=Task.GENERATE_FLOW_CONTINUATION_FROM_NLD.value)
+        )
+
         prompt = self.llm_task_manager.render_task_prompt(
             task=Task.GENERATE_FLOW_CONTINUATION_FROM_NLD,
             events=events,

nemoguardrails/logging/callbacks.py

@@ -114,7 +114,11 @@ async def on_chat_model_start(
         )
 
         log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
-        log.info("Prompt Messages :: %s", prompt, extra={"id": llm_call_info.id})
+        log.info(
+            "Prompt Messages :: %s",
+            prompt,
+            extra={"id": llm_call_info.id, "task": llm_call_info.task},
+        )
         llm_call_info.prompt = prompt
         llm_call_info.started_at = time()
 
@@ -156,7 +160,7 @@ async def on_llm_end(
         log.info(
             "Completion :: %s",
             response.generations[0][0].text,
-            extra={"id": llm_call_info.id},
+            extra={"id": llm_call_info.id, "task": llm_call_info.task},
         )
 
         llm_stats = llm_stats_var.get()
@@ -169,7 +173,9 @@ async def on_llm_end(
             for i, generation in enumerate(response.generations[0][1:]):
                 log.info("--- :: Completion %d", i + 2)
                 log.info(
-                    "Completion :: %s", generation.text, extra={"id": llm_call_info.id}
+                    "Completion :: %s",
+                    generation.text,
+                    extra={"id": llm_call_info.id, "task": llm_call_info.task},
                 )
 
         log.info("Output Stats :: %s", response.llm_output)

nemoguardrails/logging/verbose.py

@@ -66,7 +66,9 @@ def emit(self, record) -> None:
                 if verbose_llm_calls:
                     skip_print = True
                     console.print("")
-                    console.print(f"[cyan]LLM Prompt ({record.id[:5]}..)[/]")
+                    console.print(
+                        f"[cyan]LLM Prompt ({record.id[:5]}..) - {record.task}[/]"
+                    )
 
                     for line in body.split("\n"):
                         if line.strip() == "[/]":