model_serving_utils.py

import json
import logging
from mlflow.deployments import get_deploy_client
from databricks.sdk import WorkspaceClient

logger = logging.getLogger(__name__)


def get_endpoint_task_type(endpoint_name: str) -> str:
    """Get the task type of a serving endpoint."""
    w = WorkspaceClient()
    ep = w.serving_endpoints.get(endpoint_name)
    return ep.task


def is_endpoint_supported(endpoint_name: str) -> bool:
    """Check if the endpoint has a supported task type."""
    task_type = get_endpoint_task_type(endpoint_name)
    supported_task_types = ["agent/v1/chat", "agent/v2/chat", "llm/v1/chat"]
    return task_type in supported_task_types


def query_endpoint(endpoint_name: str, messages: list[dict]) -> dict:
    """
    Query a Databricks serving endpoint with messages.
    
    Supports multiple response formats:
    - Standard ChatCompletions (choices)
    - Databricks agents (messages)
    - Custom multi-agent output format
    
    Returns a dict with role and content.
    """
    client = get_deploy_client('databricks')
    
    # Count assistant messages in the request (these may be echoed in response)
    assistant_messages_sent = sum(1 for msg in messages if msg.get('role') == 'assistant')
    
    # Log the request
    logger.info("=" * 80)
    logger.info("📤 REQUEST TO ENDPOINT: %s", endpoint_name)
    logger.info("Number of messages: %d", len(messages))
    logger.info("Assistant messages already sent: %d", assistant_messages_sent)
    logger.info("Messages being sent:")
    for i, msg in enumerate(messages):
        content_preview = msg.get('content', '')[:100]
        logger.info("  [%d] %s: %s%s", i, msg.get('role'), content_preview, 
                   '...' if len(msg.get('content', '')) > 100 else '')
    
    # Try standard format first, fallback to alternative format
    try:
        inputs = {'messages': messages}
        logger.info("Using input format: 'messages'")
        response = client.predict(endpoint=endpoint_name, inputs=inputs)
    except Exception as e:
        if "Model is missing inputs ['input']" in str(e):
            inputs = {'input': messages}
            logger.info("Retrying with input format: 'input'")
            response = client.predict(endpoint=endpoint_name, inputs=inputs)
        else:
            raise
    
    # Log the raw response
    logger.info("-" * 80)
    logger.info("📥 RAW RESPONSE FROM ENDPOINT:")
    logger.info("Response type: %s", type(response).__name__)
    logger.info("Response keys: %s", list(response.keys()) if isinstance(response, dict) else 'N/A')
    logger.info("Response: %s", response)

    # Log response structure details
    if "output" in response:
        logger.info("Output array length: %d", len(response.get("output", [])))
        logger.info("Output items:")
        for i, item in enumerate(response.get("output", [])):
            logger.info("  [%d] type=%s, role=%s", i, 
                       item.get("type"), item.get("role", "N/A"))
    elif "choices" in response:
        logger.info("Choices format detected")
    elif "messages" in response:
        logger.info("Messages format detected, count: %d", len(response.get("messages", [])))
    
    # Parse response based on format
    parsed = _parse_response(response, assistant_messages_sent)
    
    # Log the parsed response
    logger.info("-" * 80)
    logger.info("✅ PARSED RESPONSE:")
    logger.info("Role: %s", parsed.get('role'))
    logger.info("Content length: %d characters", len(parsed.get('content', '')))
    logger.info("Content preview: %s...", parsed.get('content', '')[:200])
    logger.info("=" * 80)
    
    return parsed


def _parse_response(response: dict, assistant_messages_sent: int = 0) -> dict:
    """Parse different response formats and return a standardized message.
    
    Args:
        response: The response from the endpoint
        assistant_messages_sent: Number of assistant messages in the request (to skip echoed messages)
    """
    
    # Format 1: Databricks agent format with messages array
    if "messages" in response:
        return response["messages"][-1]
    
    # Format 2: Standard ChatCompletions with choices
    if "choices" in response:
        message = response["choices"][0]["message"]
        content = message.get("content")
        
        # Handle structured content (list of parts)
        if isinstance(content, list):
            text = "".join(
                part.get("text", "")
                for part in content
                if part.get("type") == "text"
            )
            return {"role": message.get("role"), "content": text}
        
        # Handle simple string content
        return message
    
    # Format 3: Custom multi-agent output format
    if "output" in response:
        messages = []
        logger.info("Parsing custom output format...")
        
        # Find the index of the last user message
        last_user_index = -1
        for i, item in enumerate(response["output"]):
            if item.get("type") == "message" and item.get("role") == "user":
                last_user_index = i
        
        logger.info("Last user message index: %d", last_user_index)
        
        # Collect ALL assistant messages first
        all_messages = []
        for i, item in enumerate(response["output"]):
            if item.get("type") == "message" and item.get("role") == "assistant":
                # Extract text from content parts
                text = "".join(
                    part.get("text", "")
                    for part in item.get("content", [])
                    if part.get("type") == "output_text"
                )
                
                # Log each message found
                is_metadata = _is_metadata_message(text)
                logger.info("  Message [%d]: %s (metadata=%s)", i, 
                           text[:80] + '...' if len(text) > 80 else text, is_metadata)
                
                # Filter out metadata messages and store with index
                if text and not is_metadata:
                    all_messages.append({"index": i, "role": "assistant", "content": text})
        
        logger.info("Total assistant messages found: %d", len(all_messages))
        
        # Strategy: Skip the first N assistant messages that were already sent in the request
        # This handles cases where the endpoint echoes back conversation history
        if assistant_messages_sent > 0 and len(all_messages) > assistant_messages_sent:
            messages = all_messages[assistant_messages_sent:]
            logger.info("Skipping first %d assistant msgs (already in request) → keeping %d new msgs", 
                       assistant_messages_sent, len(messages))
        elif last_user_index == -1:
            # Fallback: No user message found - take ONLY the LAST assistant message
            if all_messages:
                messages = [all_messages[-1]]
                logger.info("No user msg in output → taking ONLY last assistant msg at index %d", messages[0]["index"])
        else:
            # Fallback: User message found - take all messages AFTER it
            messages = [msg for msg in all_messages if msg["index"] > last_user_index]
            logger.info("User msg at index %d → taking %d assistant msgs after it", last_user_index, len(messages))
        
        logger.info("Total assistant messages extracted: %d (after filtering)", len(messages))
        
        if messages:
            # Combine all messages with separators
            combined = "\n\n---\n\n".join(msg["content"] for msg in messages)
            logger.info("Combined message length: %d characters", len(combined))
            return {"role": "assistant", "content": combined}
    
    raise Exception(
        "Unsupported response format. This app supports:\n"
        "1. Standard ChatCompletions (choices)\n"
        "2. Databricks agents (messages)\n"
        "3. Custom multi-agent format (output)"
    )


def _is_metadata_message(text: str) -> bool:
    """Check if a message is metadata that should be filtered out."""
    text_lower = text.lower()
    return (
        text.strip().startswith("<name>") or
        "transferring" in text_lower or
        "successfully transferred" in text_lower
    )