chore(logging): Log model name and base URL before invoking LLMs

Author: JashG

nemoguardrails/actions/llm/utils.py

@@ -13,6 +13,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import logging
 import re
 from typing import Any, Dict, List, Optional, Sequence, Union
 
@@ -33,6 +34,8 @@
 from nemoguardrails.logging.callbacks import logging_callbacks
 from nemoguardrails.logging.explain import LLMCallInfo
 
+log = logging.getLogger(__name__)
+
 
 class LLMCallException(Exception):
     """A wrapper around the LLM call invocation exception.
@@ -46,7 +49,7 @@ def __init__(self, inner_exception: Any):
         self.inner_exception = inner_exception
 
 
-def _infer_model_name(llm: BaseLanguageModel):
+def _infer_model_name(llm: Union[BaseLanguageModel, Runnable]) -> str:
     """Helper to infer the model name based from an LLM instance.
 
     Because not all models implement correctly _identifying_params from LangChain, we have to
@@ -142,13 +145,45 @@ def _prepare_callbacks(
     return logging_callbacks
 
 
+def _log_model_and_base_url(llm: Union[BaseLanguageModel, Runnable]) -> None:
+    """Extract and log the model and base URL from an LLM instance."""
+    model_name = _infer_model_name(llm)
+    base_url = None
+
+    # If llm is a `ChatNIM` instance, we expect its `client` to be an `OpenAI` client with a `base_url` attribute.
+    if hasattr(llm, "client"):
+        client = getattr(llm, "client")
+        if hasattr(client, "base_url"):
+            base_url = str(client.base_url)
+    else:
+        # If llm is a `ChatNVIDIA` instance or other provider, check common attribute names that store the base URL.
+        for attr in [
+            "base_url",
+            "openai_api_base",
+            "azure_endpoint",
+            "api_base",
+            "endpoint",
+        ]:
+            if hasattr(llm, attr):
+                value = getattr(llm, attr, None)
+                if value:
+                    base_url = str(value)
+                    break
+
+    if base_url:
+        log.info(f"Invoking LLM: model={model_name}, url={base_url}")
+    else:
+        log.info(f"Invoking LLM: model={model_name}")
+
+
 async def _invoke_with_string_prompt(
     llm: Union[BaseLanguageModel, Runnable],
     prompt: str,
     callbacks: BaseCallbackManager,
 ):
     """Invoke LLM with string prompt."""
     try:
+        _log_model_and_base_url(llm)
         return await llm.ainvoke(prompt, config=RunnableConfig(callbacks=callbacks))
     except Exception as e:
         raise LLMCallException(e)
@@ -163,6 +198,7 @@ async def _invoke_with_message_list(
     messages = _convert_messages_to_langchain_format(prompt)
 
     try:
+        _log_model_and_base_url(llm)
         return await llm.ainvoke(messages, config=RunnableConfig(callbacks=callbacks))
     except Exception as e:
         raise LLMCallException(e)