Add unique identifier to LLM prompt/response logging message

Author: schuellc-nvidia

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,6 +52,8 @@ 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:
@@ -86,6 +89,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 +114,7 @@ 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})
         llm_call_info.prompt = prompt
         llm_call_info.started_at = time()
 
@@ -143,12 +148,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},
+        )
 
         llm_stats = llm_stats_var.get()
         if llm_stats is None:
@@ -159,7 +168,9 @@ 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}
+                )
 
         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,7 @@ def emit(self, record) -> None:
                 if verbose_llm_calls:
                     skip_print = True
                     console.print("")
+                    console.print(f"[cyan]LLM Prompt ({record.id[:5]}..)[/]")
 
                     for line in body.split("\n"):
                         if line.strip() == "[/]":