feat(custom_logger.py): expose new async_dataset_hook for modifying… (#6331)

* feat(custom_logger.py): expose new `async_dataset_hook` for modifying/rejecting argilla items before logging

Allows user more control on what gets logged to argilla for annotations

* feat(google_ai_studio_endpoints.py): add new `/azure/*` pass through route

enables pass-through for azure provider

* feat(utils.py): support checking ollama `/api/show` endpoint for retrieving ollama model info

Fixes https://github.com/BerriAI/litellm/issues/6322

* fix(user_api_key_auth.py): add `/key/delete` to an allowed_ui_routes

Fixes https://github.com/BerriAI/litellm/issues/6236

* fix(user_api_key_auth.py): remove type ignore

* fix(user_api_key_auth.py): route ui vs. api token checks differently

Fixes https://github.com/BerriAI/litellm/issues/6238

* feat(internal_user_endpoints.py): support setting models as a default internal user param

Closes https://github.com/BerriAI/litellm/issues/6239

* fix(user_api_key_auth.py): fix exception string

* fix(user_api_key_auth.py): fix error string

* fix: fix test

tests/local_testing/test_get_model_info.py

@@ -11,6 +11,7 @@ import pytest
 
 import litellm
 from litellm import get_model_info
+from unittest.mock import AsyncMock, MagicMock, patch
 
 
 def test_get_model_info_simple_model_name():
@@ -74,3 +75,25 @@ def test_get_model_info_gemini_pro():
     info = litellm.get_model_info("gemini-1.5-pro-002")
     print("info", info)
     assert info["key"] == "gemini-1.5-pro-002"
+
+
+def test_get_model_info_ollama_chat():
+    from litellm.llms.ollama import OllamaConfig
+
+    with patch.object(
+        litellm.module_level_client,
+        "post",
+        return_value=MagicMock(
+            json=lambda: {
+                "model_info": {"llama.context_length": 32768},
+                "template": "tools",
+            }
+        ),
+    ):
+        info = OllamaConfig().get_model_info("mistral")
+        print("info", info)
+        assert info["supports_function_calling"] is True
+
+        info = get_model_info("ollama/mistral")
+        print("info", info)
+        assert info["supports_function_calling"] is True

tests/local_testing/test_proxy_utils.py

@@ -406,3 +406,29 @@ def test_add_litellm_data_for_backend_llm_call(headers, expected_data):
     data = add_litellm_data_for_backend_llm_call(headers)
 
     assert json.dumps(data, sort_keys=True) == json.dumps(expected_data, sort_keys=True)
+
+
+def test_update_internal_user_params():
+    from litellm.proxy.management_endpoints.internal_user_endpoints import (
+        _update_internal_user_params,
+    )
+    from litellm.proxy._types import NewUserRequest
+
+    litellm.default_internal_user_params = {
+        "max_budget": 100,
+        "budget_duration": "30d",
+        "models": ["gpt-3.5-turbo"],
+    }
+
+    data = NewUserRequest(user_role="internal_user", user_email="krrish3@berri.ai")
+    data_json = data.model_dump()
+    updated_data_json = _update_internal_user_params(data_json, data)
+    assert updated_data_json["models"] == litellm.default_internal_user_params["models"]
+    assert (
+        updated_data_json["max_budget"]
+        == litellm.default_internal_user_params["max_budget"]
+    )
+    assert (
+        updated_data_json["budget_duration"]
+        == litellm.default_internal_user_params["budget_duration"]
+    )

tests/local_testing/test_user_api_key_auth.py

@@ -291,3 +291,28 @@ async def test_auth_with_allowed_routes(route, should_raise_error):
         await user_api_key_auth(request=request, api_key="Bearer " + user_key)
 
     setattr(proxy_server, "general_settings", initial_general_settings)
+
+
+@pytest.mark.parametrize("route", ["/global/spend/logs", "/key/delete"])
+def test_is_ui_route_allowed(route):
+    from litellm.proxy.auth.user_api_key_auth import _is_ui_route_allowed
+    from litellm.proxy._types import LiteLLM_UserTable
+
+    received_args: dict = {
+        "route": route,
+        "user_obj": LiteLLM_UserTable(
+            user_id="3b803c0e-666e-4e99-bd5c-6e534c07e297",
+            max_budget=None,
+            spend=0.0,
+            model_max_budget={},
+            model_spend={},
+            user_email="my-test-email@1234.com",
+            models=[],
+            tpm_limit=None,
+            rpm_limit=None,
+            user_role="internal_user",
+            organization_memberships=[],
+        ),
+    }
+
+    assert _is_ui_route_allowed(**received_args)

tests/local_testing/test_utils.py

@@ -448,24 +448,19 @@ def test_token_counter():
 # test_token_counter()
 
 
-def test_supports_function_calling():
+@pytest.mark.parametrize(
+    "model, expected_bool",
+    [
+        ("gpt-3.5-turbo", True),
+        ("azure/gpt-4-1106-preview", True),
+        ("groq/gemma-7b-it", True),
+        ("anthropic.claude-instant-v1", False),
+        ("palm/chat-bison", False),
+    ],
+)
+def test_supports_function_calling(model, expected_bool):
     try:
-        assert litellm.supports_function_calling(model="gpt-3.5-turbo") == True
-        assert (
-            litellm.supports_function_calling(model="azure/gpt-4-1106-preview") == True
-        )
-        assert litellm.supports_function_calling(model="groq/gemma-7b-it") == True
-        assert (
-            litellm.supports_function_calling(model="anthropic.claude-instant-v1")
-            == False
-        )
-        assert litellm.supports_function_calling(model="palm/chat-bison") == False
-        assert litellm.supports_function_calling(model="ollama/llama2") == False
-        assert (
-            litellm.supports_function_calling(model="anthropic.claude-instant-v1")
-            == False
-        )
-        assert litellm.supports_function_calling(model="claude-2") == False
+        assert litellm.supports_function_calling(model=model) == expected_bool
     except Exception as e:
         pytest.fail(f"Error occurred: {e}")