Fixing LiteLLM With Azure OpenAI: Load_artifacts_tool Issues

by Alex Johnson 61 views

Introduction

In the realm of AI development, the load_artifacts_tool is a powerful asset, enabling agents to seamlessly interact with various file types. However, when integrating LiteLLM with Azure OpenAI, developers have encountered frustrating issues, specifically concerning null tool_call.id and unsupported text artifacts. This article dives deep into these bugs, offering comprehensive explanations, reproduction steps, suggested fixes, and a practical workaround. Understanding these challenges and their solutions is crucial for anyone aiming to leverage the full potential of Azure OpenAI within the LiteLLM framework. We'll explore the root causes, examine code snippets, and provide clear, actionable steps to ensure your load_artifacts_tool functions flawlessly.

Understanding the Bugs: Null tool_call.id and Unsupported Text Artifacts

The google.adk.models.lite_llm.LiteLlm adapter, while generally robust, exhibits two critical bugs that hinder the functionality of load_artifacts_tool when used with Azure OpenAI models. These issues manifest as follows:

  1. Null tool_call.id Values: The Azure OpenAI API vehemently rejects requests containing null values for tool_call.id, leading to errors like: Invalid type for 'messages[N].tool_calls[0].id': expected a string, but got null instead. This occurs because the tool_call.id field, which should be a unique identifier, is sometimes passed as null, causing the API to balk.
  2. Text Artifact Format Unsupported: The second bug arises when attempting to load text-based files (e.g., .txt, .json, .md). This results in a ValueError: LiteLlm(BaseLlm) does not support this content part. This error indicates that the LiteLlm adapter lacks the necessary logic to process and incorporate text artifacts into the conversation context.

These issues, while problematic, are specific to the interaction between LiteLLM and Azure OpenAI. Gemini models, for instance, do not exhibit these behaviors, highlighting the nuanced nature of AI model integrations. To fully grasp the impact and devise effective solutions, a detailed examination of these bugs is essential.

Reproducing the Bugs: Step-by-Step Guide

To effectively address these bugs, it's crucial to understand how to reproduce them. Here are detailed steps to recreate the issues with null tool_call.id and unsupported text artifacts.

Issue #1: Null tool_call.id

  1. Set up the Environment:

    • Ensure you have the necessary libraries installed, including google-adk and litellm. You can install these using pip:
      pip install google-adk litellm
    • Configure your environment variables for Azure OpenAI, including AZURE_OPENAI_API_BASE and AZURE_OPENAI_API_KEY. These are essential for authenticating your requests to the Azure OpenAI service.

  2. Write the Code:

    • Use the following Python code snippet to reproduce the bug:
      from google.adk import Agent
from google.adk.tools import load_artifacts_tool
from google.adk.models.lite_llm import LiteLlm
import os

model = LiteLlm(
    model="azure/gpt-4o",
    api_base=os.getenv("AZURE_OPENAI_API_BASE"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY")
)

agent = Agent(
    model=model,
    tools=[load_artifacts_tool],
    instruction="Load and read the user's files."
)

# Upload any artifact via artifact service, then:
async def run_agent():
    response = await agent.run("Load my files")
    print(response)

import asyncio
asyncio.run(run_agent())

  3. Run the Code:

    • Execute the script. If an artifact has been uploaded via the artifact service, the code will attempt to load it.

  4. Observe the Error:

    • You should encounter the following error message:
      Invalid type for 'messages[6].tool_calls[0].id': expected a string, but got null instead.

Issue #2: Text Artifacts Unsupported

  1. Set up the Environment:

    • Ensure the same environment setup as Issue #1.

  2. Create a Text File:

    • Create a simple text file (e.g., notes.txt) with some content.

  3. Upload the Text File:

    • Upload the .txt file artifact via your artifact service. The specifics of this process depend on your artifact service implementation.

  4. Write the Code:

    • Use the following Python code snippet:
      from google.adk import Agent
from google.adk.tools import load_artifacts_tool
from google.adk.models.lite_llm import LiteLlm
import os

model = LiteLlm(
    model="azure/gpt-4o",
    api_base=os.getenv("AZURE_OPENAI_API_BASE"),
    api_key=os.getenv("AZURE_OPENAI_API_KEY")
)

agent = Agent(
    model=model,
    tools=[load_artifacts_tool],
    instruction="Load and read the user's files."
)

# Upload a .txt file artifact via artifact service, then:
async def run_agent():
    response = await agent.run("Load my notes.txt file")
    print(response)

import asyncio
asyncio.run(run_agent())

  5. Run the Code:

    • Execute the script.

  6. Observe the Error:

    • You should encounter the following error message:
      ValueError: LiteLlm(BaseLlm) does not support this content part.

By following these steps, you can reliably reproduce both bugs, which is a critical first step in verifying any proposed fixes.

Root Causes and Suggested Fixes

Identifying the root causes of these bugs is essential for devising effective solutions. Let's delve into the underlying issues and explore suggested fixes.

Root Cause #1: Null tool_call.id

The problem lies in the google/adk/models/lite_llm.py file, specifically at line 200 within the _content_to_message_param() function:

tool_calls.append(
    ChatCompletionAssistantToolCall(
        type="function",
        id=part.function_call.id,  # ← Can be None for some models!
        function=Function(
            name=part.function_call.name,
            arguments=_safe_json_serialize(part.function_call.args),
        ),
    )
)

Here, the tool_call.id is directly assigned from part.function_call.id, which can be None for certain models. Azure OpenAI requires this ID to be a non-null string, leading to the API rejection.

Suggested Fix:

To address this, we can generate a UUID if part.function_call.id is None. This ensures that a unique, non-null ID is always provided.

import uuid

tool_call_id = part.function_call.id or f"call_{uuid.uuid4().hex[:24]}"

tool_calls.append(
    ChatCompletionAssistantToolCall(
        type="function",
        id=tool_call_id,  # ← Now guaranteed non-null
        function=Function(
            name=part.function_call.name,
            arguments=_safe_json_serialize(part.function_call.args),
        ),
    )
)

This fix introduces a check for None and generates a UUID if necessary, ensuring compliance with Azure OpenAI's requirements.

Root Cause #2: Text Artifacts Unsupported

The second issue stems from the limited support for content types in google/adk/models/lite_llm.py, specifically within the _get_content() function (lines 248-284). The function only supports inline_data for images, videos, audio, and PDFs, neglecting text files.

elif (
    part.inline_data
    and part.inline_data.data
    and part.inline_data.mime_type
):
    # ... base64 encoding ...
    if part.inline_data.mime_type.startswith("image"):
        # OK
    elif part.inline_data.mime_type.startswith("video"):
        # OK
    elif part.inline_data.mime_type.startswith("audio"):
        # OK
    elif part.inline_data.mime_type == "application/pdf":
        # OK
    else:
        raise ValueError("LiteLlm(BaseLlm) does not support this content part.")
        # ← text/plain, application/json, etc. hit this!

When load_artifacts_tool loads text files, they are provided as inline_data with MIME types like text/plain, application/json, and application/xml. These MIME types are not handled, resulting in the ValueError.

Suggested Fix:

To remedy this, we need to add support for text-based MIME types. This involves decoding the text content and appending it to the content objects.

# Add before the final else clause:
elif _is_text_mime_type(part.inline_data.mime_type):
    # Handle text-based artifacts by decoding to text
    try:
        text_content = part.inline_data.data.decode("utf-8")
        content_objects.append({
            "type": "text",
            "text": text_content,
        })
    except UnicodeDecodeError:
        raise ValueError("LiteLlm(BaseLlm) does not support this content part.")
else:
    raise ValueError("LiteLlm(BaseLlm) does not support this content part.")

# Helper function:
def _is_text_mime_type(mime_type: str) -> bool:
    """Check if MIME type represents text content."""
    text_prefixes = [
        "text/",
        "application/json",
        "application/xml",
        "application/javascript",
        "application/x-yaml",
    ]
    return any(mime_type.startswith(prefix) for prefix in text_prefixes)

This fix adds a new conditional block that checks for text-based MIME types, decodes the content, and appends it as a text object. A helper function, _is_text_mime_type, is used to efficiently check for various text-related MIME prefixes.

Workaround: A Practical Solution

While the suggested fixes provide a direct solution to the bugs, a workaround can offer immediate relief. A practical approach involves creating a wrapper class that subclasses LiteLlm and implements the necessary fixes.

import uuid
from google.adk.models.lite_llm import LiteLlm
from litellm import completion
from litellm.utils import CustomStreamWrapper
from google.ai.generativelanguage import ChatCompletionMessage, ToolCall, Function
from google.ai.generativelanguage import Part as APart
from typing import Optional, Sequence, Mapping, Any, Dict, List
from google.adk import types
import json

class AzureFixedLiteLlm(LiteLlm):
    """LiteLLM wrapper that fixes Azure OpenAI compatibility issues."""

    async def generate_content_async(self, llm_request: types.GenerateContentRequest, stream=False):
        # 1. Convert text artifacts from inline_data to text parts
        self._convert_text_artifacts(llm_request)

        # 2. Get formatted messages
        messages, tools, response_format, generation_params = (
            self._get_completion_inputs(llm_request)
        )

        # 3. Fix null tool_call IDs
        for message in messages:
            if hasattr(message, "tool_calls") and message.tool_calls:
                for tool_call in message.tool_calls:
                    if tool_call.id is None:
                        tool_call.id = f"call_{uuid.uuid4().hex[:24]}"

        # 4. Continue with normal LiteLLM flow
        completion_args = {
            "model": self.model,
            "messages": messages,
            "tools": tools,
            "response_format": response_format,
            **generation_params  # Pass all generation params to LiteLLM
        }
        if stream:
            # LiteLLM Streaming
            response = await completion(
                **completion_args,
                stream=True
            )
            return CustomStreamWrapper(response)
        else:
            # Non-streaming
            response = await completion(
                **completion_args
            )
            return self._process_response(response)

    def _convert_text_artifacts(self, llm_request: types.GenerateContentRequest):
        """Converts text artifacts from inline_data to text parts."""
        for content in llm_request.prompt.parts:
            if content.inline_data and content.inline_data.mime_type and self._is_text_mime_type(content.inline_data.mime_type):
                text_content = content.inline_data.data.decode("utf-8", errors="ignore")
                content.text = text_content
                content.inline_data = None

    def _get_completion_inputs(self, llm_request: types.GenerateContentRequest):
        """
        Formats the messages and tools for LiteLLM's completion function.
        Handles conversion from GoogleADK types to LiteLLM-compatible types.
        """
        messages = self._content_to_message_param(llm_request.prompt)
        tools = [tool.function_declarations for tool in llm_request.tools] if llm_request.tools else None
        response_format = {"type": llm_request.response_mime_type} if llm_request.response_mime_type else None

        # Extract generation parameters
        generation_params = {}
        if llm_request.generation_config:
            generation_params["temperature"] = llm_request.generation_config.temperature
            generation_params["top_p"] = llm_request.generation_config.top_p
            if llm_request.generation_config.candidate_count is not None and llm_request.generation_config.candidate_count > 0:
                generation_params["n"] = llm_request.generation_config.candidate_count
            if llm_request.generation_config.max_output_tokens:
                generation_params["max_tokens"] = llm_request.generation_config.max_output_tokens

        return messages, tools, response_format, generation_params

    def _content_to_message_param(self, content: types.Content) -> List[Dict[str, Any]]:
        """
        Converts GoogleADK Content to LiteLLM message format.
        Handles different types of content including text and tool calls.
        """
        messages = []
        if content is None:
            return messages

        message_params = {"role": content.role}
        content_objects = []

        for part in content.parts:
            if part.text:
                content_objects.append({"type": "text", "text": part.text})
            elif part.inline_data and part.inline_data.data and part.inline_data.mime_type:
                if part.inline_data.mime_type.startswith("image"):
                    # Handle images (base64 encoding)
                    base64_image = self._encode_image_base64(part.inline_data)
                    content_objects.append({"type": "image_url", "image_url": {"url": f"data:{part.inline_data.mime_type};base64,{base64_image}"}})
                elif part.inline_data.mime_type.startswith("video") or part.inline_data.mime_type.startswith("audio"):
                    # For video and audio, include a placeholder message
                    content_objects.append({"type": "text", "text": f"(Inline {part.inline_data.mime_type} content)"})  # Placeholder for video/audio
                elif part.inline_data.mime_type == "application/pdf":
                    # For PDF, include a placeholder message
                    content_objects.append({"type": "text", "text": "(Inline PDF content)"})  # Placeholder for PDF
            elif hasattr(part, "function_call") and part.function_call:
                tool_call_id = part.function_call.id or f"call_{uuid.uuid4().hex[:24]}"  # Generate UUID if None
                content_objects.append({"tool_calls": [{
                    "id": tool_call_id,
                    "function": {
                        "arguments": json.dumps(part.function_call.args),
                        "name": part.function_call.name
                    },
                    "type": "function"
                }]})
        message_params["content"] = content_objects
        messages.append(message_params)
        return messages


    def _is_text_mime_type(self, mime_type: str) -> bool:
        """Check if MIME type represents text content."""
        text_prefixes = [
            "text/",
            "application/json",
            "application/xml",
            "application/javascript",
            "application/x-yaml",
        ]
        return any(mime_type.startswith(prefix) for prefix in text_prefixes)

    def _process_response(self, response: dict) -> types.GenerateContentResponse:
        """
        Processes LiteLLM API response and converts it to GoogleADK format.
        Handles extracting text, tool calls, and other content types.
        """
        if not response or not response.get("choices"):
            return types.GenerateContentResponse()

        parts = []
        for choice in response["choices"]:
            if choice.get("message"):
                message = choice["message"]
                if message.get("content"):
                    if isinstance(message["content"], str):
                        parts.append(APart(text=message["content"]))
                    elif isinstance(message["content"], list):
                        for content_item in message["content"]:
                            if content_item.get("type") == "text":
                                parts.append(APart(text=content_item["text"]))
                if message.get("tool_calls"):
                    for tool_call in message["tool_calls"]:
                        parts.append(APart(function_call=types.FunctionCall(
                            name=tool_call["function"]["name"],
                            args=json.loads(tool_call["function"]["arguments"])
                        )))
        return types.GenerateContentResponse(content=types.Content(parts=parts))

    def _encode_image_base64(self, inline_data: types.InlineData) -> str:
        """Encodes image data to base64 format."""
        import base64
        return base64.b64encode(inline_data.data).decode("utf-8")

This wrapper, AzureFixedLiteLlm, addresses the issues by:

  1. Converting text artifacts from inline_data to text parts.
  2. Generating UUIDs for null tool_call IDs.
  3. Preserving image/video/audio artifacts unchanged.

This workaround allows developers to continue using Azure OpenAI with LiteLLM while a permanent fix is implemented.

Impact and Conclusion

The bugs discussed in this article have a significant impact on developers using LiteLLM with Azure OpenAI. They force users to:

  • Rely exclusively on Gemini models for artifact-loading agents.
  • Forego the multi-model flexibility offered by ADK.
  • Implement custom workarounds, adding complexity and maintenance overhead.

By addressing these issues, developers can unlock the full potential of Azure OpenAI within the LiteLLM framework, ensuring seamless integration and functionality.

In conclusion, the null tool_call.id and unsupported text artifacts bugs in LiteLLM's interaction with Azure OpenAI present significant challenges. However, by understanding the root causes, applying the suggested fixes, or using the provided workaround, developers can overcome these hurdles and leverage the powerful capabilities of both technologies. Embracing these solutions ensures a smoother, more efficient development experience.

For more information on LiteLLM, visit their official website: LiteLLM Documentation