From c8ae6cbd3758efd0c7291b8d596bbff0fe027d36 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 9575c18da7..c4d0e84f53 100644 --- a/litellm/tests/test_langsmith.py +++ b/litellm/tests/test_langsmith.py @@ -60,9 +60,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) @@ -165,9 +164,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)