Litellm key update fix (BerriAI#6710)

* fix(caching): convert arg to equivalent kwargs in llm caching handler prevent unexpected errors * fix(caching_handler.py): don't pass args to caching * fix(caching): remove all *args from caching.py * fix(caching): consistent function signatures + abc method * test(caching_unit_tests.py): add unit tests for llm caching ensures coverage for common caching scenarios across different implementations * refactor(litellm_logging.py): move to using cache key from hidden params instead of regenerating one * fix(router.py): drop redis password requirement * fix(proxy_server.py): fix faulty slack alerting check * fix(langfuse.py): avoid copying functions/thread lock objects in metadata fixes metadata copy error when parent otel span in metadata * test: update test * fix(key_management_endpoints.py): fix /key/update with metadata update * fix(key_management_endpoints.py): fix key_prepare_update helper * fix(key_management_endpoints.py): reset value to none if set in key update * fix: update test ' * Litellm dev 11 11 2024 (BerriAI#6693) * fix(__init__.py): add 'watsonx_text' as mapped llm api route Fixes BerriAI#6663 * fix(opentelemetry.py): fix passing parallel tool calls to otel Fixes BerriAI#6677 * refactor(test_opentelemetry_unit_tests.py): create a base set of unit tests for all logging integrations - test for parallel tool call handling reduces bugs in repo * fix(__init__.py): update provider-model mapping to include all known provider-model mappings Fixes BerriAI#6669 * feat(anthropic): support passing document in llm api call * docs(anthropic.md): add pdf anthropic call to docs + expose new 'supports_pdf_input' function * fix(factory.py): fix linting error * add clear doc string for GCS bucket logging * Add docs to export logs to Laminar (BerriAI#6674) * Add docs to export logs to Laminar * minor fix: newline at end of file * place laminar after http and grpc * (Feat) Add langsmith key based logging (BerriAI#6682) * add langsmith_api_key to StandardCallbackDynamicParams * create a file for langsmith types * langsmith add key / team based logging * add key based logging for langsmith * fix langsmith key based logging * fix linting langsmith * remove NOQA violation * add unit test coverage for all helpers in test langsmith * test_langsmith_key_based_logging * docs langsmith key based logging * run langsmith tests in logging callback tests * fix logging testing * test_langsmith_key_based_logging * test_add_callback_via_key_litellm_pre_call_utils_langsmith * add debug statement langsmith key based logging * test_langsmith_key_based_logging * (fix) OpenAI's optional messages[].name does not work with Mistral API (BerriAI#6701) * use helper for _transform_messages mistral * add test_message_with_name to base LLMChat test * fix linting * add xAI on Admin UI (BerriAI#6680) * (docs) add benchmarks on 1K RPS (BerriAI#6704) * docs litellm proxy benchmarks * docs GCS bucket * doc fix - reduce clutter on logging doc title * (feat) add cost tracking stable diffusion 3 on Bedrock (BerriAI#6676) * add cost tracking for sd3 * test_image_generation_bedrock * fix get model info for image cost * add cost_calculator for stability 1 models * add unit testing for bedrock image cost calc * test_cost_calculator_with_no_optional_params * add test_cost_calculator_basic * correctly allow size Optional * fix cost_calculator * sd3 unit tests cost calc * fix raise correct error 404 when /key/info is called on non-existent key (BerriAI#6653) * fix raise correct error on /key/info * add not_found_error error * fix key not found in DB error * use 1 helper for checking token hash * fix error code on key info * fix test key gen prisma * test_generate_and_call_key_info * test fix test_call_with_valid_model_using_all_models * fix key info tests * bump: version 1.52.4 → 1.52.5 * add defaults used for GCS logging * LiteLLM Minor Fixes & Improvements (11/12/2024) (BerriAI#6705) * fix(caching): convert arg to equivalent kwargs in llm caching handler prevent unexpected errors * fix(caching_handler.py): don't pass args to caching * fix(caching): remove all *args from caching.py * fix(caching): consistent function signatures + abc method * test(caching_unit_tests.py): add unit tests for llm caching ensures coverage for common caching scenarios across different implementations * refactor(litellm_logging.py): move to using cache key from hidden params instead of regenerating one * fix(router.py): drop redis password requirement * fix(proxy_server.py): fix faulty slack alerting check * fix(langfuse.py): avoid copying functions/thread lock objects in metadata fixes metadata copy error when parent otel span in metadata * test: update test * bump: version 1.52.5 → 1.52.6 * (feat) helm hook to sync db schema (BerriAI#6715) * v0 migration job * fix job * fix migrations job.yml * handle standalone DB on helm hook * fix argo cd annotations * fix db migration helm hook * fix migration job * doc fix Using Http/2 with Hypercorn * (fix proxy redis) Add redis sentinel support (BerriAI#6154) * add sentinel_password support * add doc for setting redis sentinel password * fix redis sentinel - use sentinel password * Fix: Update gpt-4o costs to that of gpt-4o-2024-08-06 (BerriAI#6714) Fixes BerriAI#6713 * (fix) using Anthropic `response_format={"type": "json_object"}` (BerriAI#6721) * add support for response_format=json anthropic * add test_json_response_format to baseLLM ChatTest * fix test_litellm_anthropic_prompt_caching_tools * fix test_anthropic_function_call_with_no_schema * test test_create_json_tool_call_for_response_format * (feat) Add cost tracking for Azure Dall-e-3 Image Generation + use base class to ensure basic image generation tests pass (BerriAI#6716) * add BaseImageGenTest * use 1 class for unit testing * add debugging to BaseImageGenTest * TestAzureOpenAIDalle3 * fix response_cost_calculator * test_basic_image_generation * fix img gen basic test * fix _select_model_name_for_cost_calc * fix test_aimage_generation_bedrock_with_optional_params * fix undo changes cost tracking * fix response_cost_calculator * fix test_cost_azure_gpt_35 * fix remove dup test (BerriAI#6718) * (build) update db helm hook * (build) helm db pre sync hook * (build) helm db sync hook * test: run test_team_logging firdst --------- Co-authored-by: Ishaan Jaff <ishaanjaffer0324@gmail.com> Co-authored-by: Dinmukhamed Mailibay <47117969+dinmukhamedm@users.noreply.github.com> Co-authored-by: Kilian Lieret <kilian.lieret@posteo.de>
CamdenClark · Nov 13, 2024 · 1c3dcd4 · 1c3dcd4
1 parent 70c8be5
commit 1c3dcd4
Show file tree

Hide file tree

Showing 5 changed files with 46 additions and 19 deletions.
diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml
@@ -68,7 +68,7 @@ model_list:
 
 litellm_settings:
   fallbacks: [{ "claude-3-5-sonnet-20240620": ["claude-3-5-sonnet-aihubmix"] }]
-  callbacks: ["otel", "prometheus"]
+  # callbacks: ["otel", "prometheus"]
   default_redis_batch_cache_expiry: 10
   # default_team_settings: 
   #   - team_id: "dbe2f686-a686-4896-864a-4c3924458709"

diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py
@@ -303,21 +303,17 @@ async def generate_key_fn(  # noqa: PLR0915
         )
 
 
-async def prepare_key_update_data(
+def prepare_key_update_data(
     data: Union[UpdateKeyRequest, RegenerateKeyRequest], existing_key_row
 ):
-    data_json: dict = data.dict(exclude_unset=True)
+    data_json: dict = data.model_dump(exclude_unset=True)
     data_json.pop("key", None)
     _metadata_fields = ["model_rpm_limit", "model_tpm_limit", "guardrails"]
     non_default_values = {}
     for k, v in data_json.items():
         if k in _metadata_fields:
             continue
-        if v is not None:
-            if not isinstance(v, bool) and v in ([], {}, 0):
-                pass
-            else:
-                non_default_values[k] = v
+        non_default_values[k] = v
 
     if "duration" in non_default_values:
         duration = non_default_values.pop("duration")
@@ -379,7 +375,7 @@ async def update_key_fn(
     )
 
     try:
-        data_json: dict = data.json()
+        data_json: dict = data.model_dump(exclude_unset=True)
         key = data_json.pop("key")
         # get the row from db
         if prisma_client is None:
@@ -395,7 +391,7 @@ async def update_key_fn(
                 detail={"error": f"Team not found, passed team_id={data.team_id}"},
             )
 
-        non_default_values = await prepare_key_update_data(
+        non_default_values = prepare_key_update_data(
             data=data, existing_key_row=existing_key_row
         )
 
@@ -1144,7 +1140,7 @@ async def regenerate_key_fn(
     non_default_values = {}
     if data is not None:
         # Update with any provided parameters from GenerateKeyRequest
-        non_default_values = await prepare_key_update_data(
+        non_default_values = prepare_key_update_data(
             data=data, existing_key_row=_key_in_db
         )
 

diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py
@@ -510,3 +510,23 @@ class ReturnValue(BaseModel):
                 "success_callback": "langfuse",
             }
         }
+
+
+def test_prepare_key_update_data():
+    from litellm.proxy.management_endpoints.key_management_endpoints import (
+        prepare_key_update_data,
+    )
+    from litellm.proxy._types import UpdateKeyRequest
+
+    existing_key_row = MagicMock()
+    data = UpdateKeyRequest(key="test_key", models=["gpt-4"], duration="120s")
+    updated_data = prepare_key_update_data(data, existing_key_row)
+    assert "expires" in updated_data
+
+    data = UpdateKeyRequest(key="test_key", metadata={})
+    updated_data = prepare_key_update_data(data, existing_key_row)
+    assert updated_data["metadata"] == {}
+
+    data = UpdateKeyRequest(key="test_key", metadata=None)
+    updated_data = prepare_key_update_data(data, existing_key_row)
+    assert updated_data["metadata"] == None
diff --git a/tests/test_keys.py b/tests/test_keys.py
@@ -66,6 +66,7 @@ async def generate_key(
     max_parallel_requests: Optional[int] = None,
     user_id: Optional[str] = None,
     team_id: Optional[str] = None,
+    metadata: Optional[dict] = None,
     calling_key="sk-1234",
 ):
     url = "http://0.0.0.0:4000/key/generate"
@@ -82,6 +83,7 @@ async def generate_key(
         "max_parallel_requests": max_parallel_requests,
         "user_id": user_id,
         "team_id": team_id,
+        "metadata": metadata,
     }
 
     print(f"data: {data}")
@@ -136,16 +138,21 @@ async def test_key_gen_bad_key():
             pass
 
 
-async def update_key(session, get_key):
+async def update_key(session, get_key, metadata: Optional[dict] = None):
     """
     Make sure only models user has access to are returned
     """
     url = "http://0.0.0.0:4000/key/update"
     headers = {
-        "Authorization": f"Bearer sk-1234",
+        "Authorization": "Bearer sk-1234",
         "Content-Type": "application/json",
     }
-    data = {"key": get_key, "models": ["gpt-4"], "duration": "120s"}
+    data = {"key": get_key}
+
+    if metadata is not None:
+        data["metadata"] = metadata
+    else:
+        data.update({"models": ["gpt-4"], "duration": "120s"})
 
     async with session.post(url, headers=headers, json=data) as response:
         status = response.status
@@ -276,20 +283,24 @@ async def chat_completion_streaming(session, key, model="gpt-4"):
     return prompt_tokens, completion_tokens
 
 
+@pytest.mark.parametrize("metadata", [{"test": "new"}, {}])
 @pytest.mark.asyncio
-async def test_key_update():
+async def test_key_update(metadata):
     """
     Create key
     Update key with new model
     Test key w/ model
     """
     async with aiohttp.ClientSession() as session:
-        key_gen = await generate_key(session=session, i=0)
+        key_gen = await generate_key(session=session, i=0, metadata={"test": "test"})
         key = key_gen["key"]
-        await update_key(
+        assert key_gen["metadata"]["test"] == "test"
+        updated_key = await update_key(
             session=session,
             get_key=key,
+            metadata=metadata,
         )
+        assert updated_key["metadata"] == metadata
         await update_proxy_budget(session=session)  # resets proxy spend
         await chat_completion(session=session, key=key)
 

diff --git a/tests/test_team_logging.py b/tests/test_team_logging.py
@@ -62,8 +62,8 @@ async def chat_completion(session, key, model="azure-gpt-3.5", request_metadata=
 
 
 @pytest.mark.asyncio
-@pytest.mark.flaky(retries=6, delay=1)
-async def test_team_logging():
+@pytest.mark.flaky(retries=12, delay=2)
+async def test_aaateam_logging():
     """
     -> Team 1 logs to project 1
     -> Create Key