forked from phoenix/litellm-mirror
LiteLLM Minor Fixes & Improvements (10/17/2024) (#6293)
* fix(ui_sso.py): fix faulty admin only check Fixes https://github.com/BerriAI/litellm/issues/6286 * refactor(sso_helper_utils.py): refactor /sso/callback to use helper utils, covered by unit testing Prevent future regressions * feat(prompt_factory): support 'ensure_alternating_roles' param Closes https://github.com/BerriAI/litellm/issues/6257 * fix(proxy/utils.py): add dailytagspend to expected views * feat(auth_utils.py): support setting regex for clientside auth credentials Fixes https://github.com/BerriAI/litellm/issues/6203 * build(cookbook): add tutorial for mlflow + langchain + litellm proxy tracing * feat(argilla.py): add argilla logging integration Closes https://github.com/BerriAI/litellm/issues/6201 * fix: fix linting errors * fix: fix ruff error * test: fix test * fix: update vertex ai assumption - parts not always guaranteed (#6296) * docs(configs.md): add argila env var to docs
This commit is contained in:
Krish Dholakia
GitHub
parent
5e381caf75
commit
f252350881
23 changed files with 1388 additions and 43 deletions
70
tests/local_testing/test_auth_utils.py
Normal file
70
tests/local_testing/test_auth_utils.py
Normal file
|
|
@ -0,0 +1,70 @@
|
|||
# What is this?
|
||||
## Tests if proxy/auth/auth_utils.py works as expected
|
||||
|
||||
import sys, os, asyncio, time, random, uuid
|
||||
import traceback
|
||||
from dotenv import load_dotenv
|
||||
|
||||
load_dotenv()
|
||||
import os
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
import pytest
|
||||
import litellm
|
||||
from litellm.proxy.auth.auth_utils import (
|
||||
_allow_model_level_clientside_configurable_parameters,
|
||||
)
|
||||
from litellm.router import Router
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"allowed_param, input_value, should_return_true",
|
||||
[
|
||||
("api_base", {"api_base": "http://dummy.com"}, True),
|
||||
(
|
||||
{"api_base": "https://api.openai.com/v1"},
|
||||
{"api_base": "https://api.openai.com/v1"},
|
||||
True,
|
||||
), # should return True
|
||||
(
|
||||
{"api_base": "https://api.openai.com/v1"},
|
||||
{"api_base": "https://api.anthropic.com/v1"},
|
||||
False,
|
||||
), # should return False
|
||||
(
|
||||
{"api_base": "^https://litellm.*direct\.fireworks\.ai/v1$"},
|
||||
{"api_base": "https://litellm-dev.direct.fireworks.ai/v1"},
|
||||
True,
|
||||
),
|
||||
(
|
||||
{"api_base": "^https://litellm.*novice\.fireworks\.ai/v1$"},
|
||||
{"api_base": "https://litellm-dev.direct.fireworks.ai/v1"},
|
||||
False,
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_configurable_clientside_parameters(
|
||||
allowed_param, input_value, should_return_true
|
||||
):
|
||||
router = Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "dummy-model",
|
||||
"litellm_params": {
|
||||
"model": "gpt-3.5-turbo",
|
||||
"api_key": "dummy-key",
|
||||
"configurable_clientside_auth_params": [allowed_param],
|
||||
},
|
||||
}
|
||||
]
|
||||
)
|
||||
resp = _allow_model_level_clientside_configurable_parameters(
|
||||
model="dummy-model",
|
||||
param="api_base",
|
||||
request_body_value=input_value["api_base"],
|
||||
llm_router=router,
|
||||
)
|
||||
print(resp)
|
||||
assert resp == should_return_true
|
||||
|
|
@ -22,9 +22,13 @@ from litellm.llms.prompt_templates.factory import (
|
|||
llama_2_chat_pt,
|
||||
prompt_factory,
|
||||
)
|
||||
from litellm.llms.prompt_templates.common_utils import (
|
||||
get_completion_messages,
|
||||
)
|
||||
from litellm.llms.vertex_ai_and_google_ai_studio.vertex_ai_non_gemini import (
|
||||
_gemini_convert_messages_with_history,
|
||||
)
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
||||
def test_llama_3_prompt():
|
||||
|
|
@ -457,3 +461,217 @@ def test_azure_tool_call_invoke_helper():
|
|||
"function_call": {"name": "get_weather", "arguments": ""},
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"messages, expected_messages, user_continue_message, assistant_continue_message",
|
||||
[
|
||||
(
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Hello! How can I assist you today?"},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
{"role": "assistant", "content": "I don't know anyything, do you?"},
|
||||
],
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I assist you today?",
|
||||
},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please continue.",
|
||||
},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I don't know anyything, do you?",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Please continue.",
|
||||
},
|
||||
],
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
],
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
],
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
],
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Please continue."},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
],
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
],
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Please continue."},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please continue.",
|
||||
},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
],
|
||||
None,
|
||||
None,
|
||||
),
|
||||
(
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I assist you today?",
|
||||
},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
{"role": "assistant", "content": "I don't know anyything, do you?"},
|
||||
{"role": "assistant", "content": "I can't repeat sentences."},
|
||||
],
|
||||
[
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I assist you today?",
|
||||
},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please continue",
|
||||
},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I don't know anyything, do you?",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Ok",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I can't repeat sentences.",
|
||||
},
|
||||
{"role": "user", "content": "Ok"},
|
||||
],
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Ok",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please continue",
|
||||
},
|
||||
),
|
||||
],
|
||||
)
|
||||
def test_ensure_alternating_roles(
|
||||
messages, expected_messages, user_continue_message, assistant_continue_message
|
||||
):
|
||||
|
||||
messages = get_completion_messages(
|
||||
messages=messages,
|
||||
assistant_continue_message=assistant_continue_message,
|
||||
user_continue_message=user_continue_message,
|
||||
ensure_alternating_roles=True,
|
||||
)
|
||||
|
||||
print(messages)
|
||||
|
||||
assert messages == expected_messages
|
||||
|
||||
|
||||
def test_alternating_roles_e2e():
|
||||
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
||||
import json
|
||||
|
||||
litellm.set_verbose = True
|
||||
http_handler = HTTPHandler()
|
||||
|
||||
with patch.object(http_handler, "post", new=MagicMock()) as mock_post:
|
||||
response = litellm.completion(
|
||||
**{
|
||||
"model": "databricks/databricks-meta-llama-3-1-70b-instruct",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I assist you today?",
|
||||
},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
{"role": "assistant", "content": "I don't know anyything, do you?"},
|
||||
{"role": "assistant", "content": "I can't repeat sentences."},
|
||||
],
|
||||
"user_continue_message": {
|
||||
"role": "user",
|
||||
"content": "Ok",
|
||||
},
|
||||
"assistant_continue_message": {
|
||||
"role": "assistant",
|
||||
"content": "Please continue",
|
||||
},
|
||||
"ensure_alternating_roles": True,
|
||||
},
|
||||
client=http_handler,
|
||||
)
|
||||
print(f"response: {response}")
|
||||
assert mock_post.call_args.kwargs["data"] == json.dumps(
|
||||
{
|
||||
"model": "databricks-meta-llama-3-1-70b-instruct",
|
||||
"messages": [
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Hello! How can I assist you today?",
|
||||
},
|
||||
{"role": "user", "content": "What is Databricks?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "Please continue",
|
||||
},
|
||||
{"role": "user", "content": "What is Azure?"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I don't know anyything, do you?",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Ok",
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "I can't repeat sentences.",
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": "Ok",
|
||||
},
|
||||
],
|
||||
"stream": False,
|
||||
}
|
||||
)
|
||||
|
|
|
|||
38
tests/local_testing/test_ui_sso_helper_utils.py
Normal file
38
tests/local_testing/test_ui_sso_helper_utils.py
Normal file
|
|
@ -0,0 +1,38 @@
|
|||
# What is this?
|
||||
## This tests the batch update spend logic on the proxy server
|
||||
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import random
|
||||
import sys
|
||||
import time
|
||||
import traceback
|
||||
from datetime import datetime
|
||||
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import Request
|
||||
|
||||
load_dotenv()
|
||||
|
||||
sys.path.insert(
|
||||
0, os.path.abspath("../..")
|
||||
) # Adds the parent directory to the system path
|
||||
|
||||
import logging
|
||||
from litellm.proxy.management_endpoints.sso_helper_utils import (
|
||||
check_is_admin_only_access,
|
||||
has_admin_ui_access,
|
||||
)
|
||||
from litellm.proxy._types import LitellmUserRoles
|
||||
|
||||
|
||||
def test_check_is_admin_only_access():
|
||||
assert check_is_admin_only_access("admin_only") is True
|
||||
assert check_is_admin_only_access("user_only") is False
|
||||
|
||||
|
||||
def test_has_admin_ui_access():
|
||||
assert has_admin_ui_access(LitellmUserRoles.PROXY_ADMIN.value) is True
|
||||
assert has_admin_ui_access(LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY.value) is True
|
||||
assert has_admin_ui_access(LitellmUserRoles.INTERNAL_USER_VIEW_ONLY.value) is False
|
||||
Loading…
Add table
Add a link
Reference in a new issue