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

litellm/llms/prompt_templates/common_utils.py

@@ -3,11 +3,25 @@ Common utility functions used for translating messages across providers
 """
 
 import json
-from typing import Dict, List
+from copy import deepcopy
+from typing import Dict, List, Literal, Optional
 
-from litellm.types.llms.openai import AllMessageValues
+import litellm
+from litellm.types.llms.openai import (
+    AllMessageValues,
+    ChatCompletionAssistantMessage,
+    ChatCompletionUserMessage,
+)
 from litellm.types.utils import Choices, ModelResponse, StreamingChoices
 
+DEFAULT_USER_CONTINUE_MESSAGE = ChatCompletionUserMessage(
+    content="Please continue.", role="user"
+)
+
+DEFAULT_ASSISTANT_CONTINUE_MESSAGE = ChatCompletionAssistantMessage(
+    content="Please continue.", role="assistant"
+)
+
 
 def convert_content_list_to_str(message: AllMessageValues) -> str:
     """
@@ -69,3 +83,145 @@ def get_content_from_model_response(response: ModelResponse) -> str:
         elif isinstance(choice, StreamingChoices):
             content += getattr(choice, "delta", {}).get("content", "") or ""
     return content
+
+
+def detect_first_expected_role(
+    messages: List[AllMessageValues],
+) -> Optional[Literal["user", "assistant"]]:
+    """
+    Detect the first expected role based on the message sequence.
+
+    Rules:
+    1. If messages list is empty, assume 'user' starts
+    2. If first message is from assistant, expect 'user' next
+    3. If first message is from user, expect 'assistant' next
+    4. If first message is system, look at the next non-system message
+
+    Returns:
+        str: Either 'user' or 'assistant'
+        None: If no 'user' or 'assistant' messages provided
+    """
+    if not messages:
+        return "user"
+
+    for message in messages:
+        if message["role"] == "system":
+            continue
+        return "user" if message["role"] == "assistant" else "assistant"
+
+    return None
+
+
+def _insert_user_continue_message(
+    messages: List[AllMessageValues],
+    user_continue_message: Optional[ChatCompletionUserMessage],
+    ensure_alternating_roles: bool,
+) -> List[AllMessageValues]:
+    """
+    Inserts a user continue message into the messages list.
+    Handles three cases:
+    1. Initial assistant message
+    2. Final assistant message
+    3. Consecutive assistant messages
+
+    Only inserts messages between consecutive assistant messages,
+    ignoring all other role types.
+    """
+    if not messages:
+        return messages
+
+    result_messages = messages.copy()  # Don't modify the input list
+    continue_message = user_continue_message or DEFAULT_USER_CONTINUE_MESSAGE
+
+    # Handle first message if it's an assistant message
+    if result_messages[0]["role"] == "assistant":
+        result_messages.insert(0, continue_message)
+
+    # Handle consecutive assistant messages and final message
+    i = 1  # Start from second message since we handled first message
+    while i < len(result_messages):
+        curr_message = result_messages[i]
+        prev_message = result_messages[i - 1]
+
+        # Only check for consecutive assistant messages
+        # Ignore all other role types
+        if curr_message["role"] == "assistant" and prev_message["role"] == "assistant":
+            result_messages.insert(i, continue_message)
+            i += 2  # Skip over the message we just inserted
+        else:
+            i += 1
+
+    # Handle final message
+    if result_messages[-1]["role"] == "assistant" and ensure_alternating_roles:
+        result_messages.append(continue_message)
+
+    return result_messages
+
+
+def _insert_assistant_continue_message(
+    messages: List[AllMessageValues],
+    assistant_continue_message: Optional[ChatCompletionAssistantMessage] = None,
+    ensure_alternating_roles: bool = True,
+) -> List[AllMessageValues]:
+    """
+    Add assistant continuation messages between consecutive user messages.
+
+    Args:
+        messages: List of message dictionaries
+        assistant_continue_message: Optional custom assistant message
+        ensure_alternating_roles: Whether to enforce alternating roles
+
+    Returns:
+        Modified list of messages with inserted assistant messages
+    """
+    if not ensure_alternating_roles or len(messages) <= 1:
+        return messages
+
+    # Create a new list to store modified messages
+    modified_messages: List[AllMessageValues] = []
+
+    for i, message in enumerate(messages):
+        modified_messages.append(message)
+
+        # Check if we need to insert an assistant message
+        if (
+            i < len(messages) - 1  # Not the last message
+            and message.get("role") == "user"  # Current is user
+            and messages[i + 1].get("role") == "user"
+        ):  # Next is user
+
+            # Insert assistant message
+            continue_message = (
+                assistant_continue_message or DEFAULT_ASSISTANT_CONTINUE_MESSAGE
+            )
+            modified_messages.append(continue_message)
+
+    return modified_messages
+
+
+def get_completion_messages(
+    messages: List[AllMessageValues],
+    assistant_continue_message: Optional[ChatCompletionAssistantMessage],
+    user_continue_message: Optional[ChatCompletionUserMessage],
+    ensure_alternating_roles: bool,
+) -> List[AllMessageValues]:
+    """
+    Ensures messages alternate between user and assistant roles by adding placeholders
+    only when there are consecutive messages of the same role.
+
+    1. ensure 'user' message before 1st 'assistant' message
+    2. ensure 'user' message after last 'assistant' message
+    """
+    if not ensure_alternating_roles:
+        return messages.copy()
+
+    ## INSERT USER CONTINUE MESSAGE
+    messages = _insert_user_continue_message(
+        messages, user_continue_message, ensure_alternating_roles
+    )
+
+    ## INSERT ASSISTANT CONTINUE MESSAGE
+    messages = _insert_assistant_continue_message(
+        messages, assistant_continue_message, ensure_alternating_roles
+    )
+    return messages