Merge pull request #811 from NVIDIA/feature/improve_llm_logs

Feature/improve llm logs

Author: schuellc-nvidia

nemoguardrails/actions/v2_x/generation.py

@@ -273,6 +273,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,
@@ -343,6 +347,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,
@@ -504,6 +514,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,
@@ -567,6 +581,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,
@@ -628,6 +644,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,
@@ -743,6 +761,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,
@@ -849,6 +871,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

@@ -27,6 +27,7 @@
 from nemoguardrails.logging.explain import LLMCallInfo
 from nemoguardrails.logging.processing_log import processing_log_var
 from nemoguardrails.logging.stats import LLMStats
+from nemoguardrails.utils import new_uuid
 
 log = logging.getLogger(__name__)
 
@@ -51,13 +52,19 @@ async def on_llm_start(
             llm_call_info = LLMCallInfo()
             llm_call_info_var.set(llm_call_info)
 
+        llm_call_info.id = new_uuid()
+
         # We also add it to the explain object
         explain_info = explain_info_var.get()
         if explain_info:
             explain_info.llm_calls.append(llm_call_info)
 
         log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
-        log.info("Prompt :: %s", prompts[0])
+        log.info(
+            "Prompt :: %s",
+            prompts[0],
+            extra={"id": llm_call_info.id, "task": llm_call_info.task},
+        )
         llm_call_info.prompt = prompts[0]
 
         llm_call_info.started_at = time()
@@ -86,6 +93,8 @@ async def on_chat_model_start(
             llm_call_info = LLMCallInfo()
             llm_call_info_var.set(llm_call_info)
 
+        llm_call_info.id = new_uuid()
+
         # We also add it to the explain object
         explain_info = explain_info_var.get()
         if explain_info:
@@ -109,7 +118,11 @@ async def on_chat_model_start(
         )
 
         log.info("Invocation Params :: %s", kwargs.get("invocation_params", {}))
-        log.info("Prompt Messages :: %s", prompt)
+        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()
 
@@ -143,12 +156,16 @@ async def on_llm_end(
         **kwargs: Any,
     ) -> None:
         """Run when LLM ends running."""
-        log.info("Completion :: %s", response.generations[0][0].text)
         llm_call_info = llm_call_info_var.get()
         if llm_call_info is None:
             llm_call_info = LLMCallInfo()
         llm_call_info.completion = response.generations[0][0].text
         llm_call_info.finished_at = time()
+        log.info(
+            "Completion :: %s",
+            response.generations[0][0].text,
+            extra={"id": llm_call_info.id, "task": llm_call_info.task},
+        )
 
         llm_stats = llm_stats_var.get()
         if llm_stats is None:
@@ -159,7 +176,11 @@ async def on_llm_end(
         if len(response.generations[0]) > 1:
             for i, generation in enumerate(response.generations[0][1:]):
                 log.info("--- :: Completion %d", i + 2)
-                log.info("Completion :: %s", generation.text)
+                log.info(
+                    "Completion :: %s",
+                    generation.text,
+                    extra={"id": llm_call_info.id, "task": llm_call_info.task},
+                )
 
         log.info("Output Stats :: %s", response.llm_output)
         took = llm_call_info.finished_at - llm_call_info.started_at

nemoguardrails/logging/explain.py

@@ -43,6 +43,7 @@ class LLMCallSummary(BaseModel):
 
 
 class LLMCallInfo(LLMCallSummary):
+    id: Optional[str] = Field(default=None, description="The unique prompt identifier.")
     prompt: Optional[str] = Field(
         default=None, description="The prompt that was used for the LLM call."
     )

nemoguardrails/logging/verbose.py

@@ -54,7 +54,7 @@ def emit(self, record) -> None:
                 skip_print = True
                 if verbose_llm_calls:
                     console.print("")
-                    console.print(f"[cyan]LLM {title}[/]")
+                    console.print(f"[cyan]LLM {title} ({record.id[:5]}..)[/]")
                     for line in body.split("\n"):
                         text = Text(line, style="black on #006600", end="\n")
                         text.pad_right(console.width)
@@ -66,6 +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]}..) - {record.task}[/]"
+                    )
 
                     for line in body.split("\n"):
                         if line.strip() == "[/]":