From 169e8b541c85df5a8e020f86bc45eeb3bb658857 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 8 Aug 2024 12:53:19 -0700 Subject: [PATCH] test(test_langsmith.py): fix test --- litellm/tests/test_langsmith.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/litellm/tests/test_langsmith.py b/litellm/tests/test_langsmith.py index 68182e73d52..f8009942781 100644 --- a/litellm/tests/test_langsmith.py +++ b/litellm/tests/test_langsmith.py @@ -59,9 +59,8 @@ async def test_async_langsmith_logging(): print("fields in logged_run_on_langsmith", logged_run_on_langsmith.keys()) input_fields_on_langsmith = logged_run_on_langsmith.get("inputs") - extra_fields_on_langsmith = logged_run_on_langsmith.get("extra").get( - "invocation_params" - ) + _extra = logged_run_on_langsmith.get("extra", {}) or {} + extra_fields_on_langsmith = _extra.get("invocation_params") print("\nLogged INPUT ON LANGSMITH", input_fields_on_langsmith) @@ -162,9 +161,8 @@ async def test_async_langsmith_logging_with_streaming_and_metadata(sync_mode): input_fields_on_langsmith = logged_run_on_langsmith.get("inputs") - extra_fields_on_langsmith = logged_run_on_langsmith.get("extra").get( - "invocation_params" - ) + _extra = logged_run_on_langsmith.get("extra", {}) or {} + extra_fields_on_langsmith = _extra.get("invocation_params") assert logged_run_on_langsmith.get("run_type") == "llm" print("\nLogged INPUT ON LANGSMITH", input_fields_on_langsmith)