Update tests

Author: JashG

tests/test_actions_llm_utils.py

@@ -13,11 +13,12 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import pytest
 from typing import cast
+from unittest.mock import AsyncMock
+
+import pytest
 from langchain_core.language_models import BaseLanguageModel
 from langchain_core.messages import AIMessage
-from unittest.mock import AsyncMock
 
 from nemoguardrails.actions.llm.utils import (
     _extract_reasoning_from_additional_kwargs,
@@ -27,6 +28,7 @@
     _infer_provider_from_module,
     _store_reasoning_traces,
     _store_tool_calls,
+    llm_call,
 )
 from nemoguardrails.context import reasoning_trace_var, tool_calls_var
 from nemoguardrails.exceptions import LLMCallException
@@ -581,6 +583,7 @@ def test_store_tool_calls_with_real_aimessage_multiple_tool_calls():
     assert tool_calls[0]["name"] == "foo"
     assert tool_calls[1]["name"] == "bar"
 
+
 @pytest.mark.asyncio
 async def test_llm_call_exception_enrichment_with_model_and_endpoint():
     """Test that LLM invocation errors include model and endpoint context."""
@@ -663,4 +666,4 @@ async def test_llm_call_exception_extracts_nested_client_base_url():
     exc_str = str(exc_info.value)
     assert "https://custom.endpoint.com/v1" in exc_str
     assert "gpt-4-turbo" in exc_str
-    assert "Client error" in exc_str
+    assert "Client error" in exc_str