# Copyright 2026 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License

import base64
import contextlib
import json
import logging
import os
import sys
import tempfile
import unittest
from unittest.mock import ANY
from unittest.mock import AsyncMock
from unittest.mock import MagicMock
from unittest.mock import Mock
from unittest.mock import patch
import warnings

from google.adk.models.lite_llm import _aggregate_streaming_thought_parts
from google.adk.models.lite_llm import _append_fallback_user_content_if_missing
from google.adk.models.lite_llm import _BraceDepthTracker
from google.adk.models.lite_llm import _content_to_message_param
from google.adk.models.lite_llm import _convert_reasoning_value_to_parts
from google.adk.models.lite_llm import _enforce_strict_openai_schema
from google.adk.models.lite_llm import _extract_json_from_deepseek_args
from google.adk.models.lite_llm import _extract_reasoning_value
from google.adk.models.lite_llm import _extract_thought_signature_from_tool_call
from google.adk.models.lite_llm import _FILE_ID_REQUIRED_PROVIDERS
from google.adk.models.lite_llm import _FINISH_REASON_MAPPING
from google.adk.models.lite_llm import _function_declaration_to_tool_param
from google.adk.models.lite_llm import _get_completion_inputs
from google.adk.models.lite_llm import _get_content
from google.adk.models.lite_llm import _get_provider_from_model
from google.adk.models.lite_llm import _is_anthropic_model
from google.adk.models.lite_llm import _is_anthropic_provider
from google.adk.models.lite_llm import _is_anthropic_route
from google.adk.models.lite_llm import _looks_like_openai_file_id
from google.adk.models.lite_llm import _message_to_generate_content_response
from google.adk.models.lite_llm import _MISSING_TOOL_RESULT_MESSAGE
from google.adk.models.lite_llm import _model_response_to_chunk
from google.adk.models.lite_llm import _model_response_to_generate_content_response
from google.adk.models.lite_llm import _parse_deepseek_tool_calls_from_text
from google.adk.models.lite_llm import _parse_tool_calls_from_text
from google.adk.models.lite_llm import _redact_file_uri_for_log
from google.adk.models.lite_llm import _redirect_litellm_loggers_to_stdout
from google.adk.models.lite_llm import _safe_json_serialize
from google.adk.models.lite_llm import _schema_to_dict
from google.adk.models.lite_llm import _split_message_content_and_tool_calls
from google.adk.models.lite_llm import _THOUGHT_SIGNATURE_SEPARATOR
from google.adk.models.lite_llm import _to_litellm_response_format
from google.adk.models.lite_llm import _to_litellm_role
from google.adk.models.lite_llm import FunctionChunk
from google.adk.models.lite_llm import LiteLlm
from google.adk.models.lite_llm import LiteLLMClient
from google.adk.models.lite_llm import ReasoningChunk
from google.adk.models.lite_llm import TextChunk
from google.adk.models.lite_llm import UsageMetadataChunk
from google.adk.models.llm_request import LlmRequest
from google.genai import types
import litellm
from litellm import ChatCompletionAssistantMessage
from litellm import ChatCompletionMessageToolCall
from litellm import Function
from litellm.types.utils import ChatCompletionDeltaToolCall
from litellm.types.utils import Choices
from litellm.types.utils import Delta
from litellm.types.utils import ModelResponse
from litellm.types.utils import ModelResponseStream
from litellm.types.utils import StreamingChoices
from pydantic import BaseModel
from pydantic import Field
import pytest

LLM_REQUEST_WITH_FUNCTION_DECLARATION = LlmRequest(
    contents=[
        types.Content(
            role="user", parts=[types.Part.from_text(text="Test prompt")]
        )
    ],
    config=types.GenerateContentConfig(
        tools=[
            types.Tool(
                function_declarations=[
                    types.FunctionDeclaration(
                        name="test_function",
                        description="Test function description",
                        parameters=types.Schema(
                            type=types.Type.OBJECT,
                            properties={
                                "test_arg": types.Schema(
                                    type=types.Type.STRING
                                ),
                                "array_arg": types.Schema(
                                    type=types.Type.ARRAY,
                                    items={
                                        "type": types.Type.STRING,
                                    },
                                ),
                                "nested_arg": types.Schema(
                                    type=types.Type.OBJECT,
                                    properties={
                                        "nested_key1": types.Schema(
                                            type=types.Type.STRING
                                        ),
                                        "nested_key2": types.Schema(
                                            type=types.Type.STRING
                                        ),
                                    },
                                ),
                            },
                        ),
                    )
                ]
            )
        ],
    ),
)

FILE_URI_TEST_CASES = [
    pytest.param("gs://bucket/document.pdf", "application/pdf", id="pdf"),
    pytest.param("gs://bucket/data.json", "application/json", id="json"),
    pytest.param("gs://bucket/data.txt", "text/plain", id="txt"),
]

FILE_BYTES_TEST_CASES = [
    pytest.param(
        b"test_pdf_data",
        "application/pdf",
        "data:application/pdf;base64,dGVzdF9wZGZfZGF0YQ==",
        id="pdf",
    ),
    pytest.param(
        b'{"hello":"world"}',
        "application/json",
        "data:application/json;base64,eyJoZWxsbyI6IndvcmxkIn0=",
        id="json",
    ),
]

STREAMING_MODEL_RESPONSE = [
    ModelResponseStream(
        model="test_model",
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    content="zero, ",
                ),
            )
        ],
    ),
    ModelResponseStream(
        model="test_model",
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    content="one, ",
                ),
            )
        ],
    ),
    ModelResponseStream(
        model="test_model",
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    content="two:",
                ),
            )
        ],
    ),
    ModelResponseStream(
        model="test_model",
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id="test_tool_call_id",
                            function=Function(
                                name="test_function",
                                arguments='{"test_arg": "test_',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ],
    ),
    ModelResponseStream(
        model="test_model",
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments='value"}',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ],
    ),
    ModelResponseStream(
        model="test_model",
        choices=[
            StreamingChoices(
                finish_reason="tool_use",
            )
        ],
    ),
]


class _StructuredOutput(BaseModel):
  value: int = Field(description="Value to emit")


class _ModelDumpOnly:
  """Test helper that mimics objects exposing only model_dump."""

  def __init__(self):
    self._schema = {
        "type": "object",
        "properties": {"foo": {"type": "string"}},
    }

  def model_dump(self, *, exclude_none=True, mode="json"):
    # The method signature matches pydantic BaseModel.model_dump to simulate
    # google.genai schema-like objects.
    del exclude_none
    del mode
    return self._schema


async def test_get_completion_inputs_formats_pydantic_schema_for_litellm():
  llm_request = LlmRequest(
      config=types.GenerateContentConfig(response_schema=_StructuredOutput)
  )

  _, _, response_format, _ = await _get_completion_inputs(
      llm_request, model="gemini/gemini-2.5-flash"
  )

  assert response_format == {
      "type": "json_object",
      "response_schema": _StructuredOutput.model_json_schema(),
  }


def test_to_litellm_response_format_passes_preformatted_dict():
  response_format = {
      "type": "json_object",
      "response_schema": {
          "type": "object",
          "properties": {"foo": {"type": "string"}},
      },
  }

  assert (
      _to_litellm_response_format(
          response_format, model="gemini/gemini-2.5-flash"
      )
      == response_format
  )


def test_to_litellm_response_format_wraps_json_schema_dict():
  schema = {
      "type": "object",
      "properties": {"foo": {"type": "string"}},
  }

  formatted = _to_litellm_response_format(
      schema, model="gemini/gemini-2.5-flash"
  )
  assert formatted["type"] == "json_object"
  assert formatted["response_schema"] == schema


def test_to_litellm_response_format_handles_model_dump_object():
  schema_obj = _ModelDumpOnly()

  formatted = _to_litellm_response_format(
      schema_obj, model="gemini/gemini-2.5-flash"
  )

  assert formatted["type"] == "json_object"
  assert formatted["response_schema"] == schema_obj.model_dump()


def test_to_litellm_response_format_handles_genai_schema_instance():
  schema_instance = types.Schema(
      type=types.Type.OBJECT,
      properties={"foo": types.Schema(type=types.Type.STRING)},
      required=["foo"],
  )

  formatted = _to_litellm_response_format(
      schema_instance, model="gemini/gemini-2.5-flash"
  )
  assert formatted["type"] == "json_object"
  assert formatted["response_schema"] == schema_instance.model_dump(
      exclude_none=True, mode="json"
  )


def test_to_litellm_response_format_uses_json_schema_for_openai_model():
  """Test that OpenAI models use json_schema format instead of response_schema."""
  formatted = _to_litellm_response_format(
      _StructuredOutput, model="gpt-4o-mini"
  )

  assert formatted["type"] == "json_schema"
  assert "json_schema" in formatted
  assert formatted["json_schema"]["name"] == "_StructuredOutput"
  assert formatted["json_schema"]["strict"] is True
  assert formatted["json_schema"]["schema"]["additionalProperties"] is False
  assert "additionalProperties" in formatted["json_schema"]["schema"]


def test_to_litellm_response_format_uses_response_schema_for_gemini_model():
  """Test that Gemini models continue to use response_schema format."""
  formatted = _to_litellm_response_format(
      _StructuredOutput, model="gemini/gemini-2.5-flash"
  )

  assert formatted["type"] == "json_object"
  assert "response_schema" in formatted
  assert formatted["response_schema"] == _StructuredOutput.model_json_schema()


def test_to_litellm_response_format_uses_response_schema_for_vertex_gemini():
  """Test that Vertex AI Gemini models use response_schema format."""
  formatted = _to_litellm_response_format(
      _StructuredOutput, model="vertex_ai/gemini-2.5-flash"
  )

  assert formatted["type"] == "json_object"
  assert "response_schema" in formatted
  assert formatted["response_schema"] == _StructuredOutput.model_json_schema()


def test_to_litellm_response_format_uses_json_schema_for_azure_openai():
  """Test that Azure OpenAI models use json_schema format."""
  formatted = _to_litellm_response_format(
      _StructuredOutput, model="azure/gpt-4o"
  )

  assert formatted["type"] == "json_schema"
  assert "json_schema" in formatted
  assert formatted["json_schema"]["name"] == "_StructuredOutput"
  assert formatted["json_schema"]["strict"] is True
  assert formatted["json_schema"]["schema"]["additionalProperties"] is False
  assert "additionalProperties" in formatted["json_schema"]["schema"]


def test_to_litellm_response_format_uses_json_schema_for_anthropic():
  """Test that Anthropic models use json_schema format."""
  formatted = _to_litellm_response_format(
      _StructuredOutput, model="anthropic/claude-3-5-sonnet"
  )

  assert formatted["type"] == "json_schema"
  assert "json_schema" in formatted
  assert formatted["json_schema"]["name"] == "_StructuredOutput"
  assert formatted["json_schema"]["strict"] is True
  assert formatted["json_schema"]["schema"]["additionalProperties"] is False
  assert "additionalProperties" in formatted["json_schema"]["schema"]


def test_to_litellm_response_format_with_dict_schema_for_openai():
  """Test dict schema with OpenAI model uses json_schema format."""
  schema = {
      "type": "object",
      "properties": {"foo": {"type": "string"}},
  }

  formatted = _to_litellm_response_format(schema, model="gpt-4o")

  assert formatted["type"] == "json_schema"
  assert formatted["json_schema"]["name"] == "response"
  assert formatted["json_schema"]["strict"] is True
  assert formatted["json_schema"]["schema"]["additionalProperties"] is False


class _InnerModel(BaseModel):
  value: str = Field(description="A value")
  optional_field: str | None = Field(default=None, description="Optional")


class _OuterModel(BaseModel):
  inner: _InnerModel = Field(description="Nested model")
  name: str


class _WithList(BaseModel):
  items: list[_InnerModel] = Field(description="List of items")
  label: str


def test_enforce_strict_openai_schema_adds_additional_properties_recursively():
  """additionalProperties: false must appear on all object schemas."""
  schema = _OuterModel.model_json_schema()

  _enforce_strict_openai_schema(schema)

  # Root level
  assert schema["additionalProperties"] is False
  # Nested model in $defs
  inner_def = schema["$defs"]["_InnerModel"]
  assert inner_def["additionalProperties"] is False


def test_enforce_strict_openai_schema_marks_all_properties_required():
  """All properties must appear in 'required', including optional fields."""
  schema = _InnerModel.model_json_schema()

  _enforce_strict_openai_schema(schema)

  assert sorted(schema["required"]) == ["optional_field", "value"]


def test_enforce_strict_openai_schema_strips_ref_sibling_keywords():
  """$ref nodes must have no sibling keywords like 'description'."""
  schema = _OuterModel.model_json_schema()
  # Pydantic v2 generates {"$ref": "...", "description": "..."} for nested models
  inner_prop = schema["properties"]["inner"]
  assert "$ref" in inner_prop, "Expected Pydantic to generate a $ref property"
  assert len(inner_prop) > 1, "Expected sibling keywords alongside $ref"

  _enforce_strict_openai_schema(schema)

  inner_prop = schema["properties"]["inner"]
  assert list(inner_prop.keys()) == ["$ref"]


def test_enforce_strict_openai_schema_handles_array_items():
  """Array item schemas should also be recursively transformed."""
  schema = _WithList.model_json_schema()

  _enforce_strict_openai_schema(schema)

  assert schema["additionalProperties"] is False
  inner_def = schema["$defs"]["_InnerModel"]
  assert inner_def["additionalProperties"] is False
  assert sorted(inner_def["required"]) == ["optional_field", "value"]


def test_enforce_strict_openai_schema_preserves_anyof_and_default():
  """anyOf structure and default value for Optional fields must be preserved."""
  schema = _InnerModel.model_json_schema()

  _enforce_strict_openai_schema(schema)

  opt_prop = schema["properties"]["optional_field"]
  assert opt_prop["anyOf"] == [{"type": "string"}, {"type": "null"}]
  assert opt_prop["default"] is None


def test_to_litellm_response_format_dict_input_not_mutated():
  """Passing a raw dict should not mutate the caller's original dict."""
  schema = {
      "type": "object",
      "properties": {
          "nested": {
              "type": "object",
              "properties": {"x": {"type": "string"}},
          }
      },
  }
  import copy

  original = copy.deepcopy(schema)

  _to_litellm_response_format(schema, model="gpt-4o")

  assert schema == original, "Caller's input dict was mutated"


def test_to_litellm_response_format_instance_input_for_openai():
  """Passing a BaseModel instance should produce a valid strict schema."""
  instance = _OuterModel(
      inner=_InnerModel(value="test", optional_field=None), name="foo"
  )

  formatted = _to_litellm_response_format(instance, model="gpt-4o")

  assert formatted["type"] == "json_schema"
  schema = formatted["json_schema"]["schema"]
  assert schema["additionalProperties"] is False
  inner_def = schema["$defs"]["_InnerModel"]
  assert inner_def["additionalProperties"] is False
  assert sorted(inner_def["required"]) == ["optional_field", "value"]


def test_to_litellm_response_format_nested_pydantic_for_openai():
  """Nested Pydantic model should produce a valid OpenAI strict schema."""
  formatted = _to_litellm_response_format(_OuterModel, model="gpt-4o")

  assert formatted["type"] == "json_schema"
  assert formatted["json_schema"]["strict"] is True

  schema = formatted["json_schema"]["schema"]
  assert schema["additionalProperties"] is False
  assert sorted(schema["required"]) == ["inner", "name"]

  # $defs inner model must also be strict
  inner_def = schema["$defs"]["_InnerModel"]
  assert inner_def["additionalProperties"] is False
  assert sorted(inner_def["required"]) == ["optional_field", "value"]


def test_to_litellm_response_format_nested_pydantic_for_gemini_unchanged():
  """Gemini models should NOT get the strict OpenAI transformations."""
  formatted = _to_litellm_response_format(
      _OuterModel, model="gemini/gemini-2.5-flash"
  )

  assert formatted["type"] == "json_object"
  schema = formatted["response_schema"]
  # Gemini path should pass through the raw Pydantic schema untouched
  assert schema == _OuterModel.model_json_schema()


async def test_get_completion_inputs_uses_openai_format_for_openai_model():
  """Test that _get_completion_inputs produces OpenAI-compatible format."""
  llm_request = LlmRequest(
      model="gpt-4o-mini",
      config=types.GenerateContentConfig(response_schema=_StructuredOutput),
  )

  _, _, response_format, _ = await _get_completion_inputs(
      llm_request, model="gpt-4o-mini"
  )

  assert response_format["type"] == "json_schema"
  assert "json_schema" in response_format
  assert response_format["json_schema"]["name"] == "_StructuredOutput"
  assert response_format["json_schema"]["strict"] is True
  assert (
      response_format["json_schema"]["schema"]["additionalProperties"] is False
  )


async def test_get_completion_inputs_uses_gemini_format_for_gemini_model():
  """Test that _get_completion_inputs produces Gemini-compatible format."""
  llm_request = LlmRequest(
      model="gemini/gemini-2.5-flash",
      config=types.GenerateContentConfig(response_schema=_StructuredOutput),
  )

  _, _, response_format, _ = await _get_completion_inputs(
      llm_request, model="gemini/gemini-2.5-flash"
  )

  assert response_format["type"] == "json_object"
  assert "response_schema" in response_format


async def test_get_completion_inputs_uses_passed_model_for_response_format():
  """Test that _get_completion_inputs uses the passed model parameter for response format.

  This verifies that when llm_request.model is None, the explicit model parameter
  is used to determine the correct response format (Gemini vs OpenAI).
  """
  llm_request = LlmRequest(
      model=None,  # No model in request
      config=types.GenerateContentConfig(response_schema=_StructuredOutput),
  )

  # Pass OpenAI model explicitly - should use json_schema format
  _, _, response_format, _ = await _get_completion_inputs(
      llm_request, model="gpt-4o-mini"
  )

  assert response_format["type"] == "json_schema"
  assert "json_schema" in response_format
  assert response_format["json_schema"]["name"] == "_StructuredOutput"
  assert response_format["json_schema"]["strict"] is True
  assert (
      response_format["json_schema"]["schema"]["additionalProperties"] is False
  )


async def test_get_completion_inputs_uses_passed_model_for_gemini_format():
  """Test that _get_completion_inputs uses passed model for Gemini response format.

  This verifies that when self.model is a Gemini model and passed explicitly,
  the response format uses the Gemini-specific format.
  """
  llm_request = LlmRequest(
      model=None,  # No model in request
      config=types.GenerateContentConfig(response_schema=_StructuredOutput),
  )

  # Pass Gemini model explicitly - should use response_schema format
  _, _, response_format, _ = await _get_completion_inputs(
      llm_request, model="gemini/gemini-2.5-flash"
  )

  assert response_format["type"] == "json_object"
  assert "response_schema" in response_format


@pytest.mark.asyncio
async def test_get_completion_inputs_inserts_missing_tool_results():
  user_content = types.Content(
      role="user", parts=[types.Part.from_text(text="Hi")]
  )
  assistant_content = types.Content(
      role="assistant",
      parts=[
          types.Part.from_text(text="Calling tool."),
          types.Part.from_function_call(
              name="get_weather", args={"location": "Seoul"}
          ),
      ],
  )
  assistant_content.parts[1].function_call.id = "tool_call_1"
  followup_user = types.Content(
      role="user", parts=[types.Part.from_text(text="Next question.")]
  )

  llm_request = LlmRequest(
      contents=[user_content, assistant_content, followup_user]
  )
  messages, _, _, _ = await _get_completion_inputs(
      llm_request, model="openai/gpt-4o"
  )

  assert [message["role"] for message in messages] == [
      "user",
      "assistant",
      "tool",
      "user",
  ]
  tool_message = messages[2]
  assert tool_message["tool_call_id"] == "tool_call_1"
  assert tool_message["content"] == _MISSING_TOOL_RESULT_MESSAGE


def test_schema_to_dict_filters_none_enum_values():
  # Use model_construct to bypass strict enum validation.
  top_level_schema = types.Schema.model_construct(
      type=types.Type.STRING,
      enum=["ACTIVE", None, "INACTIVE"],
  )
  nested_schema = types.Schema.model_construct(
      type=types.Type.OBJECT,
      properties={
          "status": types.Schema.model_construct(
              type=types.Type.STRING, enum=["READY", None, "DONE"]
          ),
      },
  )

  assert _schema_to_dict(top_level_schema)["enum"] == ["ACTIVE", "INACTIVE"]
  assert _schema_to_dict(nested_schema)["properties"]["status"]["enum"] == [
      "READY",
      "DONE",
  ]


def test_safe_json_serialize_serializable_object():
  assert _safe_json_serialize({"a": 1, "b": [2, 3]}) == '{"a": 1, "b": [2, 3]}'


def test_safe_json_serialize_non_serializable_object_falls_back_to_str():
  class _NotJsonable:

    def __repr__(self):
      return "<not jsonable>"

  assert _safe_json_serialize(_NotJsonable()) == "<not jsonable>"


def test_safe_json_serialize_circular_dict_falls_back_to_str():
  obj = {}
  obj["self"] = obj
  assert isinstance(_safe_json_serialize(obj), str)


def test_safe_json_serialize_circular_list_falls_back_to_str():
  obj = []
  obj.append(obj)
  assert isinstance(_safe_json_serialize(obj), str)


def test_safe_json_serialize_recursion_error_falls_back_to_str():
  with patch(
      "google.adk.models.lite_llm.json.dumps",
      side_effect=RecursionError("maximum recursion depth"),
  ):
    assert _safe_json_serialize({"a": 1}) == str({"a": 1})


MULTIPLE_FUNCTION_CALLS_STREAM = [
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id="call_1",
                            function=Function(
                                name="function_1",
                                arguments='{"arg": "val',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments='ue1"}',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id="call_2",
                            function=Function(
                                name="function_2",
                                arguments='{"arg": "val',
                            ),
                            index=1,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments='ue2"}',
                            ),
                            index=1,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason="tool_calls",
            )
        ]
    ),
]


STREAM_WITH_EMPTY_CHUNK = [
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id="call_abc",
                            function=Function(
                                name="test_function",
                                arguments='{"test_arg":',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments=' "value"}',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    # This is the problematic empty chunk that should be ignored.
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments="",
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[StreamingChoices(finish_reason="tool_calls", delta=Delta())]
    ),
]


@pytest.fixture
def mock_response():
  return ModelResponse(
      model="test_model",
      choices=[
          Choices(
              message=ChatCompletionAssistantMessage(
                  role="assistant",
                  content="Test response",
                  tool_calls=[
                      ChatCompletionMessageToolCall(
                          type="function",
                          id="test_tool_call_id",
                          function=Function(
                              name="test_function",
                              arguments='{"test_arg": "test_value"}',
                          ),
                      )
                  ],
              )
          )
      ],
  )


# Test case reflecting litellm v1.71.2, ollama v0.9.0 streaming response
# no tool call ids
# indices all 0
# finish_reason stop instead of tool_calls
NON_COMPLIANT_MULTIPLE_FUNCTION_CALLS_STREAM = [
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name="function_1",
                                arguments='{"arg": "val',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments='ue1"}',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name="function_2",
                                arguments='{"arg": "val',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason=None,
                delta=Delta(
                    role="assistant",
                    tool_calls=[
                        ChatCompletionDeltaToolCall(
                            type="function",
                            id=None,
                            function=Function(
                                name=None,
                                arguments='ue2"}',
                            ),
                            index=0,
                        )
                    ],
                ),
            )
        ]
    ),
    ModelResponseStream(
        choices=[
            StreamingChoices(
                finish_reason="stop",
            )
        ]
    ),
]


@pytest.fixture
def mock_acompletion(mock_response):
  return AsyncMock(return_value=mock_response)


@pytest.fixture
def mock_completion(mock_response):
  return Mock(return_value=mock_response)


@pytest.fixture
def mock_client(mock_acompletion, mock_completion):
  return MockLLMClient(mock_acompletion, mock_completion)


@pytest.fixture
def lite_llm_instance(mock_client):
  return LiteLlm(model="test_model", llm_client=mock_client)


class MockLLMClient(LiteLLMClient):

  def __init__(self, acompletion_mock, completion_mock):
    self.acompletion_mock = acompletion_mock
    self.completion_mock = completion_mock

  async def acompletion(self, model, messages, tools, **kwargs):
    if kwargs.get("stream", False):
      kwargs_copy = dict(kwargs)
      kwargs_copy.pop("stream", None)

      async def stream_generator():
        stream_data = self.completion_mock(
            model=model,
            messages=messages,
            tools=tools,
            stream=True,
            **kwargs_copy,
        )
        for item in stream_data:
          yield item

      return stream_generator()
    else:
      return await self.acompletion_mock(
          model=model, messages=messages, tools=tools, **kwargs
      )

  def completion(self, model, messages, tools, stream, **kwargs):
    return self.completion_mock(
        model=model, messages=messages, tools=tools, stream=stream, **kwargs
    )


@pytest.mark.asyncio
async def test_generate_content_async(mock_acompletion, lite_llm_instance):

  async for response in lite_llm_instance.generate_content_async(
      LLM_REQUEST_WITH_FUNCTION_DECLARATION
  ):
    assert response.content.role == "model"
    assert response.content.parts[0].text == "Test response"
    assert response.content.parts[1].function_call.name == "test_function"
    assert response.content.parts[1].function_call.args == {
        "test_arg": "test_value"
    }
    assert response.content.parts[1].function_call.id == "test_tool_call_id"
    assert response.model_version == "test_model"

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  assert kwargs["model"] == "test_model"
  assert kwargs["messages"][0]["role"] == "user"
  assert kwargs["messages"][0]["content"] == "Test prompt"
  assert kwargs["tools"][0]["function"]["name"] == "test_function"
  assert (
      kwargs["tools"][0]["function"]["description"]
      == "Test function description"
  )
  assert (
      kwargs["tools"][0]["function"]["parameters"]["properties"]["test_arg"][
          "type"
      ]
      == "string"
  )


@pytest.mark.asyncio
async def test_generate_content_async_with_model_override(
    mock_acompletion, lite_llm_instance
):
  llm_request = LlmRequest(
      model="overridden_model",
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
  )

  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"
    assert response.content.parts[0].text == "Test response"

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  assert kwargs["model"] == "overridden_model"
  assert kwargs["messages"][0]["role"] == "user"
  assert kwargs["messages"][0]["content"] == "Test prompt"


@pytest.mark.asyncio
async def test_generate_content_async_without_model_override(
    mock_acompletion, lite_llm_instance
):
  llm_request = LlmRequest(
      model=None,
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
  )

  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  assert kwargs["model"] == "test_model"


@pytest.mark.asyncio
async def test_generate_content_async_adds_fallback_user_message(
    mock_acompletion, lite_llm_instance
):
  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user",
              parts=[],
          )
      ]
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  user_messages = [
      message for message in kwargs["messages"] if message["role"] == "user"
  ]
  assert any(
      message.get("content")
      == "Handle the requests as specified in the System Instruction."
      for message in user_messages
  )
  assert (
      sum(1 for content in llm_request.contents if content.role == "user") == 1
  )
  assert llm_request.contents[-1].parts[0].text == (
      "Handle the requests as specified in the System Instruction."
  )


def test_append_fallback_user_content_ignores_function_response_parts():
  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user",
              parts=[
                  types.Part.from_function_response(
                      name="add", response={"result": 6}
                  )
              ],
          )
      ]
  )

  _append_fallback_user_content_if_missing(llm_request)

  assert len(llm_request.contents) == 1
  assert len(llm_request.contents[0].parts) == 1
  assert llm_request.contents[0].parts[0].function_response is not None
  assert llm_request.contents[0].parts[0].text is None


litellm_append_user_content_test_cases = [
    pytest.param(
        LlmRequest(
            contents=[
                types.Content(
                    role="developer",
                    parts=[types.Part.from_text(text="Test prompt")],
                )
            ]
        ),
        2,
        id="litellm request without user content",
    ),
    pytest.param(
        LlmRequest(
            contents=[
                types.Content(
                    role="user",
                    parts=[types.Part.from_text(text="user prompt")],
                )
            ]
        ),
        1,
        id="litellm request with user content",
    ),
    pytest.param(
        LlmRequest(
            contents=[
                types.Content(
                    role="model",
                    parts=[types.Part.from_text(text="model prompt")],
                ),
                types.Content(
                    role="user",
                    parts=[types.Part.from_text(text="user prompt")],
                ),
                types.Content(
                    role="model",
                    parts=[types.Part.from_text(text="model prompt")],
                ),
            ]
        ),
        4,
        id="user content is not the last message scenario",
    ),
]


@pytest.mark.parametrize(
    "llm_request, expected_output", litellm_append_user_content_test_cases
)
def test_maybe_append_user_content(
    lite_llm_instance, llm_request, expected_output
):

  lite_llm_instance._maybe_append_user_content(llm_request)

  assert len(llm_request.contents) == expected_output


function_declaration_test_cases = [
    (
        "simple_function",
        types.FunctionDeclaration(
            name="test_function",
            description="Test function description",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "test_arg": types.Schema(type=types.Type.STRING),
                    "array_arg": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(
                            type=types.Type.STRING,
                        ),
                    ),
                    "nested_arg": types.Schema(
                        type=types.Type.OBJECT,
                        properties={
                            "nested_key1": types.Schema(type=types.Type.STRING),
                            "nested_key2": types.Schema(type=types.Type.STRING),
                        },
                        required=["nested_key1"],
                    ),
                },
                required=["nested_arg"],
            ),
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function",
                "description": "Test function description",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "test_arg": {"type": "string"},
                        "array_arg": {
                            "items": {"type": "string"},
                            "type": "array",
                        },
                        "nested_arg": {
                            "properties": {
                                "nested_key1": {"type": "string"},
                                "nested_key2": {"type": "string"},
                            },
                            "type": "object",
                            "required": ["nested_key1"],
                        },
                    },
                    "required": ["nested_arg"],
                },
            },
        },
    ),
    (
        "no_description",
        types.FunctionDeclaration(
            name="test_function_no_description",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "test_arg": types.Schema(type=types.Type.STRING),
                },
            ),
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function_no_description",
                "description": "",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "test_arg": {"type": "string"},
                    },
                },
            },
        },
    ),
    (
        "empty_parameters",
        types.FunctionDeclaration(
            name="test_function_empty_params",
            parameters=types.Schema(type=types.Type.OBJECT, properties={}),
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function_empty_params",
                "description": "",
                "parameters": {
                    "type": "object",
                    "properties": {},
                },
            },
        },
    ),
    (
        "nested_array",
        types.FunctionDeclaration(
            name="test_function_nested_array",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "array_arg": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(
                            type=types.Type.OBJECT,
                            properties={
                                "nested_key": types.Schema(
                                    type=types.Type.STRING
                                )
                            },
                        ),
                    ),
                },
            ),
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function_nested_array",
                "description": "",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "array_arg": {
                            "items": {
                                "properties": {
                                    "nested_key": {"type": "string"}
                                },
                                "type": "object",
                            },
                            "type": "array",
                        },
                    },
                },
            },
        },
    ),
    (
        "nested_properties",
        types.FunctionDeclaration(
            name="test_function_nested_properties",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "array_arg": types.Schema(
                        type=types.Type.ARRAY,
                        items=types.Schema(
                            type=types.Type.OBJECT,
                            properties={
                                "nested_key": types.Schema(
                                    type=types.Type.OBJECT,
                                    properties={
                                        "inner_key": types.Schema(
                                            type=types.Type.STRING,
                                        )
                                    },
                                )
                            },
                        ),
                    ),
                },
            ),
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function_nested_properties",
                "description": "",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "array_arg": {
                            "items": {
                                "type": "object",
                                "properties": {
                                    "nested_key": {
                                        "type": "object",
                                        "properties": {
                                            "inner_key": {"type": "string"},
                                        },
                                    },
                                },
                            },
                            "type": "array",
                        },
                    },
                },
            },
        },
    ),
    (
        "no_parameters",
        types.FunctionDeclaration(
            name="test_function_no_params",
            description="Test function with no parameters",
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function_no_params",
                "description": "Test function with no parameters",
                "parameters": {
                    "type": "object",
                    "properties": {},
                },
            },
        },
    ),
    (
        "parameters_without_required",
        types.FunctionDeclaration(
            name="test_function_no_required",
            description="Test function with parameters but no required field",
            parameters=types.Schema(
                type=types.Type.OBJECT,
                properties={
                    "optional_arg": types.Schema(type=types.Type.STRING),
                },
            ),
        ),
        {
            "type": "function",
            "function": {
                "name": "test_function_no_required",
                "description": (
                    "Test function with parameters but no required field"
                ),
                "parameters": {
                    "type": "object",
                    "properties": {
                        "optional_arg": {"type": "string"},
                    },
                },
            },
        },
    ),
]


@pytest.mark.parametrize(
    "_, function_declaration, expected_output",
    function_declaration_test_cases,
    ids=[case[0] for case in function_declaration_test_cases],
)
def test_function_declaration_to_tool_param(
    _, function_declaration, expected_output
):
  assert (
      _function_declaration_to_tool_param(function_declaration)
      == expected_output
  )


def test_function_declaration_to_tool_param_without_required_attribute():
  """Ensure tools without a required field attribute don't raise errors."""

  class SchemaWithoutRequired:
    """Mimics a Schema object that lacks the required attribute."""

    def __init__(self):
      self.properties = {
          "optional_arg": types.Schema(type=types.Type.STRING),
      }

  func_decl = types.FunctionDeclaration(
      name="function_without_required_attr",
      description="Function missing required attribute",
  )
  func_decl.parameters = SchemaWithoutRequired()

  expected = {
      "type": "function",
      "function": {
          "name": "function_without_required_attr",
          "description": "Function missing required attribute",
          "parameters": {
              "type": "object",
              "properties": {
                  "optional_arg": {"type": "string"},
              },
          },
      },
  }

  assert _function_declaration_to_tool_param(func_decl) == expected


def test_function_declaration_to_tool_param_with_parameters_json_schema():
  """Ensure function declarations using parameters_json_schema are handled.

  This verifies that when a FunctionDeclaration includes a raw
  `parameters_json_schema` dict, it is used directly as the function
  parameters in the resulting tool param.
  """

  func_decl = types.FunctionDeclaration(
      name="fn_with_json",
      description="desc",
      parameters_json_schema={
          "type": "object",
          "properties": {
              "a": {"type": "string"},
              "b": {"type": "array", "items": {"type": "string"}},
          },
          "required": ["a"],
      },
  )

  expected = {
      "type": "function",
      "function": {
          "name": "fn_with_json",
          "description": "desc",
          "parameters": {
              "type": "object",
              "properties": {
                  "a": {"type": "string"},
                  "b": {"type": "array", "items": {"type": "string"}},
              },
              "required": ["a"],
          },
      },
  }

  assert _function_declaration_to_tool_param(func_decl) == expected


@pytest.mark.asyncio
async def test_generate_content_async_with_system_instruction(
    lite_llm_instance, mock_acompletion
):
  mock_response_with_system_instruction = ModelResponse(
      choices=[
          Choices(
              message=ChatCompletionAssistantMessage(
                  role="assistant",
                  content="Test response",
              )
          )
      ]
  )
  mock_acompletion.return_value = mock_response_with_system_instruction

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          system_instruction="Test system instruction"
      ),
  )

  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"
    assert response.content.parts[0].text == "Test response"

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  assert kwargs["model"] == "test_model"
  assert kwargs["messages"][0]["role"] == "system"
  assert kwargs["messages"][0]["content"] == "Test system instruction"
  assert kwargs["messages"][1]["role"] == "user"
  assert kwargs["messages"][1]["content"] == "Test prompt"


@pytest.mark.asyncio
async def test_generate_content_async_with_tool_response(
    lite_llm_instance, mock_acompletion
):
  mock_response_with_tool_response = ModelResponse(
      choices=[
          Choices(
              message=ChatCompletionAssistantMessage(
                  role="tool",
                  content='{"result": "test_result"}',
                  tool_call_id="test_tool_call_id",
              )
          )
      ]
  )
  mock_acompletion.return_value = mock_response_with_tool_response

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          ),
          types.Content(
              role="tool",
              parts=[
                  types.Part.from_function_response(
                      name="test_function",
                      response={"result": "test_result"},
                  )
              ],
          ),
      ],
      config=types.GenerateContentConfig(
          system_instruction="test instruction",
      ),
  )
  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"
    assert response.content.parts[0].text == '{"result": "test_result"}'

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  assert kwargs["model"] == "test_model"

  assert kwargs["messages"][2]["role"] == "tool"
  assert kwargs["messages"][2]["content"] == '{"result": "test_result"}'


@pytest.mark.asyncio
async def test_generate_content_async_with_usage_metadata(
    lite_llm_instance, mock_acompletion
):
  mock_response_with_usage_metadata = ModelResponse(
      choices=[
          Choices(
              message=ChatCompletionAssistantMessage(
                  role="assistant",
                  content="Test response",
              )
          )
      ],
      usage={
          "prompt_tokens": 10,
          "completion_tokens": 5,
          "total_tokens": 15,
          "cached_tokens": 8,
          "completion_tokens_details": {"reasoning_tokens": 5},
      },
  )
  mock_acompletion.return_value = mock_response_with_usage_metadata

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          ),
      ],
      config=types.GenerateContentConfig(
          system_instruction="test instruction",
      ),
  )
  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"
    assert response.content.parts[0].text == "Test response"
    assert response.usage_metadata.prompt_token_count == 10
    assert response.usage_metadata.candidates_token_count == 5
    assert response.usage_metadata.total_token_count == 15
    assert response.usage_metadata.cached_content_token_count == 8
    assert response.usage_metadata.thoughts_token_count == 5

  mock_acompletion.assert_called_once()


@pytest.mark.asyncio
async def test_generate_content_async_ollama_chat_preserves_multimodal_content(
    mock_acompletion, mock_completion
):
  llm_client = MockLLMClient(mock_acompletion, mock_completion)
  lite_llm_instance = LiteLlm(
      model="ollama_chat/qwen2.5:7b", llm_client=llm_client
  )
  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user",
              parts=[
                  types.Part.from_text(text="Describe this image."),
                  types.Part.from_bytes(
                      data=b"test_image", mime_type="image/png"
                  ),
              ],
          )
      ]
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once_with(
      model="ollama_chat/qwen2.5:7b",
      messages=ANY,
      tools=ANY,
      response_format=ANY,
  )
  _, kwargs = mock_acompletion.call_args
  message_content = kwargs["messages"][0]["content"]
  # Multimodal content (text + image) should be kept as a list so LiteLLM
  # can convert it to Ollama's native images field.
  assert isinstance(message_content, list)
  text_blocks = [
      b
      for b in message_content
      if isinstance(b, dict) and b.get("type") == "text"
  ]
  image_blocks = [
      b
      for b in message_content
      if isinstance(b, dict) and b.get("type") == "image_url"
  ]
  assert len(text_blocks) >= 1
  assert "Describe this image." in text_blocks[0].get("text", "")
  assert len(image_blocks) >= 1


@pytest.mark.asyncio
async def test_generate_content_async_custom_provider_preserves_multimodal(
    mock_acompletion, mock_completion
):
  llm_client = MockLLMClient(mock_acompletion, mock_completion)
  lite_llm_instance = LiteLlm(
      model="qwen2.5:7b",
      llm_client=llm_client,
      custom_llm_provider="ollama_chat",
  )
  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user",
              parts=[
                  types.Part.from_text(text="Describe this image."),
                  types.Part.from_bytes(
                      data=b"test_image", mime_type="image/png"
                  ),
              ],
          )
      ]
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert kwargs["custom_llm_provider"] == "ollama_chat"
  assert kwargs["model"] == "qwen2.5:7b"
  message_content = kwargs["messages"][0]["content"]
  # Multimodal content should be preserved as a list.
  assert isinstance(message_content, list)
  text_blocks = [
      b
      for b in message_content
      if isinstance(b, dict) and b.get("type") == "text"
  ]
  assert any("Describe this image." in b.get("text", "") for b in text_blocks)


def test_flatten_ollama_content_accepts_tuple_blocks():
  from google.adk.models.lite_llm import _flatten_ollama_content

  content = (
      {"type": "text", "text": "first"},
      {"type": "text", "text": "second"},
  )
  flattened = _flatten_ollama_content(content)
  assert flattened == "first\nsecond"


@pytest.mark.parametrize(
    "content, expected",
    [
        (None, None),
        ("hello", "hello"),
        (
            [
                {"type": "text", "text": "first"},
                {"type": "text", "text": "second"},
            ],
            "first\nsecond",
        ),
    ],
)
def test_flatten_ollama_content_returns_str_or_none(content, expected):
  from google.adk.models.lite_llm import _flatten_ollama_content

  flattened = _flatten_ollama_content(content)
  assert flattened == expected
  assert flattened is None or isinstance(flattened, str)


def test_flatten_ollama_content_preserves_image_url_blocks():
  """Media blocks should be kept as a list so LiteLLM can convert them."""
  from google.adk.models.lite_llm import _flatten_ollama_content

  blocks = [
      {"type": "image_url", "image_url": {"url": "http://example.com/img.png"}},
  ]
  result = _flatten_ollama_content(blocks)
  assert isinstance(result, list)
  assert result == blocks


def test_flatten_ollama_content_preserves_mixed_text_and_image():
  """Text + image_url should return the full list, not just the text."""
  from google.adk.models.lite_llm import _flatten_ollama_content

  blocks = [
      {"type": "text", "text": "Describe this image."},
      {
          "type": "image_url",
          "image_url": {"url": "data:image/png;base64,iVBORw0KGgo="},
      },
  ]
  result = _flatten_ollama_content(blocks)
  assert isinstance(result, list)
  assert len(result) == 2
  assert result[0]["type"] == "text"
  assert result[1]["type"] == "image_url"


def test_flatten_ollama_content_preserves_video_url_blocks():
  from google.adk.models.lite_llm import _flatten_ollama_content

  blocks = [
      {"type": "text", "text": "What happens in this clip?"},
      {"type": "video_url", "video_url": {"url": "http://example.com/v.mp4"}},
  ]
  result = _flatten_ollama_content(blocks)
  assert isinstance(result, list)
  assert len(result) == 2


def test_flatten_ollama_content_serializes_non_media_non_text_blocks_to_json():
  """Blocks with unknown types and no media should still serialize to JSON."""
  from google.adk.models.lite_llm import _flatten_ollama_content

  blocks = [
      {"type": "custom_block", "data": "something"},
  ]
  result = _flatten_ollama_content(blocks)
  assert isinstance(result, str)
  assert json.loads(result) == blocks


def test_flatten_ollama_content_serializes_dict_to_json():
  from google.adk.models.lite_llm import _flatten_ollama_content

  content = {"type": "image_url", "image_url": {"url": "http://example.com"}}
  flattened = _flatten_ollama_content(content)
  assert isinstance(flattened, str)
  assert json.loads(flattened) == content


@pytest.mark.asyncio
async def test_content_to_message_param_user_message():
  content = types.Content(
      role="user", parts=[types.Part.from_text(text="Test prompt")]
  )
  message = await _content_to_message_param(content)
  assert message["role"] == "user"
  assert message["content"] == "Test prompt"


@pytest.mark.asyncio
@pytest.mark.parametrize("file_uri,mime_type", FILE_URI_TEST_CASES)
async def test_content_to_message_param_user_message_with_file_uri(
    file_uri, mime_type
):
  file_part = types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)
  content = types.Content(
      role="user",
      parts=[
          types.Part.from_text(text="Summarize this file."),
          file_part,
      ],
  )

  message = await _content_to_message_param(content)
  assert message == {
      "role": "user",
      "content": [
          {"type": "text", "text": "Summarize this file."},
          {"type": "file", "file": {"file_id": file_uri, "format": mime_type}},
      ],
  }


@pytest.mark.asyncio
@pytest.mark.parametrize("file_uri,mime_type", FILE_URI_TEST_CASES)
async def test_content_to_message_param_user_message_file_uri_only(
    file_uri, mime_type
):
  file_part = types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)
  content = types.Content(
      role="user",
      parts=[
          file_part,
      ],
  )

  message = await _content_to_message_param(content)
  assert message == {
      "role": "user",
      "content": [
          {"type": "file", "file": {"file_id": file_uri, "format": mime_type}},
      ],
  }


@pytest.mark.asyncio
async def test_content_to_message_param_user_message_file_uri_without_mime_type():
  """Test that file_data without an inferable mime_type raises ValueError.

  When using GcsArtifactService, artifacts may have file_uri (gs://...) but
  without mime_type set. When the MIME type cannot be determined from the URI
  extension or display_name, ADK raises a clear ValueError rather than
  forwarding an unsupported 'application/octet-stream' to LiteLLM.
  """
  file_part = types.Part(
      file_data=types.FileData(
          file_uri="gs://agent-artifact-bucket/app/user/session/artifact/0"
      )
  )
  content = types.Content(
      role="user",
      parts=[
          types.Part.from_text(text="Analyze this file."),
          file_part,
      ],
  )

  with pytest.raises(ValueError, match="Cannot process file_uri"):
    await _content_to_message_param(content)


@pytest.mark.asyncio
async def test_content_to_message_param_user_message_file_uri_explicit_octet_stream():
  """Test that an explicit application/octet-stream MIME type raises ValueError.

  Upstream callers may explicitly set mime_type to 'application/octet-stream'
  when the true type is unknown. ADK treats this identically to a missing MIME
  type and raises early rather than forwarding the unsupported type to LiteLLM.
  """
  file_part = types.Part(
      file_data=types.FileData(
          file_uri="gs://agent-artifact-bucket/app/user/session/artifact/0",
          mime_type="application/octet-stream",
      )
  )
  content = types.Content(
      role="user",
      parts=[
          types.Part.from_text(text="Analyze this file."),
          file_part,
      ],
  )

  with pytest.raises(ValueError, match="application/octet-stream"):
    await _content_to_message_param(content)


@pytest.mark.asyncio
async def test_content_to_message_param_user_message_file_uri_infer_mime_type():
  """Test MIME type inference from file_uri extension.

  When file_data has a file_uri with a recognizable extension but no explicit
  mime_type, the MIME type should be inferred from the extension.
  """
  file_part = types.Part(
      file_data=types.FileData(
          file_uri="gs://bucket/path/to/document.pdf",
      )
  )
  content = types.Content(
      role="user",
      parts=[file_part],
  )

  message = await _content_to_message_param(content)
  assert message == {
      "role": "user",
      "content": [
          {
              "type": "file",
              "file": {
                  "file_id": "gs://bucket/path/to/document.pdf",
                  "format": "application/pdf",
              },
          },
      ],
  }


@pytest.mark.asyncio
async def test_content_to_message_param_multi_part_function_response():
  part1 = types.Part.from_function_response(
      name="function_one",
      response={"result": "result_one"},
  )
  part1.function_response.id = "tool_call_1"

  part2 = types.Part.from_function_response(
      name="function_two",
      response={"value": 123},
  )
  part2.function_response.id = "tool_call_2"

  content = types.Content(
      role="tool",
      parts=[part1, part2],
  )
  messages = await _content_to_message_param(content)
  assert isinstance(messages, list)
  assert len(messages) == 2

  assert messages[0]["role"] == "tool"
  assert messages[0]["tool_call_id"] == "tool_call_1"
  assert messages[0]["content"] == '{"result": "result_one"}'

  assert messages[1]["role"] == "tool"
  assert messages[1]["tool_call_id"] == "tool_call_2"
  assert messages[1]["content"] == '{"value": 123}'


@pytest.mark.asyncio
async def test_content_to_message_param_function_response_with_extra_parts():
  tool_part = types.Part.from_function_response(
      name="load_image",
      response={"status": "success"},
  )
  tool_part.function_response.id = "tool_call_1"

  text_part = types.Part.from_text(text="[Image: img_123.png]")
  image_bytes = b"test_image_data"
  image_part = types.Part.from_bytes(data=image_bytes, mime_type="image/png")

  content = types.Content(
      role="user",
      parts=[tool_part, text_part, image_part],
  )

  messages = await _content_to_message_param(content)
  assert isinstance(messages, list)
  assert messages == [
      {
          "role": "tool",
          "tool_call_id": "tool_call_1",
          "content": '{"status": "success"}',
      },
      {
          "role": "user",
          "content": [
              {"type": "text", "text": "[Image: img_123.png]"},
              {
                  "type": "image_url",
                  "image_url": {
                      "url": "data:image/png;base64,dGVzdF9pbWFnZV9kYXRh"
                  },
              },
          ],
      },
  ]


@pytest.mark.asyncio
async def test_content_to_message_param_function_response_preserves_string():
  """Tests that string responses are used directly without double-serialization.

  The google.genai FunctionResponse.response field is typed as dict, but
  _content_to_message_param defensively handles string responses to avoid
  double-serialization. This test verifies that behavior by mocking a
  function_response with a string response attribute.
  """
  response_payload = '{"type": "files", "count": 2}'

  # Create a Part with a dict response, then mock the response to be a string
  # to simulate edge cases where response might be set directly as a string
  part = types.Part.from_function_response(
      name="list_files",
      response={"placeholder": "will be mocked"},
  )

  # Mock the response attribute to return a string
  # Using Mock without spec_set to allow setting response to a string,
  # which simulates the edge case we're testing
  mock_function_response = Mock(spec=types.FunctionResponse)
  mock_function_response.response = response_payload
  mock_function_response.id = "tool_call_1"
  part.function_response = mock_function_response

  content = types.Content(
      role="tool",
      parts=[part],
  )
  message = await _content_to_message_param(content)

  assert message["role"] == "tool"
  assert message["tool_call_id"] == "tool_call_1"
  assert message["content"] == response_payload


@pytest.mark.asyncio
async def test_content_to_message_param_assistant_message():
  content = types.Content(
      role="assistant", parts=[types.Part.from_text(text="Test response")]
  )
  message = await _content_to_message_param(content)
  assert message["role"] == "assistant"
  assert message["content"] == "Test response"


@pytest.mark.asyncio
async def test_content_to_message_param_user_filters_thought_parts():
  thought_part = types.Part.from_text(text="internal reasoning")
  thought_part.thought = True
  content_part = types.Part.from_text(text="visible content")
  content = types.Content(role="user", parts=[thought_part, content_part])

  message = await _content_to_message_param(content)

  assert message["role"] == "user"
  assert message["content"] == "visible content"


@pytest.mark.asyncio
async def test_content_to_message_param_assistant_thought_message():
  part = types.Part.from_text(text="internal reasoning")
  part.thought = True
  content = types.Content(role="assistant", parts=[part])

  message = await _content_to_message_param(content)

  assert message["role"] == "assistant"
  assert message["content"] is None
  assert message["reasoning_content"] == "internal reasoning"


@pytest.mark.asyncio
async def test_content_to_message_param_merges_reasoning_chunks_without_separator():
  first_part = types.Part.from_text(text="Let")
  first_part.thought = True
  second_part = types.Part.from_text(text=" me think")
  second_part.thought = True
  third_part = types.Part.from_text(text=" this through.")
  third_part.thought = True
  content = types.Content(
      role="assistant", parts=[first_part, second_part, third_part]
  )

  message = await _content_to_message_param(content)

  assert message["role"] == "assistant"
  assert message["content"] is None
  assert message["reasoning_content"] == "Let me think this through."


@pytest.mark.asyncio
async def test_content_to_message_param_model_thought_message():
  part = types.Part.from_text(text="internal reasoning")
  part.thought = True
  content = types.Content(role="model", parts=[part])

  message = await _content_to_message_param(content)

  assert message["role"] == "assistant"
  assert message["content"] is None
  assert message["reasoning_content"] == "internal reasoning"


@pytest.mark.asyncio
async def test_content_to_message_param_assistant_thought_and_content_message():
  thought_part = types.Part.from_text(text="internal reasoning")
  thought_part.thought = True
  content_part = types.Part.from_text(text="visible content")
  content = types.Content(role="assistant", parts=[thought_part, content_part])

  message = await _content_to_message_param(content)

  assert message["role"] == "assistant"
  assert message["content"] == "visible content"
  assert message["reasoning_content"] == "internal reasoning"


@pytest.mark.asyncio
async def test_content_to_message_param_preserves_chunked_reasoning_deltas():
  thought_part_1 = types.Part.from_text(text="Hel")
  thought_part_1.thought = True
  thought_part_2 = types.Part.from_text(text="lo")
  thought_part_2.thought = True
  content = types.Content(
      role="assistant", parts=[thought_part_1, thought_part_2]
  )

  message = await _content_to_message_param(content)

  assert message["role"] == "assistant"
  assert message["content"] is None
  assert message["reasoning_content"] == "Hello"


@pytest.mark.asyncio
async def test_content_to_message_param_preserves_reasoning_newlines():
  thought_part_1 = types.Part.from_text(text="line 1\n")
  thought_part_1.thought = True
  thought_part_2 = types.Part.from_text(text="line 2")
  thought_part_2.thought = True
  content = types.Content(
      role="assistant", parts=[thought_part_1, thought_part_2]
  )

  message = await _content_to_message_param(content)

  assert message["reasoning_content"] == "line 1\nline 2"


@pytest.mark.asyncio
async def test_content_to_message_param_function_call():
  content = types.Content(
      role="assistant",
      parts=[
          types.Part.from_text(text="test response"),
          types.Part.from_function_call(
              name="test_function", args={"test_arg": "test_value"}
          ),
      ],
  )
  content.parts[1].function_call.id = "test_tool_call_id"
  message = await _content_to_message_param(content)
  assert message["role"] == "assistant"
  assert message["content"] == "test response"

  tool_call = message["tool_calls"][0]
  assert tool_call["type"] == "function"
  assert tool_call["id"] == "test_tool_call_id"
  assert tool_call["function"]["name"] == "test_function"
  assert tool_call["function"]["arguments"] == '{"test_arg": "test_value"}'


@pytest.mark.asyncio
async def test_content_to_message_param_multipart_content():
  """Test handling of multipart content where final_content is a list with text objects."""
  content = types.Content(
      role="assistant",
      parts=[
          types.Part.from_text(text="text part"),
          types.Part.from_bytes(data=b"test_image_data", mime_type="image/png"),
      ],
  )
  message = await _content_to_message_param(content)
  assert message["role"] == "assistant"
  # When content is a list and the first element is a text object with type "text",
  # it should extract the text (for providers like ollama_chat that don't handle lists well)
  # This is the behavior implemented in the fix
  assert message["content"] == "text part"
  assert message["tool_calls"] is None


@pytest.mark.asyncio
async def test_content_to_message_param_single_text_object_in_list(mocker):
  """Test extraction of text from single text object in list (for ollama_chat compatibility)."""
  from google.adk.models import lite_llm

  # Mock _get_content to return a list with single text object
  async def mock_get_content(*args, **kwargs):
    return [{"type": "text", "text": "single text"}]

  mocker.patch.object(lite_llm, "_get_content", side_effect=mock_get_content)

  content = types.Content(
      role="assistant",
      parts=[types.Part.from_text(text="single text")],
  )
  message = await _content_to_message_param(content)
  assert message["role"] == "assistant"
  # Should extract the text from the single text object
  assert message["content"] == "single text"
  assert message["tool_calls"] is None


def test_message_to_generate_content_response_text():
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content="Test response",
  )
  response = _message_to_generate_content_response(message)
  assert response.content.role == "model"
  assert response.content.parts[0].text == "Test response"


def test_message_to_generate_content_response_tool_call():
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content=None,
      tool_calls=[
          ChatCompletionMessageToolCall(
              type="function",
              id="test_tool_call_id",
              function=Function(
                  name="test_function",
                  arguments='{"test_arg": "test_value"}',
              ),
          )
      ],
  )

  response = _message_to_generate_content_response(message)
  assert response.content.role == "model"
  assert response.content.parts[0].function_call.name == "test_function"
  assert response.content.parts[0].function_call.args == {
      "test_arg": "test_value"
  }
  assert response.content.parts[0].function_call.id == "test_tool_call_id"


def test_message_to_generate_content_response_tool_call_accepts_python_literal_arguments():
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content=None,
      tool_calls=[
          ChatCompletionMessageToolCall(
              type="function",
              id="test_tool_call_id",
              function=Function(
                  name="test_function",
                  arguments="{'query': 'MATCH (n) RETURN n'}",
              ),
          )
      ],
  )

  response = _message_to_generate_content_response(message)

  assert response.content.role == "model"
  assert response.content.parts[0].function_call.name == "test_function"
  assert response.content.parts[0].function_call.args == {
      "query": "MATCH (n) RETURN n"
  }


def test_message_to_generate_content_response_tool_call_accepts_unquoted_json_keys():
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content=None,
      tool_calls=[
          ChatCompletionMessageToolCall(
              type="function",
              id="test_tool_call_id",
              function=Function(
                  name="test_function",
                  arguments='{query: "MATCH (n) RETURN n", limit: 5}',
              ),
          )
      ],
  )

  response = _message_to_generate_content_response(message)

  assert response.content.role == "model"
  assert response.content.parts[0].function_call.name == "test_function"
  assert response.content.parts[0].function_call.args == {
      "query": "MATCH (n) RETURN n",
      "limit": 5,
  }


def test_message_to_generate_content_response_inline_tool_call_text():
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content=(
          '{"id":"inline_call","name":"get_current_time",'
          '"arguments":{"timezone_str":"Asia/Taipei"}} <|im_end|>system'
      ),
  )

  response = _message_to_generate_content_response(message)
  assert len(response.content.parts) == 2
  text_part = response.content.parts[0]
  tool_part = response.content.parts[1]
  assert text_part.text == "<|im_end|>system"
  assert tool_part.function_call.name == "get_current_time"
  assert tool_part.function_call.args == {"timezone_str": "Asia/Taipei"}
  assert tool_part.function_call.id == "inline_call"


def test_message_to_generate_content_response_with_model():
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content="Test response",
  )
  response = _message_to_generate_content_response(
      message, model_version="gemini-2.5-pro"
  )
  assert response.content.role == "model"
  assert response.content.parts[0].text == "Test response"
  assert response.model_version == "gemini-2.5-pro"


def test_message_to_generate_content_response_reasoning_content():
  message = {
      "role": "assistant",
      "content": "Visible text",
      "reasoning_content": "Hidden chain",
  }
  response = _message_to_generate_content_response(message)

  assert len(response.content.parts) == 2
  thought_part = response.content.parts[0]
  text_part = response.content.parts[1]
  assert thought_part.text == "Hidden chain"
  assert thought_part.thought is True
  assert text_part.text == "Visible text"


def test_model_response_to_generate_content_response_reasoning_content():
  model_response = ModelResponse(
      model="thinking-model",
      choices=[{
          "message": {
              "role": "assistant",
              "content": "Answer",
              "reasoning_content": "Step-by-step",
          },
          "finish_reason": "stop",
      }],
  )

  response = _model_response_to_generate_content_response(model_response)

  assert response.content.parts[0].text == "Step-by-step"
  assert response.content.parts[0].thought is True
  assert response.content.parts[1].text == "Answer"


def test_message_to_generate_content_response_reasoning_field():
  """Test that the 'reasoning' field is supported (LM Studio, vLLM)."""
  message = {
      "role": "assistant",
      "content": "Final answer",
      "reasoning": "Thinking process",
  }
  response = _message_to_generate_content_response(message)

  assert len(response.content.parts) == 2
  thought_part = response.content.parts[0]
  text_part = response.content.parts[1]
  assert thought_part.text == "Thinking process"
  assert thought_part.thought is True
  assert text_part.text == "Final answer"


def test_model_response_to_generate_content_response_reasoning_field():
  """Test that 'reasoning' field is supported in ModelResponse."""
  model_response = ModelResponse(
      model="test-model",
      choices=[{
          "message": {
              "role": "assistant",
              "content": "Result",
              "reasoning": "Chain of thought",
          },
          "finish_reason": "stop",
      }],
  )

  response = _model_response_to_generate_content_response(model_response)

  assert response.content.parts[0].text == "Chain of thought"
  assert response.content.parts[0].thought is True
  assert response.content.parts[1].text == "Result"


def test_model_response_to_generate_content_response_grounding_metadata_dict():
  """vertex_ai_grounding_metadata as a dict is propagated to the LlmResponse."""
  model_response = ModelResponse(
      model="gemini/gemini-2.5-flash",
      choices=[{
          "message": {"role": "assistant", "content": "Answer"},
          "finish_reason": "stop",
      }],
  )
  model_response.vertex_ai_grounding_metadata = {
      "grounding_chunks": [
          {"web": {"uri": "https://example.com", "title": "Example"}}
      ],
  }

  response = _model_response_to_generate_content_response(model_response)

  assert response.grounding_metadata is not None
  assert (
      response.grounding_metadata.grounding_chunks[0].web.uri
      == "https://example.com"
  )


def test_model_response_to_generate_content_response_grounding_metadata_list():
  """LiteLLM may emit a list (per candidate); the first entry is used."""
  model_response = ModelResponse(
      model="gemini/gemini-2.5-flash",
      choices=[{
          "message": {"role": "assistant", "content": "Answer"},
          "finish_reason": "stop",
      }],
  )
  model_response.vertex_ai_grounding_metadata = [
      {"grounding_chunks": [{"web": {"uri": "https://a.test", "title": "A"}}]},
      {"grounding_chunks": [{"web": {"uri": "https://b.test", "title": "B"}}]},
  ]

  response = _model_response_to_generate_content_response(model_response)

  assert response.grounding_metadata is not None
  assert (
      response.grounding_metadata.grounding_chunks[0].web.uri
      == "https://a.test"
  )


def test_model_response_to_generate_content_response_no_grounding_metadata():
  """Without vertex_ai_grounding_metadata, grounding_metadata stays None."""
  model_response = ModelResponse(
      model="gemini/gemini-2.5-flash",
      choices=[{
          "message": {"role": "assistant", "content": "Answer"},
          "finish_reason": "stop",
      }],
  )

  response = _model_response_to_generate_content_response(model_response)

  assert response.grounding_metadata is None


def test_reasoning_content_takes_precedence_over_reasoning():
  """Test that 'reasoning_content' is prioritized over 'reasoning'."""
  message = {
      "role": "assistant",
      "content": "Answer",
      "reasoning_content": "LiteLLM standard reasoning",
      "reasoning": "Alternative reasoning",
  }
  response = _message_to_generate_content_response(message)

  assert len(response.content.parts) == 2
  thought_part = response.content.parts[0]
  assert thought_part.text == "LiteLLM standard reasoning"
  assert thought_part.thought is True


def test_extract_reasoning_value_from_reasoning_content():
  """Test extraction from reasoning_content (LiteLLM standard)."""
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content="Answer",
      reasoning_content="LiteLLM reasoning",
  )
  result = _extract_reasoning_value(message)
  assert result == "LiteLLM reasoning"


def test_extract_reasoning_value_from_reasoning():
  """Test extraction from reasoning (LM Studio, vLLM)."""

  class MockMessage:

    def __init__(self):
      self.role = "assistant"
      self.content = "Answer"
      self.reasoning = "Alternative reasoning"

    def get(self, key, default=None):
      return getattr(self, key, default)

  message = MockMessage()
  result = _extract_reasoning_value(message)
  assert result == "Alternative reasoning"


def test_extract_reasoning_value_dict_reasoning_content():
  """Test extraction from dict with reasoning_content field."""
  message = {
      "role": "assistant",
      "content": "Answer",
      "reasoning_content": "Dict reasoning content",
  }
  result = _extract_reasoning_value(message)
  assert result == "Dict reasoning content"


def test_extract_reasoning_value_dict_reasoning():
  """Test extraction from dict with reasoning field."""
  message = {
      "role": "assistant",
      "content": "Answer",
      "reasoning": "Dict reasoning",
  }
  result = _extract_reasoning_value(message)
  assert result == "Dict reasoning"


def test_extract_reasoning_value_dict_prefers_reasoning_content():
  """Test that reasoning_content takes precedence over reasoning in dicts."""
  message = {
      "role": "assistant",
      "content": "Answer",
      "reasoning_content": "Primary",
      "reasoning": "Secondary",
  }
  result = _extract_reasoning_value(message)
  assert result == "Primary"


def test_extract_reasoning_value_none_message():
  """Test that None message returns None."""
  result = _extract_reasoning_value(None)
  assert result is None


def test_extract_reasoning_value_no_reasoning_fields():
  """Test that None is returned when no reasoning fields exist."""
  message = {
      "role": "assistant",
      "content": "Answer only",
  }
  result = _extract_reasoning_value(message)
  assert result is None


def test_extract_thought_signature_from_extra_content():
  """Extracts thought_signature from extra_content (OpenAI-compatible path)."""
  sig_b64 = base64.b64encode(b"test_signature").decode("utf-8")
  tc = ChatCompletionMessageToolCall(
      type="function",
      id="call_123",
      function=Function(name="test_fn", arguments="{}"),
      extra_content={"google": {"thought_signature": sig_b64}},
  )
  result = _extract_thought_signature_from_tool_call(tc)
  assert result == b"test_signature"


def test_extract_thought_signature_from_provider_specific_fields():
  """Extracts thought_signature from provider_specific_fields (Vertex path)."""
  sig_b64 = base64.b64encode(b"vertex_sig").decode("utf-8")
  tc = ChatCompletionMessageToolCall(
      type="function",
      id="call_456",
      function=Function(name="test_fn", arguments="{}"),
      provider_specific_fields={"thought_signature": sig_b64},
  )
  result = _extract_thought_signature_from_tool_call(tc)
  assert result == b"vertex_sig"


def test_extract_thought_signature_from_function_provider_fields():
  """Extracts thought_signature from function's provider_specific_fields.

  When provider_specific_fields is set directly on the function object
  (e.g. by litellm internals), the extraction should find it.
  """
  sig_b64 = base64.b64encode(b"func_sig").decode("utf-8")
  tc = ChatCompletionMessageToolCall(
      type="function",
      id="call_func",
      function=Function(name="test_fn", arguments="{}"),
  )
  # Simulate litellm setting provider_specific_fields on the function
  tc.function.provider_specific_fields = {
      "thought_signature": sig_b64,
  }
  result = _extract_thought_signature_from_tool_call(tc)
  assert result == b"func_sig"


def test_extract_thought_signature_from_id():
  """Extracts thought_signature from tool call ID (__thought__ separator)."""
  sig_b64 = base64.b64encode(b"id_sig").decode("utf-8")
  tc = ChatCompletionMessageToolCall(
      type="function",
      id=f"call_789{_THOUGHT_SIGNATURE_SEPARATOR}{sig_b64}",
      function=Function(name="test_fn", arguments="{}"),
  )
  result = _extract_thought_signature_from_tool_call(tc)
  assert result == b"id_sig"


def test_extract_thought_signature_returns_none_when_absent():
  """Returns None when no thought_signature is present."""
  tc = ChatCompletionMessageToolCall(
      type="function",
      id="call_plain",
      function=Function(name="test_fn", arguments="{}"),
  )
  result = _extract_thought_signature_from_tool_call(tc)
  assert result is None


def test_extract_thought_signature_corrupted_base64_returns_none():
  """Returns None gracefully for corrupted base64 signatures."""
  tc = ChatCompletionMessageToolCall(
      type="function",
      id="call_bad",
      function=Function(name="test_fn", arguments="{}"),
      extra_content={"google": {"thought_signature": "!!!not_valid_base64!!!"}},
  )
  result = _extract_thought_signature_from_tool_call(tc)
  assert result is None


def test_message_to_generate_content_response_preserves_thought_signature():
  """thought_signature from tool call is preserved on the output Part."""
  sig_b64 = base64.b64encode(b"round_trip_sig").decode("utf-8")
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content=None,
      tool_calls=[
          ChatCompletionMessageToolCall(
              type="function",
              id="call_ts_1",
              function=Function(
                  name="load_skill",
                  arguments='{"skill": "my_skill"}',
              ),
              extra_content={"google": {"thought_signature": sig_b64}},
          )
      ],
  )

  response = _message_to_generate_content_response(message)
  fc_part = response.content.parts[0]
  assert fc_part.function_call.name == "load_skill"
  assert fc_part.function_call.id == "call_ts_1"
  assert fc_part.thought_signature == b"round_trip_sig"


def test_message_to_generate_content_response_no_thought_signature():
  """Parts without thought_signature have thought_signature=None."""
  message = ChatCompletionAssistantMessage(
      role="assistant",
      content=None,
      tool_calls=[
          ChatCompletionMessageToolCall(
              type="function",
              id="call_no_ts",
              function=Function(
                  name="plain_tool",
                  arguments="{}",
              ),
          )
      ],
  )

  response = _message_to_generate_content_response(message)
  fc_part = response.content.parts[0]
  assert fc_part.function_call.name == "plain_tool"
  assert fc_part.thought_signature is None


@pytest.mark.asyncio
async def test_content_to_message_param_preserves_thought_signature():
  """thought_signature on Part is emitted on both tool call metadata paths."""
  sig_bytes = b"preserved_sig"
  sig_b64 = base64.b64encode(sig_bytes).decode("utf-8")
  content = types.Content(
      role="model",
      parts=[
          types.Part(
              function_call=types.FunctionCall(
                  name="load_skill",
                  args={"skill": "my_skill"},
                  id="call_rt",
              ),
              thought_signature=sig_bytes,
          ),
      ],
  )

  message = await _content_to_message_param(content)
  assert message["role"] == "assistant"
  tc = message["tool_calls"][0]
  assert tc["function"]["name"] == "load_skill"
  assert tc["id"] == "call_rt"
  assert tc["provider_specific_fields"] == {"thought_signature": sig_b64}
  assert tc["extra_content"] == {"google": {"thought_signature": sig_b64}}


@pytest.mark.asyncio
async def test_content_to_message_param_no_thought_signature():
  """Tool calls without thought_signature have no signature metadata."""
  content = types.Content(
      role="model",
      parts=[
          types.Part.from_function_call(name="plain_tool", args={"key": "val"}),
      ],
  )
  content.parts[0].function_call.id = "call_plain"

  message = await _content_to_message_param(content)
  tc = message["tool_calls"][0]
  assert tc["id"] == "call_plain"
  assert "provider_specific_fields" not in tc
  assert "extra_content" not in tc


@pytest.mark.asyncio
async def test_thought_signature_round_trip():
  """thought_signature survives a full round trip through ADK conversions.

  Simulates the flow: litellm response → types.Part → litellm request.
  """
  sig_b64 = base64.b64encode(b"full_round_trip").decode("utf-8")

  # Step 1: Incoming litellm message with thought_signature
  incoming_message = ChatCompletionAssistantMessage(
      role="assistant",
      content=None,
      tool_calls=[
          ChatCompletionMessageToolCall(
              type="function",
              id="call_round",
              function=Function(
                  name="load_skill",
                  arguments='{"skill_name": "test"}',
              ),
              extra_content={"google": {"thought_signature": sig_b64}},
          )
      ],
  )

  # Step 2: Convert to ADK internal format (types.Content)
  llm_response = _message_to_generate_content_response(incoming_message)
  fc_part = llm_response.content.parts[0]
  assert fc_part.thought_signature == b"full_round_trip"

  # Step 3: Convert back to litellm format
  outgoing_message = await _content_to_message_param(llm_response.content)
  out_tc = outgoing_message["tool_calls"][0]
  assert out_tc["provider_specific_fields"] == {"thought_signature": sig_b64}
  assert out_tc["extra_content"] == {"google": {"thought_signature": sig_b64}}


def test_parse_tool_calls_from_text_multiple_calls():
  text = (
      '{"name":"alpha","arguments":{"value":1}}\n'
      "Some filler text "
      '{"id":"custom","name":"beta","arguments":{"timezone":"Asia/Taipei"}} '
      "ignored suffix"
  )
  tool_calls, remainder = _parse_tool_calls_from_text(text)
  assert len(tool_calls) == 2
  assert tool_calls[0].function.name == "alpha"
  assert json.loads(tool_calls[0].function.arguments) == {"value": 1}
  assert tool_calls[1].id == "custom"
  assert tool_calls[1].function.name == "beta"
  assert json.loads(tool_calls[1].function.arguments) == {
      "timezone": "Asia/Taipei"
  }
  assert remainder == "Some filler text  ignored suffix"


def test_parse_tool_calls_from_text_invalid_json_returns_remainder():
  text = 'Leading {"unused": "payload"} trailing text'
  tool_calls, remainder = _parse_tool_calls_from_text(text)
  assert tool_calls == []
  assert remainder == 'Leading {"unused": "payload"} trailing text'


# ---------------------------------------------------------------------------
# DeepSeek proprietary inline tool-call format tests
# ---------------------------------------------------------------------------

_DS_BEGIN_CALLS = "\u003c\uff5ctool\u2581calls\u2581begin\uff5c\u003e"
_DS_END_CALLS = "\u003c\uff5ctool\u2581calls\u2581end\uff5c\u003e"
_DS_BEGIN_CALL = "\u003c\uff5ctool\u2581call\u2581begin\uff5c\u003e"
_DS_END_CALL = "\u003c\uff5ctool\u2581call\u2581end\uff5c\u003e"
_DS_SEP = "\u003c\uff5ctool\u2581sep\uff5c\u003e"


def _ds_tool_call(name: str, args_json: str) -> str:
  """Build a single DeepSeek-style tool-call block."""
  return (
      f"{_DS_BEGIN_CALL}function{_DS_SEP}{name}\n"
      f"```json\n{args_json}\n```"
      f"{_DS_END_CALL}"
  )


def _ds_wrapped(inner: str) -> str:
  """Wrap content in <｜tool▁calls▁begin｜>...<｜tool▁calls▁end｜>."""
  return f"{_DS_BEGIN_CALLS}{inner}{_DS_END_CALLS}"


def test_parse_deepseek_single_tool_call():
  """Single DeepSeek tool call with code-fenced JSON args."""
  text = _ds_wrapped(
      _ds_tool_call("get_weather", '{"city": "Beijing", "unit": "celsius"}')
  )
  tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
  assert len(tool_calls) == 1
  assert tool_calls[0].function.name == "get_weather"
  assert json.loads(tool_calls[0].function.arguments) == {
      "city": "Beijing",
      "unit": "celsius",
  }
  assert remainder is None


def test_parse_deepseek_multi_tool_call():
  """Multiple DeepSeek tool calls in a single wrapped block."""
  inner = _ds_tool_call("func_a", '{"x": 1}') + _ds_tool_call(
      "func_b", '{"y": 2}'
  )
  text = _ds_wrapped(inner)
  tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
  assert len(tool_calls) == 2
  assert tool_calls[0].function.name == "func_a"
  assert json.loads(tool_calls[0].function.arguments) == {"x": 1}
  assert tool_calls[1].function.name == "func_b"
  assert json.loads(tool_calls[1].function.arguments) == {"y": 2}
  assert remainder is None


def test_parse_deepseek_plain_json_args():
  """DeepSeek tool call without Markdown code fences around args."""
  inner = (
      f"{_DS_BEGIN_CALL}function{_DS_SEP}search\n"
      f'{{"query": "天气"}}'
      f"{_DS_END_CALL}"
  )
  text = _ds_wrapped(inner)
  tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
  assert len(tool_calls) == 1
  assert tool_calls[0].function.name == "search"
  assert json.loads(tool_calls[0].function.arguments) == {"query": "天气"}


def test_parse_deepseek_with_surrounding_text():
  """DeepSeek tool call embedded in surrounding non-tool text."""
  prefix = "Let me think about this.\n"
  suffix = "\nI'll proceed now."
  inner = _ds_tool_call("calculate", '{"expr": "2+2"}')
  text = prefix + _ds_wrapped(inner) + suffix
  tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
  assert len(tool_calls) == 1
  assert tool_calls[0].function.name == "calculate"
  assert remainder == "Let me think about this.\n\nI'll proceed now."


def test_parse_deepseek_no_tokens_returns_empty():
  """Text without DeepSeek tokens returns no tool calls and None remainder."""
  text = "Just a regular response, no special tokens here."
  tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
  assert tool_calls == []
  assert remainder is None


def test_parse_tool_calls_from_text_handles_deepseek_format():
  """Integration: the generic parser delegates to the DeepSeek parser."""
  text = _ds_wrapped(
      _ds_tool_call("fetch_page", '{"url": "https://example.com"}')
  )
  tool_calls, remainder = _parse_tool_calls_from_text(text)
  assert len(tool_calls) == 1
  assert tool_calls[0].function.name == "fetch_page"
  assert json.loads(tool_calls[0].function.arguments) == {
      "url": "https://example.com"
  }
  assert remainder is None


def test_parse_tool_calls_from_text_mixed_formats():
  """DeepSeek tokens + standard inline JSON in the same text."""
  ds_part = _ds_wrapped(_ds_tool_call("ds_func", '{"a": 1}'))
  standard_part = '{"name": "std_func", "arguments": {"b": 2}}'
  text = ds_part + " some text " + standard_part
  tool_calls, remainder = _parse_tool_calls_from_text(text)
  assert len(tool_calls) == 2
  assert tool_calls[0].function.name == "ds_func"
  assert tool_calls[1].function.name == "std_func"
  assert remainder == "some text"


def test_parse_deepseek_empty_text():
  """Empty or whitespace-only text returns no tool calls."""
  for text in ("", "   ", "\n\n"):
    tool_calls, remainder = _parse_deepseek_tool_calls_from_text(text)
    assert tool_calls == []
    assert remainder is None


def test_parse_deepseek_unwrapped_call_before_wrapped_block():
  """Unwrapped call preceding a wrapped block is not dropped."""
  unwrapped = _ds_tool_call("first", '{"x": 1}')
  wrapped = _ds_wrapped(_ds_tool_call("second", '{"y": 2}'))
  tool_calls, remainder = _parse_deepseek_tool_calls_from_text(
      unwrapped + wrapped
  )
  assert [tc.function.name for tc in tool_calls] == ["first", "second"]
  assert remainder is None


def test_extract_json_from_deepseek_args_invalid_fence_returns_none():
  """Invalid JSON inside a code fence is rejected rather than returned."""
  assert _extract_json_from_deepseek_args('```json\n{"a": 1,}\n```') is None


def test_split_message_content_and_tool_calls_inline_text():
  message = {
      "role": "assistant",
      "content": (
          'Intro {"name":"alpha","arguments":{"value":1}} trailing content'
      ),
  }
  content, tool_calls = _split_message_content_and_tool_calls(message)
  assert content == "Intro  trailing content"
  assert len(tool_calls) == 1
  assert tool_calls[0].function.name == "alpha"
  assert json.loads(tool_calls[0].function.arguments) == {"value": 1}


def test_split_message_content_prefers_existing_structured_calls():
  tool_call = ChatCompletionMessageToolCall(
      type="function",
      id="existing",
      function=Function(
          name="existing_call",
          arguments='{"arg": "value"}',
      ),
  )
  message = {
      "role": "assistant",
      "content": "ignored",
      "tool_calls": [tool_call],
  }
  content, tool_calls = _split_message_content_and_tool_calls(message)
  assert content == "ignored"
  assert tool_calls == [tool_call]


@pytest.mark.asyncio
async def test_get_content_does_not_filter_thought_parts():
  """Test that _get_content does not drop thought parts.

  Thought filtering is handled by the caller (e.g., _content_to_message_param)
  to avoid duplicating logic across helpers.
  """
  thought_part = types.Part(text="Internal reasoning...", thought=True)
  regular_part = types.Part.from_text(text="Visible response")

  content = await _get_content([thought_part, regular_part])

  assert content == [
      {"type": "text", "text": "Internal reasoning..."},
      {"type": "text", "text": "Visible response"},
  ]


@pytest.mark.asyncio
async def test_get_content_all_thought_parts():
  """Test that thought parts convert like regular text parts."""
  thought_part1 = types.Part(text="First reasoning...", thought=True)
  thought_part2 = types.Part(text="Second reasoning...", thought=True)

  content = await _get_content([thought_part1, thought_part2])

  assert content == [
      {"type": "text", "text": "First reasoning..."},
      {"type": "text", "text": "Second reasoning..."},
  ]


@pytest.mark.asyncio
async def test_get_content_text():
  parts = [types.Part.from_text(text="Test text")]
  content = await _get_content(parts)
  assert content == "Test text"


@pytest.mark.asyncio
async def test_get_content_text_inline_data_single_part():
  parts = [
      types.Part.from_bytes(
          data="Inline text".encode("utf-8"), mime_type="text/plain"
      )
  ]
  content = await _get_content(parts)
  assert content == "Inline text"


@pytest.mark.asyncio
async def test_get_content_text_inline_data_multiple_parts():
  parts = [
      types.Part.from_bytes(
          data="First part".encode("utf-8"), mime_type="text/plain"
      ),
      types.Part.from_text(text="Second part"),
  ]
  content = await _get_content(parts)
  assert content[0]["type"] == "text"
  assert content[0]["text"] == "First part"
  assert content[1]["type"] == "text"
  assert content[1]["text"] == "Second part"


@pytest.mark.asyncio
async def test_get_content_text_inline_data_fallback_decoding():
  parts = [
      types.Part.from_bytes(data=b"\xff", mime_type="text/plain"),
  ]
  content = await _get_content(parts)
  assert content == "ÿ"


@pytest.mark.asyncio
async def test_get_content_image():
  parts = [
      types.Part.from_bytes(data=b"test_image_data", mime_type="image/png")
  ]
  content = await _get_content(parts)
  assert content[0]["type"] == "image_url"
  assert (
      content[0]["image_url"]["url"]
      == "data:image/png;base64,dGVzdF9pbWFnZV9kYXRh"
  )
  assert "format" not in content[0]["image_url"]


@pytest.mark.asyncio
async def test_get_content_video():
  parts = [
      types.Part.from_bytes(data=b"test_video_data", mime_type="video/mp4")
  ]
  content = await _get_content(parts)
  assert content[0]["type"] == "video_url"
  assert (
      content[0]["video_url"]["url"]
      == "data:video/mp4;base64,dGVzdF92aWRlb19kYXRh"
  )
  assert "format" not in content[0]["video_url"]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "file_data,mime_type,expected_base64", FILE_BYTES_TEST_CASES
)
async def test_get_content_file_bytes(file_data, mime_type, expected_base64):
  parts = [types.Part.from_bytes(data=file_data, mime_type=mime_type)]
  content = await _get_content(parts)
  assert content[0]["type"] == "file"
  assert content[0]["file"]["file_data"] == expected_base64
  assert "format" not in content[0]["file"]


@pytest.mark.asyncio
@pytest.mark.parametrize("file_uri,mime_type", FILE_URI_TEST_CASES)
async def test_get_content_file_uri(file_uri, mime_type):
  parts = [types.Part.from_uri(file_uri=file_uri, mime_type=mime_type)]
  content = await _get_content(parts)
  assert content[0] == {
      "type": "file",
      "file": {"file_id": file_uri, "format": mime_type},
  }


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "provider,model",
    [
        ("openai", "openai/gpt-4o"),
        ("azure", "azure/gpt-4"),
    ],
)
async def test_get_content_file_uri_file_id_required_raises_error(
    provider, model
):
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/path/to/document.pdf",
              mime_type="application/pdf",
              display_name="document.pdf",
          )
      )
  ]
  with pytest.raises(
      ValueError,
      match=f"File URI `document.pdf` not supported for provider: {provider}",
  ):
    _ = await _get_content(parts, provider=provider, model=model)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "provider,model",
    [
        ("openai", "openai/gpt-4o"),
        ("azure", "azure/gpt-4"),
    ],
)
@pytest.mark.parametrize(
    "file_uri,mime_type,expected_type",
    [
        pytest.param(
            "https://example.com/image.png",
            "image/png",
            "image_url",
            id="image",
        ),
        pytest.param(
            "https://example.com/video.mp4",
            "video/mp4",
            "video_url",
            id="video",
        ),
    ],
)
async def test_get_content_file_uri_media_url_file_id_required_uses_url_type(
    provider, model, file_uri, mime_type, expected_type
):
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri=file_uri,
              mime_type=mime_type,
          )
      )
  ]
  content = await _get_content(parts, provider=provider, model=model)
  assert content == [{
      "type": expected_type,
      expected_type: {"url": file_uri},
  }]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "provider,model",
    [
        ("openai", "openai/gpt-4o"),
        ("azure", "azure/gpt-4"),
    ],
)
async def test_get_content_file_uri_file_id_required_preserves_file_id(
    provider, model
):
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="file-abc123",
              mime_type="application/pdf",
          )
      )
  ]
  content = await _get_content(parts, provider=provider, model=model)
  assert content == [{"type": "file", "file": {"file_id": "file-abc123"}}]


@pytest.mark.asyncio
async def test_get_content_file_uri_azure_preserves_assistant_file_id():
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="assistant-abc123",
              mime_type="application/pdf",
          )
      )
  ]
  content = await _get_content(parts, provider="azure", model="azure/gpt-4.1")
  assert content == [{"type": "file", "file": {"file_id": "assistant-abc123"}}]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "provider,model",
    [
        ("openai", "openai/gpt-4o"),
        ("azure", "azure/gpt-4"),
    ],
)
async def test_get_content_file_uri_http_pdf_file_id_required_raises_error(
    provider, model
):
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="https://example.com/document.pdf",
              mime_type="application/pdf",
              display_name="document.pdf",
          )
      )
  ]
  with pytest.raises(
      ValueError,
      match=f"File URI `document.pdf` not supported for provider: {provider}",
  ):
    _ = await _get_content(parts, provider=provider, model=model)


@pytest.mark.asyncio
async def test_get_content_file_uri_http_pdf_non_file_id_provider_uses_file():
  file_uri = "https://example.com/document.pdf"
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri=file_uri,
              mime_type="application/pdf",
          )
      )
  ]
  content = await _get_content(
      parts, provider="vertex_ai", model="vertex_ai/gemini-2.5-flash"
  )
  assert content == [{
      "type": "file",
      "file": {"file_id": file_uri, "format": "application/pdf"},
  }]


@pytest.mark.asyncio
async def test_get_content_file_uri_anthropic_raises_error():
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/path/to/document.pdf",
              mime_type="application/pdf",
              display_name="document.pdf",
          )
      )
  ]
  with pytest.raises(
      ValueError,
      match="File URI `document.pdf` not supported for provider: anthropic",
  ):
    _ = await _get_content(
        parts, provider="anthropic", model="anthropic/claude-3-5"
    )


@pytest.mark.asyncio
async def test_get_content_file_uri_anthropic_openai_file_id_raises_error():
  parts = [types.Part(file_data=types.FileData(file_uri="file-abc123"))]
  with pytest.raises(
      ValueError,
      match="File URI `file-<redacted>` not supported for provider: anthropic",
  ):
    _ = await _get_content(
        parts, provider="anthropic", model="anthropic/claude-3-5"
    )


@pytest.mark.asyncio
async def test_get_content_file_uri_vertex_ai_non_gemini_raises_error():
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/path/to/document.pdf",
              mime_type="application/pdf",
              display_name="document.pdf",
          )
      )
  ]
  with pytest.raises(
      ValueError,
      match="File URI `document.pdf` not supported for provider: vertex_ai",
  ):
    _ = await _get_content(
        parts, provider="vertex_ai", model="vertex_ai/claude-3-5"
    )


@pytest.mark.asyncio
async def test_get_content_file_uri_vertex_ai_gemini_keeps_file_block():
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/path/to/document.pdf",
              mime_type="application/pdf",
          )
      )
  ]
  content = await _get_content(
      parts, provider="vertex_ai", model="vertex_ai/gemini-2.5-flash"
  )
  assert content == [{
      "type": "file",
      "file": {
          "file_id": "gs://bucket/path/to/document.pdf",
          "format": "application/pdf",
      },
  }]


@pytest.mark.asyncio
async def test_get_content_file_uri_infer_mime_type():
  """Test MIME type inference from file_uri extension.

  When file_data has a file_uri with a recognizable extension but no explicit
  mime_type, the MIME type should be inferred from the extension.
  """
  # Use Part constructor directly to test MIME type inference in _get_content
  # (types.Part.from_uri does its own inference, so we bypass it)
  parts = [
      types.Part(
          file_data=types.FileData(file_uri="gs://bucket/path/to/document.pdf")
      )
  ]
  content = await _get_content(parts)
  assert content[0] == {
      "type": "file",
      "file": {
          "file_id": "gs://bucket/path/to/document.pdf",
          "format": "application/pdf",
      },
  }


@pytest.mark.asyncio
async def test_get_content_file_uri_versioned_infer_mime_type():
  """Test MIME type inference from versioned artifact URIs."""
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/path/to/document.pdf/0"
          )
      )
  ]
  content = await _get_content(parts)
  assert content[0]["file"]["format"] == "application/pdf"


@pytest.mark.asyncio
async def test_get_content_file_uri_infers_from_display_name():
  """Test MIME type inference from display_name when URI lacks extension."""
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/artifact/0",
              display_name="document.pdf",
          )
      )
  ]
  content = await _get_content(parts)
  assert content[0]["file"]["format"] == "application/pdf"


@pytest.mark.asyncio
async def test_get_content_file_uri_default_mime_type():
  """Test that file_uri without an inferable extension raises ValueError.

  When file_data has a file_uri without a recognizable extension and no explicit
  mime_type, ADK raises a clear ValueError instead of forwarding the unsupported
  'application/octet-stream' MIME type to LiteLLM.
  """
  parts = [
      types.Part(file_data=types.FileData(file_uri="gs://bucket/artifact/0"))
  ]
  with pytest.raises(ValueError, match="Cannot process file_uri"):
    await _get_content(parts)


@pytest.mark.asyncio
async def test_get_content_file_uri_explicit_octet_stream_raises():
  """Test that an explicit application/octet-stream MIME type raises ValueError.

  'application/octet-stream' is semantically equivalent to an unknown type and
  causes the same downstream ValueError from LiteLLM whether it arrives as a
  default fallback or is set explicitly by the caller. ADK raises early with
  an actionable message in both cases.
  """
  parts = [
      types.Part(
          file_data=types.FileData(
              file_uri="gs://bucket/artifact/0",
              mime_type="application/octet-stream",
          )
      )
  ]
  with pytest.raises(ValueError, match="application/octet-stream"):
    await _get_content(parts)


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "uri,expected_mime_type",
    [
        ("gs://bucket/file.pdf", "application/pdf"),
        ("gs://bucket/path/to/document.json", "application/json"),
        ("gs://bucket/image.png", "image/png"),
        ("gs://bucket/image.jpg", "image/jpeg"),
        ("gs://bucket/audio.mp3", "audio/mpeg"),
        ("gs://bucket/video.mp4", "video/mp4"),
    ],
)
async def test_get_content_file_uri_mime_type_inference(
    uri, expected_mime_type
):
  """Test MIME type inference from various file extensions."""
  # Use Part constructor directly to test MIME type inference in _get_content
  parts = [types.Part(file_data=types.FileData(file_uri=uri))]
  content = await _get_content(parts)
  assert content[0]["file"]["format"] == expected_mime_type


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "mime_type,expected_format",
    [
        ("audio/mpeg", "mp3"),
        ("audio/mp3", "mp3"),
        ("audio/wav", "wav"),
        ("audio/x-wav", "wav"),
        ("audio/wave", "wav"),
        ("audio/flac", "flac"),
        ("audio/ogg", "ogg"),
        ("audio/mp4", "mp4"),
    ],
)
async def test_get_content_audio_inline_data_emits_input_audio(
    mime_type, expected_format
):
  """Audio inline_data is serialised as `input_audio` with raw base64 + format."""
  parts = [types.Part.from_bytes(data=b"test_audio_data", mime_type=mime_type)]
  content = await _get_content(parts)
  assert content == [{
      "type": "input_audio",
      "input_audio": {
          "data": "dGVzdF9hdWRpb19kYXRh",
          "format": expected_format,
      },
  }]


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "provider,model",
    [
        ("openai", "openai/gpt-4o"),
        ("azure", "azure/gpt-4"),
    ],
)
async def test_get_content_audio_file_uri_http_raises_error(provider, model):
  """Audio HTTP file_uri raises an error for openai/azure."""
  file_uri = "https://example.com/audio.mp3"
  parts = [
      types.Part(
          file_data=types.FileData(file_uri=file_uri, mime_type="audio/mpeg")
      )
  ]
  with pytest.raises(
      ValueError,
      match=(
          "File URI `https://<redacted>/audio.mp3` not supported for provider:"
          f" {provider}"
      ),
  ):
    _ = await _get_content(parts, provider=provider, model=model)


def test_to_litellm_role():
  assert _to_litellm_role("model") == "assistant"
  assert _to_litellm_role("assistant") == "assistant"
  assert _to_litellm_role("user") == "user"
  assert _to_litellm_role(None) == "user"


@pytest.mark.parametrize(
    "response, expected_chunks, expected_usage_chunk, expected_finished",
    [
        (
            ModelResponse(
                choices=[
                    {
                        "message": {
                            "content": "this is a test",
                        }
                    }
                ],
                usage={
                    "prompt_tokens": 0,
                    "completion_tokens": 0,
                    "total_tokens": 0,
                },
            ),
            [TextChunk(text="this is a test")],
            UsageMetadataChunk(
                prompt_tokens=0, completion_tokens=0, total_tokens=0
            ),
            "stop",
        ),
        (
            ModelResponse(
                choices=[
                    {
                        "message": {
                            "content": "this is a test",
                        }
                    }
                ],
                usage={
                    "prompt_tokens": 3,
                    "completion_tokens": 5,
                    "total_tokens": 8,
                },
            ),
            [TextChunk(text="this is a test")],
            UsageMetadataChunk(
                prompt_tokens=3, completion_tokens=5, total_tokens=8
            ),
            "stop",
        ),
        (
            ModelResponseStream(
                choices=[
                    StreamingChoices(
                        finish_reason=None,
                        delta=Delta(
                            role="assistant",
                            tool_calls=[
                                ChatCompletionDeltaToolCall(
                                    type="function",
                                    id="1",
                                    function=Function(
                                        name="test_function",
                                        arguments='{"key": "va',
                                    ),
                                    index=0,
                                )
                            ],
                        ),
                    )
                ]
            ),
            [FunctionChunk(id="1", name="test_function", args='{"key": "va')],
            None,
            # LiteLLM 1.81+ defaults finish_reason to "stop" for partial chunks,
            # older versions return None. Both are valid for streaming chunks.
            (None, "stop"),
        ),
        (
            ModelResponse(choices=[{"finish_reason": "tool_calls"}]),
            [None],
            (
                None,
                UsageMetadataChunk(
                    prompt_tokens=0, completion_tokens=0, total_tokens=0
                ),
            ),
            "tool_calls",
        ),
        (
            ModelResponse(choices=[{}]),
            [None],
            (
                None,
                UsageMetadataChunk(
                    prompt_tokens=0, completion_tokens=0, total_tokens=0
                ),
            ),
            "stop",
        ),
        (
            ModelResponse(
                choices=[{
                    "finish_reason": "tool_calls",
                    "message": {
                        "role": "assistant",
                        "content": (
                            '{"id":"call_1","name":"get_current_time",'
                            '"arguments":{"timezone_str":"Asia/Taipei"}}'
                        ),
                    },
                }],
                usage={
                    "prompt_tokens": 7,
                    "completion_tokens": 9,
                    "total_tokens": 16,
                },
            ),
            [
                FunctionChunk(
                    id="call_1",
                    name="get_current_time",
                    args='{"timezone_str": "Asia/Taipei"}',
                    index=0,
                ),
            ],
            UsageMetadataChunk(
                prompt_tokens=7, completion_tokens=9, total_tokens=16
            ),
            "tool_calls",
        ),
        (
            ModelResponse(
                choices=[{
                    "finish_reason": "tool_calls",
                    "message": {
                        "role": "assistant",
                        "content": (
                            'Intro {"id":"call_2","name":"alpha",'
                            '"arguments":{"foo":"bar"}} wrap'
                        ),
                    },
                }],
                usage={
                    "prompt_tokens": 11,
                    "completion_tokens": 13,
                    "total_tokens": 24,
                },
            ),
            [
                TextChunk(text="Intro  wrap"),
                FunctionChunk(
                    id="call_2",
                    name="alpha",
                    args='{"foo": "bar"}',
                    index=0,
                ),
            ],
            UsageMetadataChunk(
                prompt_tokens=11, completion_tokens=13, total_tokens=24
            ),
            "tool_calls",
        ),
        (
            ModelResponseStream(
                choices=[
                    StreamingChoices(
                        finish_reason=None,
                        delta=Delta(role="assistant", content="Hello"),
                    )
                ],
                usage=None,
            ),
            [TextChunk(text="Hello")],
            None,
            (None, "stop"),
        ),
        (
            ModelResponseStream(
                choices=[
                    StreamingChoices(
                        finish_reason="stop",
                        delta=Delta(
                            role="assistant", reasoning_content="thinking..."
                        ),
                    )
                ],
                usage=None,
            ),
            [
                ReasoningChunk(
                    parts=[types.Part(text="thinking...", thought=True)]
                )
            ],
            None,
            "stop",
        ),
    ],
)
def test_model_response_to_chunk(
    response, expected_chunks, expected_usage_chunk, expected_finished
):
  result = list(_model_response_to_chunk(response))
  observed_chunks = []
  usage_chunk = None
  for chunk, finished in result:
    if isinstance(chunk, UsageMetadataChunk):
      usage_chunk = chunk
      continue
    observed_chunks.append((chunk, finished))

  assert len(observed_chunks) == len(expected_chunks)
  for (chunk, finished), expected_chunk in zip(
      observed_chunks, expected_chunks
  ):
    if expected_chunk is None:
      assert chunk is None
    else:
      assert isinstance(chunk, type(expected_chunk))
      assert chunk == expected_chunk
    if isinstance(expected_finished, tuple):
      assert finished in expected_finished
    else:
      assert finished == expected_finished

  if isinstance(expected_usage_chunk, tuple):
    assert usage_chunk in expected_usage_chunk
  elif expected_usage_chunk is None:
    assert usage_chunk is None
  else:
    assert usage_chunk is not None
    assert usage_chunk == expected_usage_chunk


def test_model_response_to_chunk_does_not_mutate_delta_object():
  """Verify that _model_response_to_chunk doesn't mutate the Delta object.

  In real streaming responses, LiteLLM's StreamingChoices only has 'delta'
  (message is explicitly popped in StreamingChoices constructor). The delta
  object itself carries reasoning_content when present.
  """
  delta = Delta(
      role="assistant", content="Hello", reasoning_content="thinking..."
  )
  response = ModelResponseStream(
      choices=[StreamingChoices(delta=delta, finish_reason=None)]
  )

  chunks = [chunk for chunk, _ in _model_response_to_chunk(response) if chunk]

  assert (
      ReasoningChunk(parts=[types.Part(text="thinking...", thought=True)])
      in chunks
  )
  assert TextChunk(text="Hello") in chunks

  # Verify we don't accidentally mutate the original delta object.
  assert delta.content == "Hello"
  assert delta.reasoning_content == "thinking..."


def test_model_response_to_chunk_rejects_dict_response():
  with pytest.raises(TypeError):
    list(_model_response_to_chunk({"choices": []}))


@pytest.mark.asyncio
async def test_acompletion_additional_args(mock_acompletion, mock_client):
  lite_llm_instance = LiteLlm(
      # valid args
      model="vertex_ai/test_model",
      llm_client=mock_client,
      api_key="test_key",
      api_base="some://url",
      api_version="2024-09-12",
      headers={"custom": "header"},  # Add custom header to test merge
      # invalid args (ignored)
      stream=True,
      messages=[{"role": "invalid", "content": "invalid"}],
      tools=[{
          "type": "function",
          "function": {
              "name": "invalid",
          },
      }],
  )

  async for response in lite_llm_instance.generate_content_async(
      LLM_REQUEST_WITH_FUNCTION_DECLARATION
  ):
    assert response.content.role == "model"
    assert response.content.parts[0].text == "Test response"
    assert response.content.parts[1].function_call.name == "test_function"
    assert response.content.parts[1].function_call.args == {
        "test_arg": "test_value"
    }
    assert response.content.parts[1].function_call.id == "test_tool_call_id"

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args

  assert kwargs["model"] == "vertex_ai/test_model"
  assert kwargs["messages"][0]["role"] == "user"
  assert kwargs["messages"][0]["content"] == "Test prompt"
  assert kwargs["tools"][0]["function"]["name"] == "test_function"
  assert "stream" not in kwargs
  assert "llm_client" not in kwargs
  assert kwargs["api_base"] == "some://url"
  assert "headers" in kwargs
  assert kwargs["headers"]["custom"] == "header"
  assert "x-goog-api-client" in kwargs["headers"]
  assert "user-agent" in kwargs["headers"]


@pytest.mark.asyncio
async def test_acompletion_additional_args_non_vertex(
    mock_acompletion, mock_client
):
  """Test that tracking headers are not added for non-Vertex AI models."""
  lite_llm_instance = LiteLlm(
      model="openai/gpt-4o",
      llm_client=mock_client,
      api_key="test_key",
      headers={"custom": "header"},
  )

  async for _ in lite_llm_instance.generate_content_async(
      LLM_REQUEST_WITH_FUNCTION_DECLARATION
  ):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert kwargs["model"] == "openai/gpt-4o"
  assert "headers" in kwargs
  assert kwargs["headers"]["custom"] == "header"
  assert "x-goog-api-client" not in kwargs["headers"]
  assert "user-agent" not in kwargs["headers"]


@pytest.mark.asyncio
async def test_acompletion_with_drop_params(mock_acompletion, mock_client):
  lite_llm_instance = LiteLlm(
      model="test_model", llm_client=mock_client, drop_params=True
  )

  async for _ in lite_llm_instance.generate_content_async(
      LLM_REQUEST_WITH_FUNCTION_DECLARATION
  ):
    pass

  mock_acompletion.assert_called_once()

  _, kwargs = mock_acompletion.call_args
  assert kwargs["drop_params"] is True


@pytest.mark.asyncio
async def test_completion_additional_args(mock_completion, mock_client):
  lite_llm_instance = LiteLlm(
      # valid args
      model="test_model",
      llm_client=mock_client,
      api_key="test_key",
      api_base="some://url",
      api_version="2024-09-12",
      # invalid args (ignored)
      stream=False,
      messages=[{"role": "invalid", "content": "invalid"}],
      tools=[{
          "type": "function",
          "function": {
              "name": "invalid",
          },
      }],
  )

  mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE)

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]
  assert len(responses) == 4
  mock_completion.assert_called_once()

  _, kwargs = mock_completion.call_args

  assert kwargs["model"] == "test_model"
  assert kwargs["messages"][0]["role"] == "user"
  assert kwargs["messages"][0]["content"] == "Test prompt"
  assert kwargs["tools"][0]["function"]["name"] == "test_function"
  assert kwargs["stream"]
  assert "llm_client" not in kwargs
  assert kwargs["api_base"] == "some://url"


@pytest.mark.asyncio
async def test_completion_with_drop_params(mock_completion, mock_client):
  lite_llm_instance = LiteLlm(
      model="test_model", llm_client=mock_client, drop_params=True
  )

  mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE)

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]
  assert len(responses) == 4

  mock_completion.assert_called_once()

  _, kwargs = mock_completion.call_args
  assert kwargs["drop_params"] is True


@pytest.mark.asyncio
async def test_generate_content_async_stream_grounding_metadata(
    mock_completion, lite_llm_instance
):
  final_chunk = ModelResponseStream(
      model="test_model",
      choices=[StreamingChoices(finish_reason="stop", delta=Delta())],
  )
  final_chunk.vertex_ai_grounding_metadata = {
      "grounding_chunks": [
          {"web": {"uri": "https://example.com", "title": "Example"}}
      ],
  }
  mock_completion.return_value = iter([
      ModelResponseStream(
          model="test_model",
          choices=[
              StreamingChoices(
                  finish_reason=None,
                  delta=Delta(role="assistant", content="Grounded answer"),
              )
          ],
      ),
      final_chunk,
  ])

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
  )

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          llm_request, stream=True
      )
  ]

  assert responses[-1].partial is False
  assert responses[-1].grounding_metadata is not None
  assert (
      responses[-1].grounding_metadata.grounding_chunks[0].web.uri
      == "https://example.com"
  )


@pytest.mark.asyncio
async def test_generate_content_async_stream_with_usage_metadata(
    mock_completion, lite_llm_instance
):

  mock_completion.return_value = iter(STREAMING_MODEL_RESPONSE)

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]
  assert len(responses) == 4
  assert responses[0].content.role == "model"
  assert responses[0].content.parts[0].text == "zero, "
  assert responses[0].model_version == "test_model"
  assert responses[1].content.role == "model"
  assert responses[1].content.parts[0].text == "one, "
  assert responses[1].model_version == "test_model"
  assert responses[2].content.role == "model"
  assert responses[2].content.parts[0].text == "two:"
  assert responses[2].model_version == "test_model"
  assert responses[3].content.role == "model"
  assert responses[3].content.parts[-1].function_call.name == "test_function"
  assert responses[3].content.parts[-1].function_call.args == {
      "test_arg": "test_value"
  }
  assert responses[3].content.parts[-1].function_call.id == "test_tool_call_id"
  assert responses[3].finish_reason == types.FinishReason.STOP
  assert responses[3].model_version == "test_model"
  mock_completion.assert_called_once()

  _, kwargs = mock_completion.call_args
  assert kwargs["model"] == "test_model"
  assert kwargs["messages"][0]["role"] == "user"
  assert kwargs["messages"][0]["content"] == "Test prompt"
  assert kwargs["tools"][0]["function"]["name"] == "test_function"
  assert (
      kwargs["tools"][0]["function"]["description"]
      == "Test function description"
  )
  assert (
      kwargs["tools"][0]["function"]["parameters"]["properties"]["test_arg"][
          "type"
      ]
      == "string"
  )


@pytest.mark.asyncio
async def test_generate_content_async_stream_sets_finish_reason(
    mock_completion, lite_llm_instance
):
  mock_completion.return_value = iter([
      ModelResponseStream(
          model="test_model",
          choices=[
              StreamingChoices(
                  finish_reason=None,
                  delta=Delta(role="assistant", content="Hello "),
              )
          ],
      ),
      ModelResponseStream(
          model="test_model",
          choices=[
              StreamingChoices(
                  finish_reason=None,
                  delta=Delta(role="assistant", content="world"),
              )
          ],
      ),
      ModelResponseStream(
          model="test_model",
          choices=[StreamingChoices(finish_reason="stop", delta=Delta())],
      ),
  ])

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
  )

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          llm_request, stream=True
      )
  ]

  assert responses[-1].partial is False
  assert responses[-1].finish_reason == types.FinishReason.STOP
  assert responses[-1].content.parts[0].text == "Hello world"


@pytest.mark.asyncio
async def test_generate_content_async_stream_with_usage_metadata(
    mock_completion, lite_llm_instance
):

  streaming_model_response_with_usage_metadata = [
      *STREAMING_MODEL_RESPONSE,
      ModelResponseStream(
          usage={
              "prompt_tokens": 10,
              "completion_tokens": 5,
              "total_tokens": 15,
              "completion_tokens_details": {"reasoning_tokens": 5},
          },
          choices=[
              StreamingChoices(
                  finish_reason=None,
              )
          ],
      ),
  ]

  mock_completion.return_value = iter(
      streaming_model_response_with_usage_metadata
  )

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]
  assert len(responses) == 4
  assert responses[0].content.role == "model"
  assert responses[0].content.parts[0].text == "zero, "
  assert responses[1].content.role == "model"
  assert responses[1].content.parts[0].text == "one, "
  assert responses[2].content.role == "model"
  assert responses[2].content.parts[0].text == "two:"
  assert responses[3].content.role == "model"
  assert responses[3].content.parts[-1].function_call.name == "test_function"
  assert responses[3].content.parts[-1].function_call.args == {
      "test_arg": "test_value"
  }
  assert responses[3].content.parts[-1].function_call.id == "test_tool_call_id"
  assert responses[3].finish_reason == types.FinishReason.STOP

  assert responses[3].usage_metadata.prompt_token_count == 10
  assert responses[3].usage_metadata.candidates_token_count == 5
  assert responses[3].usage_metadata.total_token_count == 15
  assert responses[3].usage_metadata.thoughts_token_count == 5

  mock_completion.assert_called_once()

  _, kwargs = mock_completion.call_args
  assert kwargs["model"] == "test_model"
  assert kwargs["messages"][0]["role"] == "user"
  assert kwargs["messages"][0]["content"] == "Test prompt"
  assert kwargs["tools"][0]["function"]["name"] == "test_function"
  assert (
      kwargs["tools"][0]["function"]["description"]
      == "Test function description"
  )
  assert (
      kwargs["tools"][0]["function"]["parameters"]["properties"]["test_arg"][
          "type"
      ]
      == "string"
  )


@pytest.mark.asyncio
async def test_generate_content_async_stream_with_usage_metadata(
    mock_completion, lite_llm_instance
):
  """Tests that cached prompt tokens are propagated in streaming mode."""
  streaming_model_response_with_usage_metadata = [
      *STREAMING_MODEL_RESPONSE,
      ModelResponseStream(
          usage={
              "prompt_tokens": 10,
              "completion_tokens": 5,
              "total_tokens": 15,
              "cached_tokens": 8,
              "completion_tokens_details": {"reasoning_tokens": 5},
          },
          choices=[
              StreamingChoices(
                  finish_reason=None,
              )
          ],
      ),
  ]

  mock_completion.return_value = iter(
      streaming_model_response_with_usage_metadata
  )

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]
  assert len(responses) == 4
  assert responses[3].usage_metadata.prompt_token_count == 10
  assert responses[3].usage_metadata.candidates_token_count == 5
  assert responses[3].usage_metadata.total_token_count == 15
  assert responses[3].usage_metadata.cached_content_token_count == 8
  assert responses[3].usage_metadata.thoughts_token_count == 5


@pytest.mark.asyncio
async def test_generate_content_async_multiple_function_calls(
    mock_completion, lite_llm_instance
):
  """Test handling of multiple function calls with different indices in streaming mode.

  This test verifies that:
  1. Multiple function calls with different indices are handled correctly
  2. Arguments and names are properly accumulated for each function call
  3. The final response contains all function calls with correct indices
  """
  mock_completion.return_value = MULTIPLE_FUNCTION_CALLS_STREAM

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user",
              parts=[types.Part.from_text(text="Test multiple function calls")],
          )
      ],
      config=types.GenerateContentConfig(
          tools=[
              types.Tool(
                  function_declarations=[
                      types.FunctionDeclaration(
                          name="function_1",
                          description="First test function",
                          parameters=types.Schema(
                              type=types.Type.OBJECT,
                              properties={
                                  "arg": types.Schema(type=types.Type.STRING),
                              },
                          ),
                      ),
                      types.FunctionDeclaration(
                          name="function_2",
                          description="Second test function",
                          parameters=types.Schema(
                              type=types.Type.OBJECT,
                              properties={
                                  "arg": types.Schema(type=types.Type.STRING),
                              },
                          ),
                      ),
                  ]
              )
          ],
      ),
  )

  responses = []
  async for response in lite_llm_instance.generate_content_async(
      llm_request, stream=True
  ):
    responses.append(response)

  # Verify we got the final response with both function calls
  assert len(responses) > 0
  final_response = responses[-1]
  assert final_response.content.role == "model"
  assert len(final_response.content.parts) == 2

  # Verify first function call
  assert final_response.content.parts[0].function_call.name == "function_1"
  assert final_response.content.parts[0].function_call.id == "call_1"
  assert final_response.content.parts[0].function_call.args == {"arg": "value1"}

  # Verify second function call
  assert final_response.content.parts[1].function_call.name == "function_2"
  assert final_response.content.parts[1].function_call.id == "call_2"
  assert final_response.content.parts[1].function_call.args == {"arg": "value2"}


@pytest.mark.asyncio
async def test_generate_content_async_non_compliant_multiple_function_calls(
    mock_completion, lite_llm_instance
):
  """Test handling of multiple function calls with same 0 indices in streaming mode.

  This test verifies that:
  1. Multiple function calls with same indices (0) are handled correctly
  2. Arguments and names are properly accumulated for each function call
  3. The final response contains all function calls with correct incremented
  indices
  """
  mock_completion.return_value = NON_COMPLIANT_MULTIPLE_FUNCTION_CALLS_STREAM

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user",
              parts=[types.Part.from_text(text="Test multiple function calls")],
          )
      ],
      config=types.GenerateContentConfig(
          tools=[
              types.Tool(
                  function_declarations=[
                      types.FunctionDeclaration(
                          name="function_1",
                          description="First test function",
                          parameters=types.Schema(
                              type=types.Type.OBJECT,
                              properties={
                                  "arg": types.Schema(type=types.Type.STRING),
                              },
                          ),
                      ),
                      types.FunctionDeclaration(
                          name="function_2",
                          description="Second test function",
                          parameters=types.Schema(
                              type=types.Type.OBJECT,
                              properties={
                                  "arg": types.Schema(type=types.Type.STRING),
                              },
                          ),
                      ),
                  ]
              )
          ],
      ),
  )

  responses = []
  async for response in lite_llm_instance.generate_content_async(
      llm_request, stream=True
  ):
    responses.append(response)

  # Verify we got the final response with both function calls
  assert len(responses) > 0
  final_response = responses[-1]
  assert final_response.content.role == "model"
  assert len(final_response.content.parts) == 2

  # Verify first function call
  assert final_response.content.parts[0].function_call.name == "function_1"
  assert final_response.content.parts[0].function_call.id == "0"
  assert final_response.content.parts[0].function_call.args == {"arg": "value1"}

  # Verify second function call
  assert final_response.content.parts[1].function_call.name == "function_2"
  assert final_response.content.parts[1].function_call.id == "1"
  assert final_response.content.parts[1].function_call.args == {"arg": "value2"}


@pytest.mark.asyncio
async def test_generate_content_async_stream_with_empty_chunk(
    mock_completion, lite_llm_instance
):
  """Tests that empty tool call chunks in a stream are ignored."""
  mock_completion.return_value = iter(STREAM_WITH_EMPTY_CHUNK)

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]

  assert len(responses) == 1
  final_response = responses[0]
  assert final_response.content.role == "model"

  # Crucially, assert that only ONE tool call was generated,
  # proving the empty chunk was ignored.
  assert len(final_response.content.parts) == 1

  function_call = final_response.content.parts[0].function_call
  assert function_call.name == "test_function"
  assert function_call.id == "call_abc"
  assert function_call.args == {"test_arg": "value"}


@pytest.mark.asyncio
async def test_streaming_tool_call_truncated_by_max_tokens(
    mock_completion, lite_llm_instance
):
  """Tests that truncated tool calls with finish_reason='length' yield an error LlmResponse."""
  stream_chunks = [
      ModelResponseStream(
          choices=[
              StreamingChoices(
                  finish_reason=None,
                  delta=Delta(
                      role="assistant",
                      tool_calls=[
                          ChatCompletionDeltaToolCall(
                              type="function",
                              id="call_123",
                              function=Function(
                                  name="test_function",
                                  arguments='{"test_arg":',
                              ),
                              index=0,
                          )
                      ],
                  ),
              )
          ]
      ),
      ModelResponseStream(
          choices=[StreamingChoices(finish_reason="length", delta=Delta())]
      ),
  ]
  mock_completion.return_value = iter(stream_chunks)

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]

  assert len(responses) == 1
  error_response = responses[0]
  assert error_response.error_code == types.FinishReason.MAX_TOKENS
  assert error_response.finish_reason == types.FinishReason.MAX_TOKENS
  assert "truncated" in error_response.error_message
  assert "max_output_tokens" in error_response.error_message


@pytest.mark.asyncio
async def test_streaming_tool_call_complete_with_length_finish_reason(
    mock_completion, lite_llm_instance
):
  """Tests that complete tool calls with finish_reason='length' are yielded normally."""
  stream_chunks = [
      ModelResponseStream(
          choices=[
              StreamingChoices(
                  finish_reason=None,
                  delta=Delta(
                      role="assistant",
                      tool_calls=[
                          ChatCompletionDeltaToolCall(
                              type="function",
                              id="call_456",
                              function=Function(
                                  name="test_function",
                                  arguments='{"test_arg": "value"}',
                              ),
                              index=0,
                          )
                      ],
                  ),
              )
          ]
      ),
      ModelResponseStream(
          choices=[StreamingChoices(finish_reason="length", delta=Delta())]
      ),
  ]
  mock_completion.return_value = iter(stream_chunks)

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]

  assert len(responses) == 1
  final_response = responses[0]
  assert final_response.content.role == "model"
  assert len(final_response.content.parts) == 1

  function_call = final_response.content.parts[0].function_call
  assert function_call.name == "test_function"
  assert function_call.id == "call_456"
  assert function_call.args == {"test_arg": "value"}
  assert final_response.finish_reason == types.FinishReason.MAX_TOKENS
  assert final_response.error_code == types.FinishReason.MAX_TOKENS


@pytest.mark.asyncio
async def test_streaming_text_truncated_by_max_tokens(
    mock_completion, lite_llm_instance
):
  """Tests that text responses with finish_reason='length' set MAX_TOKENS error."""
  stream_chunks = [
      ModelResponseStream(
          choices=[
              StreamingChoices(
                  finish_reason=None,
                  delta=Delta(
                      role="assistant",
                      content="Hello, I am",
                  ),
              )
          ]
      ),
      ModelResponseStream(
          choices=[StreamingChoices(finish_reason="length", delta=Delta())]
      ),
  ]
  mock_completion.return_value = iter(stream_chunks)

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Say hello")]
          )
      ],
  )

  responses = [
      response
      async for response in lite_llm_instance.generate_content_async(
          llm_request, stream=True
      )
  ]

  partial_responses = [r for r in responses if r.partial]
  aggregated_responses = [r for r in responses if not r.partial]

  assert len(partial_responses) == 1
  assert len(aggregated_responses) == 1
  aggregated = aggregated_responses[0]
  assert aggregated.finish_reason == types.FinishReason.MAX_TOKENS
  assert aggregated.error_code == types.FinishReason.MAX_TOKENS
  assert "Maximum tokens reached" in aggregated.error_message


@pytest.mark.asyncio
async def test_get_completion_inputs_generation_params():
  # Test that generation_params are extracted and mapped correctly
  req = LlmRequest(
      contents=[
          types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
      ],
      config=types.GenerateContentConfig(
          temperature=0.33,
          max_output_tokens=123,
          top_p=0.88,
          top_k=7,
          stop_sequences=["foo", "bar"],
          presence_penalty=0.1,
          frequency_penalty=0.2,
      ),
  )

  _, _, _, generation_params = await _get_completion_inputs(
      req, model="gpt-4o-mini"
  )
  assert generation_params["temperature"] == 0.33
  assert generation_params["max_completion_tokens"] == 123
  assert generation_params["top_p"] == 0.88
  assert generation_params["top_k"] == 7
  assert generation_params["stop"] == ["foo", "bar"]
  assert generation_params["presence_penalty"] == 0.1
  assert generation_params["frequency_penalty"] == 0.2
  # Should not include max_output_tokens
  assert "max_output_tokens" not in generation_params
  assert "stop_sequences" not in generation_params


@pytest.mark.asyncio
async def test_get_completion_inputs_empty_generation_params():
  # Test that generation_params is None when no generation parameters are set
  req = LlmRequest(
      contents=[
          types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
      ],
      config=types.GenerateContentConfig(),
  )

  _, _, _, generation_params = await _get_completion_inputs(
      req, model="gpt-4o-mini"
  )
  assert generation_params is None


@pytest.mark.asyncio
async def test_get_completion_inputs_minimal_config():
  # Test that generation_params is None when config has no generation parameters
  req = LlmRequest(
      contents=[
          types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
      ],
      config=types.GenerateContentConfig(
          system_instruction="test instruction"  # Non-generation parameter
      ),
  )

  _, _, _, generation_params = await _get_completion_inputs(
      req, model="gpt-4o-mini"
  )
  assert generation_params is None


@pytest.mark.asyncio
async def test_get_completion_inputs_partial_generation_params():
  # Test that generation_params is correctly built even with only some parameters
  req = LlmRequest(
      contents=[
          types.Content(role="user", parts=[types.Part.from_text(text="hi")]),
      ],
      config=types.GenerateContentConfig(
          temperature=0.7,
          # Only temperature is set, others are None/default
      ),
  )

  _, _, _, generation_params = await _get_completion_inputs(
      req, model="gpt-4o-mini"
  )
  assert generation_params is not None
  assert generation_params["temperature"] == 0.7
  # Should only contain the temperature parameter
  assert len(generation_params) == 1


def test_function_declaration_to_tool_param_edge_cases():
  """Test edge cases for function declaration conversion that caused the original bug."""
  from google.adk.models.lite_llm import _function_declaration_to_tool_param

  # Test function with None parameters (the original bug scenario)
  func_decl = types.FunctionDeclaration(
      name="test_function_none_params",
      description="Function with None parameters",
      parameters=None,
  )
  result = _function_declaration_to_tool_param(func_decl)
  expected = {
      "type": "function",
      "function": {
          "name": "test_function_none_params",
          "description": "Function with None parameters",
          "parameters": {
              "type": "object",
              "properties": {},
          },
      },
  }
  assert result == expected

  # Verify no 'required' field is added when parameters is None
  assert "required" not in result["function"]["parameters"]


@pytest.mark.parametrize(
    "usage, expected_tokens",
    [
        ({"prompt_tokens_details": {"cached_tokens": 123}}, 123),
        (
            {
                "prompt_tokens_details": [
                    {"cached_tokens": 50},
                    {"cached_tokens": 25},
                ]
            },
            75,
        ),
        ({"cached_prompt_tokens": 45}, 45),
        ({"cached_tokens": 67}, 67),
        ({"prompt_tokens": 100}, 0),
        ({}, 0),
        ("not a dict", 0),
        (None, 0),
        ({"prompt_tokens_details": {"cached_tokens": "not a number"}}, 0),
        (json.dumps({"cached_tokens": 89}), 89),
        (json.dumps({"some_key": "some_value"}), 0),
    ],
)
def test_extract_cached_prompt_tokens(usage, expected_tokens):
  from google.adk.models.lite_llm import _extract_cached_prompt_tokens

  assert _extract_cached_prompt_tokens(usage) == expected_tokens


def test_gemini_via_litellm_warning(monkeypatch):
  """Test that Gemini via LiteLLM shows warning."""
  # Ensure environment variable is not set
  monkeypatch.delenv("ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", raising=False)
  with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    # Test with Google AI Studio Gemini via LiteLLM
    LiteLlm(model="gemini/gemini-2.5-pro-exp-03-25")
    assert len(w) == 1
    assert issubclass(w[0].category, UserWarning)
    assert "[GEMINI_VIA_LITELLM]" in str(w[0].message)
    assert "better performance" in str(w[0].message)
    assert "gemini-2.5-pro-exp-03-25" in str(w[0].message)
    assert "ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS" in str(w[0].message)


def test_gemini_via_litellm_warning_vertex_ai(monkeypatch):
  """Test that Vertex AI Gemini via LiteLLM shows warning."""
  # Ensure environment variable is not set
  monkeypatch.delenv("ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", raising=False)
  with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    # Test with Vertex AI Gemini via LiteLLM
    LiteLlm(model="vertex_ai/gemini-2.5-flash")
    assert len(w) == 1
    assert issubclass(w[0].category, UserWarning)
    assert "[GEMINI_VIA_LITELLM]" in str(w[0].message)
    assert "vertex_ai/gemini-2.5-flash" in str(w[0].message)


def test_gemini_via_litellm_warning_suppressed(monkeypatch):
  """Test that Gemini via LiteLLM warning can be suppressed."""
  monkeypatch.setenv("ADK_SUPPRESS_GEMINI_LITELLM_WARNINGS", "true")
  with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    LiteLlm(model="gemini/gemini-2.5-pro-exp-03-25")
    assert len(w) == 0


def test_non_gemini_litellm_no_warning():
  """Test that non-Gemini models via LiteLLM don't show warning."""
  with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    # Test with non-Gemini model
    LiteLlm(model="openai/gpt-4o")
    assert len(w) == 0


@pytest.mark.parametrize(
    "finish_reason,response_content,expected_content,has_tool_calls",
    [
        ("length", "Test response", "Test response", False),
        ("stop", "Complete response", "Complete response", False),
        (
            "tool_calls",
            "",
            "",
            True,
        ),
        ("content_filter", "", "", False),
    ],
    ids=["length", "stop", "tool_calls", "content_filter"],
)
@pytest.mark.asyncio
async def test_finish_reason_propagation(
    mock_acompletion,
    lite_llm_instance,
    finish_reason,
    response_content,
    expected_content,
    has_tool_calls,
):
  """Test that finish_reason is properly propagated from LiteLLM response."""
  tool_calls = None
  if has_tool_calls:
    tool_calls = [
        ChatCompletionMessageToolCall(
            type="function",
            id="test_id",
            function=Function(
                name="test_function",
                arguments='{"arg": "value"}',
            ),
        )
    ]

  mock_response = ModelResponse(
      choices=[
          Choices(
              message=ChatCompletionAssistantMessage(
                  role="assistant",
                  content=response_content,
                  tool_calls=tool_calls,
              ),
              finish_reason=finish_reason,
          )
      ]
  )
  mock_acompletion.return_value = mock_response

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
  )

  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"
    # Verify finish_reason is mapped to FinishReason enum
    assert isinstance(response.finish_reason, types.FinishReason)
    # Verify correct enum mapping using the actual mapping from lite_llm
    assert response.finish_reason == _FINISH_REASON_MAPPING[finish_reason]
    if expected_content:
      assert response.content.parts[0].text == expected_content
    if has_tool_calls:
      assert len(response.content.parts) > 0
      assert response.content.parts[-1].function_call.name == "test_function"

  mock_acompletion.assert_called_once()


def test_model_response_to_generate_content_response_no_message_with_finish_reason():
  """Test response with no message but finish_reason returns empty LlmResponse.

  This test covers issue #3618: when a turn ends with tool calls and no final
  message, we should return an empty LlmResponse instead of raising ValueError.
  """
  response = ModelResponse(
      model="test_model",
      choices=[{
          "finish_reason": "tool_calls",
          # message is missing/None
      }],
      usage={
          "prompt_tokens": 10,
          "completion_tokens": 5,
          "total_tokens": 15,
      },
  )
  # Force message to be None to guarantee hitting the else branch
  response.choices[0].message = None

  llm_response = _model_response_to_generate_content_response(response)

  # Should return empty LlmResponse, not raise ValueError
  assert llm_response.content is not None
  assert llm_response.content.role == "model"
  assert len(llm_response.content.parts) == 0
  # tool_calls maps to STOP
  assert llm_response.finish_reason == types.FinishReason.STOP
  assert llm_response.usage_metadata is not None
  assert llm_response.usage_metadata.prompt_token_count == 10
  assert llm_response.usage_metadata.candidates_token_count == 5
  assert llm_response.model_version == "test_model"


def test_model_response_to_generate_content_response_no_message_no_finish_reason():
  """Test response with no message and no finish_reason returns empty LlmResponse."""
  response = ModelResponse(
      model="test_model",
      choices=[{
          # Both message and finish_reason are missing
      }],
  )
  # Force message to be None to guarantee hitting the else branch
  response.choices[0].message = None

  llm_response = _model_response_to_generate_content_response(response)

  # Should return empty LlmResponse, not raise ValueError
  assert llm_response.content is not None
  assert llm_response.content.role == "model"
  assert len(llm_response.content.parts) == 0
  # finish_reason may be None or have a default value - the important thing
  # is that we don't raise ValueError
  assert llm_response.model_version == "test_model"


def test_model_response_to_generate_content_response_empty_message_dict():
  """Test response with empty message dict returns empty LlmResponse."""
  response = ModelResponse(
      model="test_model",
      choices=[{
          "message": {},  # Empty dict is falsy
          "finish_reason": "stop",
      }],
      usage={
          "prompt_tokens": 5,
          "completion_tokens": 3,
          "total_tokens": 8,
      },
  )
  # Ensure we test the parsing of an empty message dictionary rather than None.

  llm_response = _model_response_to_generate_content_response(response)

  # Should return empty LlmResponse, not raise ValueError
  assert llm_response.content is not None
  assert llm_response.content.role == "model"
  assert len(llm_response.content.parts) == 0
  assert llm_response.finish_reason == types.FinishReason.STOP
  assert llm_response.usage_metadata is not None


def test_model_response_to_generate_content_response_safety_finish_reason():
  """Test that SAFETY finish reason sets error_code and error_message."""
  response = ModelResponse(
      model="test_model",
      choices=[{
          "finish_reason": "content_filter",
      }],
  )
  # Force message to be None to guarantee hitting the else branch
  response.choices[0].message = None

  llm_response = _model_response_to_generate_content_response(response)

  assert llm_response.finish_reason == types.FinishReason.SAFETY
  assert llm_response.error_code == types.FinishReason.SAFETY
  assert llm_response.error_message == "Finished with SAFETY"


@pytest.mark.skip(reason="LiteLLM finish_reason mapping behaviour changed")
@pytest.mark.asyncio
async def test_finish_reason_unknown_maps_to_other(
    mock_acompletion, lite_llm_instance
):
  """Test that unmapped finish_reason values map to FinishReason.OTHER."""
  # LiteLLM's Choices model normalizes finish_reason values (e.g., "eos" ->
  # "stop") before ADK processes them. To test ADK's own fallback mapping,
  # construct a mock response that bypasses LiteLLM's normalization and
  # returns a raw unmapped finish_reason string.
  mock_choice = MagicMock()
  mock_choice.get = lambda key, default=None: {
      "message": ChatCompletionAssistantMessage(
          role="assistant",
          content="Test response",
      ),
      "finish_reason": "totally_unknown_reason",
  }.get(key, default)

  mock_response = MagicMock()
  mock_response.get = lambda key, default=None: {
      "choices": [mock_choice],
  }.get(key, default)
  mock_response.model = "test_model"

  mock_acompletion.return_value = mock_response

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
  )

  async for response in lite_llm_instance.generate_content_async(llm_request):
    assert response.content.role == "model"
    # Unknown finish_reason should map to OTHER
    assert isinstance(response.finish_reason, types.FinishReason)
    assert response.finish_reason == types.FinishReason.OTHER

  mock_acompletion.assert_called_once()


# Tests for provider detection and file_id support


@pytest.mark.parametrize(
    "model_string, expected_provider",
    [
        # Standard provider/model format
        ("openai/gpt-4o", "openai"),
        ("azure/gpt-4", "azure"),
        ("groq/llama3-70b", "groq"),
        ("anthropic/claude-3", "anthropic"),
        ("vertex_ai/gemini-pro", "vertex_ai"),
        # Fallback heuristics
        ("gpt-4o", "openai"),
        ("o1-preview", "openai"),
        ("azure-gpt-4", "azure"),
        # Unknown models
        ("custom-model", ""),
        ("", ""),
        (None, ""),
    ],
)
def test_get_provider_from_model(model_string, expected_provider):
  """Test provider extraction from model strings."""
  assert _get_provider_from_model(model_string) == expected_provider


@pytest.mark.parametrize(
    "provider, expected_in_list",
    [
        ("openai", True),
        ("azure", True),
        ("anthropic", False),
        ("vertex_ai", False),
    ],
)
def test_file_id_required_providers(provider, expected_in_list):
  """Test that the correct providers require file_id."""
  assert (provider in _FILE_ID_REQUIRED_PROVIDERS) == expected_in_list


@pytest.mark.asyncio
async def test_get_content_pdf_openai_uses_file_id(mocker):
  """Test that PDF files use file_id for OpenAI provider."""
  mock_file_response = mocker.create_autospec(litellm.FileObject)
  mock_file_response.id = "file-abc123"
  mock_acreate_file = AsyncMock(return_value=mock_file_response)
  mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)

  parts = [
      types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf")
  ]
  content = await _get_content(parts, provider="openai")

  assert content[0]["type"] == "file"
  assert content[0]["file"]["file_id"] == "file-abc123"
  assert "file_data" not in content[0]["file"]

  mock_acreate_file.assert_called_once_with(
      file=b"test_pdf_data",
      purpose="assistants",
      custom_llm_provider="openai",
  )


@pytest.mark.asyncio
async def test_get_content_pdf_non_openai_uses_file_data():
  """Test that PDF files use file_data for non-OpenAI providers."""
  parts = [
      types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf")
  ]
  content = await _get_content(parts, provider="anthropic")

  assert content[0]["type"] == "file"
  assert "file_data" in content[0]["file"]
  assert content[0]["file"]["file_data"].startswith(
      "data:application/pdf;base64,"
  )
  assert "file_id" not in content[0]["file"]


@pytest.mark.asyncio
async def test_get_content_pdf_azure_uses_file_id(mocker):
  """Test that PDF files use file_id for Azure provider."""
  mock_file_response = mocker.create_autospec(litellm.FileObject)
  mock_file_response.id = "file-xyz789"
  mock_acreate_file = AsyncMock(return_value=mock_file_response)
  mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)

  parts = [
      types.Part.from_bytes(data=b"test_pdf_data", mime_type="application/pdf")
  ]
  content = await _get_content(parts, provider="azure")

  assert content[0]["type"] == "file"
  assert content[0]["file"]["file_id"] == "file-xyz789"

  mock_acreate_file.assert_called_once_with(
      file=b"test_pdf_data",
      purpose="assistants",
      custom_llm_provider="azure",
  )


@pytest.mark.asyncio
async def test_get_completion_inputs_openai_file_upload(mocker):
  """Test that _get_completion_inputs uploads files for OpenAI models."""
  mock_file_response = mocker.create_autospec(litellm.FileObject)
  mock_file_response.id = "file-uploaded123"
  mock_acreate_file = AsyncMock(return_value=mock_file_response)
  mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)

  pdf_part = types.Part.from_bytes(
      data=b"test_pdf_content", mime_type="application/pdf"
  )
  llm_request = LlmRequest(
      model="openai/gpt-4o",
      contents=[
          types.Content(
              role="user",
              parts=[
                  types.Part.from_text(text="Analyze this PDF"),
                  pdf_part,
              ],
          )
      ],
      config=types.GenerateContentConfig(tools=[]),
  )

  messages, tools, response_format, generation_params = (
      await _get_completion_inputs(llm_request, model="openai/gpt-4o")
  )

  assert len(messages) == 1
  assert messages[0]["role"] == "user"
  content = messages[0]["content"]
  assert len(content) == 2
  assert content[0]["type"] == "text"
  assert content[0]["text"] == "Analyze this PDF"
  assert content[1]["type"] == "file"
  assert content[1]["file"]["file_id"] == "file-uploaded123"

  mock_acreate_file.assert_called_once()


@pytest.mark.asyncio
async def test_get_completion_inputs_non_openai_no_file_upload(mocker):
  """Test that _get_completion_inputs does not upload files for non-OpenAI models."""
  mock_acreate_file = AsyncMock()
  mocker.patch.object(litellm, "acreate_file", new=mock_acreate_file)

  pdf_part = types.Part.from_bytes(
      data=b"test_pdf_content", mime_type="application/pdf"
  )
  llm_request = LlmRequest(
      model="anthropic/claude-3-opus",
      contents=[
          types.Content(
              role="user",
              parts=[
                  types.Part.from_text(text="Analyze this PDF"),
                  pdf_part,
              ],
          )
      ],
      config=types.GenerateContentConfig(tools=[]),
  )

  messages, tools, response_format, generation_params = (
      await _get_completion_inputs(llm_request, model="anthropic/claude-3-opus")
  )

  assert len(messages) == 1
  content = messages[0]["content"]
  assert content[1]["type"] == "file"
  assert "file_data" in content[1]["file"]
  assert "file_id" not in content[1]["file"]

  mock_acreate_file.assert_not_called()


class TestRedirectLitellmLoggersToStdout(unittest.TestCase):
  """Tests for _redirect_litellm_loggers_to_stdout function."""

  def test_redirects_stderr_handler_to_stdout(self):
    """Test that handlers pointing to stderr are redirected to stdout."""
    test_logger = logging.getLogger("LiteLLM")
    # Create a handler pointing to stderr
    handler = logging.StreamHandler(sys.stderr)
    test_logger.addHandler(handler)

    try:
      self.assertIs(handler.stream, sys.stderr)

      _redirect_litellm_loggers_to_stdout()

      self.assertIs(handler.stream, sys.stdout)
    finally:
      # Clean up
      test_logger.removeHandler(handler)

  def test_preserves_stdout_handler(self):
    """Test that handlers already pointing to stdout are not modified."""
    test_logger = logging.getLogger("LiteLLM Proxy")
    # Create a handler already pointing to stdout
    handler = logging.StreamHandler(sys.stdout)
    test_logger.addHandler(handler)

    try:
      _redirect_litellm_loggers_to_stdout()

      self.assertIs(handler.stream, sys.stdout)
    finally:
      # Clean up
      test_logger.removeHandler(handler)

  def test_does_not_affect_non_stream_handlers(self):
    """Test that non-StreamHandler handlers are not affected."""
    test_logger = logging.getLogger("LiteLLM Router")
    # Create a FileHandler (not a StreamHandler)
    with tempfile.NamedTemporaryFile(delete=False) as temp_file:
      temp_file_name = temp_file.name
    with contextlib.closing(
        logging.FileHandler(temp_file_name)
    ) as file_handler:
      test_logger.addHandler(file_handler)

      try:
        _redirect_litellm_loggers_to_stdout()
        # FileHandler should not be modified (it doesn't point to stderr or stdout)
        self.assertEqual(file_handler.baseFilename, temp_file_name)
      finally:
        # Clean up
        test_logger.removeHandler(file_handler)
    os.unlink(temp_file_name)


@pytest.mark.parametrize(
    "logger_name",
    ["LiteLLM", "LiteLLM Proxy", "LiteLLM Router"],
    ids=["LiteLLM", "LiteLLM Proxy", "LiteLLM Router"],
)
def test_handles_litellm_logger_names(logger_name):
  """Test that LiteLLM logger names are processed."""
  test_logger = logging.getLogger(logger_name)
  handler = logging.StreamHandler(sys.stderr)
  test_logger.addHandler(handler)

  try:
    _redirect_litellm_loggers_to_stdout()

    assert handler.stream is sys.stdout
  finally:
    # Clean up
    test_logger.removeHandler(handler)


# ── Anthropic thinking_blocks tests ─────────────────────────────


def test_is_anthropic_provider():
  """Verify _is_anthropic_provider matches known Claude provider prefixes."""
  assert _is_anthropic_provider("anthropic")
  assert _is_anthropic_provider("bedrock")
  assert _is_anthropic_provider("vertex_ai")
  assert _is_anthropic_provider("ANTHROPIC")  # case-insensitive
  assert not _is_anthropic_provider("openai")
  assert not _is_anthropic_provider("")
  assert not _is_anthropic_provider(None)


@pytest.mark.parametrize(
    "model_string,expected",
    [
        ("anthropic/claude-4-sonnet", True),
        ("anthropic/claude-3-5-sonnet-20241022", True),
        ("Anthropic/Claude-4-Opus", True),
        ("bedrock/anthropic.claude-3-5-sonnet", True),
        ("bedrock/us.anthropic.claude-3-5-sonnet-20241022-v2:0", True),
        ("bedrock/claude-3-5-sonnet", True),
        ("vertex_ai/claude-3-5-sonnet@20241022", True),
        ("openai/gpt-4o", False),
        ("gemini/gemini-2.5-pro", False),
        ("vertex_ai/gemini-2.5-flash", False),
        ("bedrock/amazon.titan-text-express-v1", False),
    ],
    ids=[
        "anthropic-prefix",
        "anthropic-versioned",
        "anthropic-uppercase",
        "bedrock-anthropic-dot",
        "bedrock-us-anthropic",
        "bedrock-claude",
        "vertex-claude",
        "openai-no-match",
        "gemini-no-match",
        "vertex-gemini-no-match",
        "bedrock-non-anthropic",
    ],
)
def test_is_anthropic_model(model_string, expected):
  assert _is_anthropic_model(model_string) is expected


def test_extract_reasoning_value_prefers_thinking_blocks():
  """thinking_blocks (Anthropic format with signatures) take priority."""
  thinking_blocks = [
      {"type": "thinking", "thinking": "step 1", "signature": "c2lnX2E="},
      {"type": "thinking", "thinking": "step 2", "signature": "c2lnX2I="},
  ]
  message = {
      "role": "assistant",
      "content": "Answer",
      "thinking_blocks": thinking_blocks,
      "reasoning_content": "flat reasoning",
  }
  result = _extract_reasoning_value(message)
  assert result is thinking_blocks


def test_extract_reasoning_value_falls_back_without_thinking_blocks():
  """When thinking_blocks is absent, falls back to reasoning_content."""
  message = {
      "role": "assistant",
      "content": "Answer",
      "reasoning_content": "flat reasoning",
  }
  result = _extract_reasoning_value(message)
  assert result == "flat reasoning"


def test_convert_reasoning_value_to_parts_preserves_base64_signature():
  """Base64 signatures are decoded to raw bytes on thought parts."""
  thinking_blocks = [
      {"type": "thinking", "thinking": "step 1", "signature": "c2lnX2E="},
      {"type": "thinking", "thinking": "step 2", "signature": "c2lnX2I="},
  ]
  parts = _convert_reasoning_value_to_parts(thinking_blocks)
  assert len(parts) == 2
  assert parts[0].text == "step 1"
  assert parts[0].thought is True
  assert parts[0].thought_signature == b"sig_a"
  assert parts[1].text == "step 2"
  assert parts[1].thought_signature == b"sig_b"


def test_convert_reasoning_value_to_parts_raw_signature_falls_back_to_utf8():
  """Non-base64 signatures are preserved as utf-8 bytes."""
  thinking_blocks = [
      {"type": "thinking", "thinking": "step 1", "signature": "sig_raw"},
  ]
  parts = _convert_reasoning_value_to_parts(thinking_blocks)
  assert len(parts) == 1
  assert parts[0].text == "step 1"
  assert parts[0].thought_signature == b"sig_raw"


def test_convert_reasoning_value_to_parts_skips_redacted_blocks():
  """Redacted thinking blocks are excluded from parts."""
  thinking_blocks = [
      {"type": "thinking", "thinking": "visible", "signature": "c2lnMQ=="},
      {"type": "redacted", "data": "hidden"},
  ]
  parts = _convert_reasoning_value_to_parts(thinking_blocks)
  assert len(parts) == 1
  assert parts[0].text == "visible"


def test_convert_reasoning_value_to_parts_preserves_signature_only_blocks():
  """Signature-only blocks (empty text) are preserved for streaming aggregation.

  Anthropic emits the block_stop signature as a delta with empty thinking text.
  Dropping it would lose the signature, breaking multi-turn thinking continuity.
  Blocks with neither text nor signature are still skipped.
  """
  thinking_blocks = [
      {"type": "thinking", "thinking": "", "signature": "c2lnMQ=="},
      {"type": "thinking", "thinking": "real thought", "signature": "c2lnMg=="},
      {
          "type": "thinking",
          "thinking": "",
          "signature": "",
      },  # fully empty: drop
  ]
  parts = _convert_reasoning_value_to_parts(thinking_blocks)
  assert len(parts) == 2
  assert parts[0].text == ""
  assert parts[0].thought is True
  assert parts[0].thought_signature == b"sig1"
  assert parts[1].text == "real thought"
  assert parts[1].thought_signature == b"sig2"


def test_aggregate_streaming_thought_parts():
  """Tests aggregating fragmented streaming thought parts and multiple blocks."""
  parts = [
      types.Part(text="First block ", thought=True),
      types.Part(text="text.", thought=True),
      types.Part(text="", thought=True, thought_signature=b"sig1"),
      types.Part(text="Second block", thought=True, thought_signature=b"sig2"),
      types.Part(text="Trailing without sig", thought=True),
  ]
  aggregated = _aggregate_streaming_thought_parts(parts)
  assert len(aggregated) == 3
  assert aggregated[0].text == "First block text."
  assert aggregated[0].thought_signature == b"sig1"
  assert aggregated[1].text == "Second block"
  assert aggregated[1].thought_signature == b"sig2"
  assert aggregated[2].text == "Trailing without sig"
  assert aggregated[2].thought_signature is None


def test_convert_reasoning_value_to_parts_flat_string_unchanged():
  """Flat string reasoning still produces thought parts without signature."""
  parts = _convert_reasoning_value_to_parts("simple reasoning text")
  assert len(parts) == 1
  assert parts[0].text == "simple reasoning text"
  assert parts[0].thought is True
  assert parts[0].thought_signature is None


@pytest.mark.asyncio
async def test_content_to_message_param_anthropic_outputs_thinking_blocks():
  """Anthropic model messages base64-encode thought signatures."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(
              text="deep thought",
              thought=True,
              thought_signature=b"sig_round_trip",
          ),
          types.Part(text="Hello!"),
      ],
  )
  result = await _content_to_message_param(
      content, model="anthropic/claude-4-sonnet"
  )
  assert result["role"] == "assistant"
  assert result["thinking_blocks"] == [{
      "type": "thinking",
      "thinking": "deep thought",
      "signature": "c2lnX3JvdW5kX3RyaXA=",
  }]
  assert result.get("reasoning_content") is None
  assert result["content"] == "Hello!"


@pytest.mark.asyncio
async def test_content_to_message_param_anthropic_model_round_trip_preserves_signature():
  """Decoded signatures are re-encoded when rebuilding Anthropic messages."""
  response_message = {
      "role": "assistant",
      "content": "Final answer",
      "thinking_blocks": [{
          "type": "thinking",
          "thinking": "Let me reason...",
          "signature": "c2lnX2E=",
      }],
  }

  parts = _convert_reasoning_value_to_parts(
      _extract_reasoning_value(response_message)
  )
  content = types.Content(
      role="model",
      parts=parts + [types.Part(text="Final answer")],
  )

  result = await _content_to_message_param(
      content,
      provider="anthropic",
      model="anthropic/claude-4-sonnet",
  )

  assert result["thinking_blocks"] == [{
      "type": "thinking",
      "thinking": "Let me reason...",
      "signature": "c2lnX2E=",
  }]
  assert result.get("reasoning_content") is None


@pytest.mark.asyncio
async def test_content_to_message_param_anthropic_split_thinking_and_signature():
  """Combines separate thinking and signature parts into a single thinking_block."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="deep thought", thought=True),
          types.Part(
              text="", thought=True, thought_signature=b"sig_round_trip"
          ),
          types.Part(text="Hello!"),
      ],
  )
  result = await _content_to_message_param(
      content, model="anthropic/claude-4-sonnet"
  )
  assert result["role"] == "assistant"
  assert "thinking_blocks" in result
  assert result.get("reasoning_content") is None
  blocks = result["thinking_blocks"]
  assert len(blocks) == 1
  assert blocks[0]["type"] == "thinking"
  assert blocks[0]["thinking"] == "deep thought"
  assert blocks[0]["signature"] == "c2lnX3JvdW5kX3RyaXA="
  assert result["content"] == "Hello!"


@pytest.mark.asyncio
async def test_content_to_message_param_non_anthropic_uses_reasoning_content():
  """For non-Anthropic models, reasoning_content is used as before."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="thinking text", thought=True),
          types.Part(text="Answer"),
      ],
  )
  result = await _content_to_message_param(content, model="openai/gpt-4o")
  assert result["role"] == "assistant"
  assert result.get("reasoning_content") == "thinking text"
  assert "thinking_blocks" not in result


@pytest.mark.asyncio
async def test_anthropic_provider_thinking_blocks_round_trip():
  """End-to-end: thinking_blocks in response stay intact for Anthropic provider."""
  response_message = {
      "role": "assistant",
      "content": "Final answer",
      "thinking_blocks": [
          {
              "type": "thinking",
              "thinking": "Let me reason...",
              "signature": "c2lnX2E=",
          },
      ],
  }

  reasoning_value = _extract_reasoning_value(response_message)
  assert isinstance(reasoning_value, list)

  parts = _convert_reasoning_value_to_parts(reasoning_value)
  assert len(parts) == 1
  assert parts[0].thought_signature == b"sig_a"

  all_parts = parts + [
      types.Part(text="Final answer"),
      types.Part.from_function_call(name="add", args={"a": 1, "b": 2}),
  ]
  content = types.Content(role="model", parts=all_parts)

  msg = await _content_to_message_param(content, provider="anthropic")
  assert isinstance(msg["content"], list)
  assert msg["content"][0] == {
      "type": "thinking",
      "thinking": "Let me reason...",
      "signature": "c2lnX2E=",
  }
  assert msg["content"][1] == {"type": "text", "text": "Final answer"}
  assert msg["tool_calls"] is not None
  assert len(msg["tool_calls"]) == 1
  assert msg["tool_calls"][0]["function"]["name"] == "add"
  assert msg.get("reasoning_content") is None


@pytest.mark.asyncio
async def test_content_to_message_param_anthropic_no_signature_falls_back():
  """Anthropic model with thought parts but no signatures uses reasoning_content."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="thinking without sig", thought=True),
          types.Part(text="Response"),
      ],
  )
  result = await _content_to_message_param(
      content, model="anthropic/claude-4-sonnet"
  )
  assert result.get("reasoning_content") == "thinking without sig"
  assert "thinking_blocks" not in result


@pytest.mark.parametrize(
    "provider,model,expected",
    [
        ("anthropic", "anthropic/claude-3-5-sonnet", True),
        ("anthropic", "", True),  # anthropic always routes to Claude
        ("bedrock", "bedrock/anthropic.claude-3-5-sonnet", True),
        ("bedrock", "bedrock/meta.llama3-70b-instruct-v1:0", False),
        ("vertex_ai", "vertex_ai/claude-3-5-sonnet@20241022", True),
        ("vertex_ai", "vertex_ai/gemini-2.5-flash", False),
        ("openai", "openai/gpt-4o", False),
        ("", "", False),
    ],
)
def test_is_anthropic_route(provider, model, expected):
  assert _is_anthropic_route(provider, model) is expected


def test_convert_reasoning_value_to_parts_empty_thinking_does_not_fall_through():
  """An empty thinking block is skipped, not parsed via the text fallback."""
  thinking_blocks = [
      {
          "type": "thinking",
          "thinking": "",
          "text": "leaked",
          "signature": "",
      },
  ]
  parts = _convert_reasoning_value_to_parts(thinking_blocks)
  assert parts == []


@pytest.mark.asyncio
async def test_content_to_message_param_bedrock_non_claude_no_thinking_blocks():
  """bedrock + non-Claude model must not get Anthropic thinking-block formatting."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="thinking text", thought=True),
          types.Part(text="Answer"),
      ],
  )
  result = await _content_to_message_param(
      content,
      provider="bedrock",
      model="bedrock/meta.llama3-70b-instruct-v1:0",
  )
  assert result.get("reasoning_content") == "thinking text"
  assert "thinking_blocks" not in result
  body = result.get("content")
  assert not (
      isinstance(body, list)
      and any(isinstance(b, dict) and b.get("type") == "thinking" for b in body)
  )


@pytest.mark.asyncio
async def test_content_to_message_param_bedrock_claude_embeds_thinking_blocks():
  """bedrock + Claude model embeds thinking blocks in the content list."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="thinking text", thought=True),
          types.Part(text="Answer"),
      ],
  )
  result = await _content_to_message_param(
      content,
      provider="bedrock",
      model="bedrock/anthropic.claude-3-5-sonnet-20241022-v2:0",
  )
  assert isinstance(result["content"], list)
  assert result["content"][0] == {
      "type": "thinking",
      "thinking": "thinking text",
  }
  assert result.get("reasoning_content") is None


@pytest.mark.asyncio
async def test_content_to_message_param_vertex_gemini_no_thinking_blocks():
  """vertex_ai + Gemini model must not get Anthropic thinking-block formatting."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="thinking text", thought=True),
          types.Part(text="Answer"),
      ],
  )
  result = await _content_to_message_param(
      content,
      provider="vertex_ai",
      model="vertex_ai/gemini-2.5-flash",
  )
  assert result.get("reasoning_content") == "thinking text"
  assert "thinking_blocks" not in result
  body = result.get("content")
  assert not (
      isinstance(body, list)
      and any(isinstance(b, dict) and b.get("type") == "thinking" for b in body)
  )


@pytest.mark.asyncio
async def test_content_to_message_param_anthropic_provider_embeds_thinking_blocks():
  """provider 'anthropic' always embeds thinking blocks in the content list."""
  content = types.Content(
      role="model",
      parts=[
          types.Part(text="thinking text", thought=True),
          types.Part(text="Answer"),
      ],
  )
  result = await _content_to_message_param(
      content,
      provider="anthropic",
      model="anthropic/claude-3-5-sonnet",
  )
  assert isinstance(result["content"], list)
  assert result["content"][0] == {
      "type": "thinking",
      "thinking": "thinking text",
  }
  assert result.get("reasoning_content") is None


@pytest.mark.asyncio
async def test_content_to_message_param_anthropic_aggregates_streaming_split_thinking():
  """Streaming splits one Anthropic thinking block across many parts:
  text-only chunks followed by a signature-only chunk at block_stop.
  _content_to_message_param must re-join them into one thinking_block.
  """
  content = types.Content(
      role="model",
      parts=[
          # Text-only chunks from streaming deltas (no signature)
          types.Part(text="The user wants ", thought=True),
          types.Part(text="GST research ", thought=True),
          types.Part(text="on secondment.", thought=True),
          # Final signature-only chunk (empty text, signature carries the whole block)
          types.Part(
              text="", thought=True, thought_signature=b"ErEDClsIDBACGAIfull"
          ),
          # Non-thought response content
          types.Part.from_function_call(name="create_plan", args={"q": "test"}),
      ],
  )
  result = await _content_to_message_param(
      content, model="anthropic/claude-4-sonnet"
  )
  # One aggregated thinking block with combined text and the block's signature
  blocks = result["thinking_blocks"]
  assert len(blocks) == 1
  assert blocks[0]["type"] == "thinking"
  assert blocks[0]["thinking"] == "The user wants GST research on secondment."
  assert blocks[0]["signature"] == "RXJFRENsc0lEQkFDR0FJZnVsbA=="
  # Legacy reasoning_content is not set when the Anthropic branch takes
  assert result.get("reasoning_content") is None


def test_model_response_to_chunk_preserves_signature_only_delta():
  """Anthropic streams a final thinking delta where content and
  reasoning_content are empty but thinking_blocks carries the signature.
  _has_meaningful_signal must recognize thinking_blocks as signal so the
  signature survives into a ReasoningChunk.
  """
  stream = ModelResponseStream(
      id="x",
      created=0,
      model="claude",
      choices=[
          StreamingChoices(
              index=0,
              delta=Delta(
                  role=None,
                  content="",
                  reasoning_content="",
                  thinking_blocks=[{
                      "type": "thinking",
                      "thinking": "",
                      "signature": "SignatureOnlyChunk",
                  }],
              ),
          )
      ],
  )
  chunks = list(_model_response_to_chunk(stream))
  reasoning_chunks = [c for c, _ in chunks if isinstance(c, ReasoningChunk)]
  assert len(reasoning_chunks) == 1
  parts = reasoning_chunks[0].parts
  assert len(parts) == 1
  assert parts[0].text == ""
  assert parts[0].thought is True
  assert parts[0].thought_signature == b"SignatureOnlyChunk"


@pytest.mark.asyncio
@pytest.mark.parametrize(
    "log_level,should_call",
    [
        (logging.WARNING, False),
        (logging.INFO, False),
        (logging.DEBUG, True),
    ],
)
async def test_generate_content_async_skips_request_log_build_above_debug(
    mock_acompletion, lite_llm_instance, log_level, should_call
):
  del mock_acompletion  # unused; lite_llm_instance is wired to it
  litellm_logger = logging.getLogger("google_adk.google.adk.models.lite_llm")
  original_level = litellm_logger.level
  litellm_logger.setLevel(log_level)
  try:
    with patch(
        "google.adk.models.lite_llm._build_request_log",
        return_value="log",
    ) as mock_build:
      async for _ in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION
      ):
        pass

      assert mock_build.called is should_call
  finally:
    litellm_logger.setLevel(original_level)


@pytest.mark.parametrize(
    "file_uri, expected",
    [
        ("file-abc123", True),
        ("assistant-abc123", True),
        ("https://example.com/file.pdf", False),
        ("not-a-file-id", False),
        ("", False),
        ("FILE-abc123", False),
    ],
)
def test_looks_like_openai_file_id(file_uri, expected):
  """Both `file-` and `assistant-` (Azure assistants) prefixes count as OpenAI file IDs."""
  assert _looks_like_openai_file_id(file_uri) is expected


@pytest.mark.parametrize(
    "file_uri, expected",
    [
        ("file-abc123", "file-<redacted>"),
        ("assistant-abc123", "assistant-<redacted>"),
    ],
)
def test_redact_file_uri_for_log_openai_prefixes(file_uri, expected):
  """OpenAI-style IDs are redacted while preserving the prefix kind."""
  assert _redact_file_uri_for_log(file_uri) == expected


def test_redact_file_uri_for_log_uses_display_name_when_provided():
  assert (
      _redact_file_uri_for_log("file-abc123", display_name="my.pdf") == "my.pdf"
  )


def test_redact_file_uri_for_log_http_url_keeps_scheme_and_tail():
  assert (
      _redact_file_uri_for_log("https://example.com/path/file.pdf")
      == "https://<redacted>/file.pdf"
  )


@pytest.mark.asyncio
async def test_generate_content_async_passes_http_options_headers_as_extra_headers(
    mock_acompletion, lite_llm_instance
):
  """Test that http_options.headers from LlmRequest are forwarded to litellm."""
  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          http_options=types.HttpOptions(
              headers={"X-User-Id": "user-123", "X-Trace-Id": "trace-abc"}
          )
      ),
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert "extra_headers" in kwargs
  assert kwargs["extra_headers"]["X-User-Id"] == "user-123"
  assert kwargs["extra_headers"]["X-Trace-Id"] == "trace-abc"


@pytest.mark.asyncio
async def test_generate_content_async_merges_http_options_with_existing_extra_headers(
    mock_response,
):
  """Test that http_options.headers merge with pre-existing extra_headers."""
  mock_acompletion = AsyncMock(return_value=mock_response)
  mock_client = MockLLMClient(mock_acompletion, Mock())
  # Create instance with pre-existing extra_headers via kwargs
  lite_llm_with_extra = LiteLlm(
      model="test_model",
      llm_client=mock_client,
      extra_headers={"X-Api-Key": "secret-key"},
  )

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          http_options=types.HttpOptions(headers={"X-User-Id": "user-456"})
      ),
  )

  async for _ in lite_llm_with_extra.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert "extra_headers" in kwargs
  # Both existing and new headers should be present
  assert kwargs["extra_headers"]["X-Api-Key"] == "secret-key"
  assert kwargs["extra_headers"]["X-User-Id"] == "user-456"


@pytest.mark.asyncio
async def test_generate_content_async_http_options_headers_override_existing(
    mock_response,
):
  """Test that http_options.headers override same-key extra_headers from init."""
  mock_acompletion = AsyncMock(return_value=mock_response)
  mock_client = MockLLMClient(mock_acompletion, Mock())
  lite_llm_with_extra = LiteLlm(
      model="test_model",
      llm_client=mock_client,
      extra_headers={"X-Override-Me": "old-value"},
  )

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          http_options=types.HttpOptions(headers={"X-Override-Me": "new-value"})
      ),
  )

  async for _ in lite_llm_with_extra.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  # Request-level headers should override init-level headers
  assert kwargs["extra_headers"]["X-Override-Me"] == "new-value"


@pytest.mark.asyncio
async def test_generate_content_async_passes_http_options_timeout(
    mock_acompletion, lite_llm_instance
):
  """Test that http_options.timeout is forwarded to litellm."""

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          http_options=types.HttpOptions(timeout=30000)
      ),
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert "timeout" in kwargs
  assert kwargs["timeout"] == 30000


@pytest.mark.asyncio
async def test_generate_content_async_passes_http_options_retry_options(
    mock_acompletion, lite_llm_instance
):
  """Test that http_options.retry_options is forwarded to litellm."""

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          http_options=types.HttpOptions(
              retry_options=types.HttpRetryOptions(
                  attempts=3,
              )
          )
      ),
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert "num_retries" in kwargs
  assert kwargs["num_retries"] == 3


@pytest.mark.asyncio
async def test_generate_content_async_passes_http_options_extra_body(
    mock_acompletion, lite_llm_instance
):
  """Test that http_options.extra_body is forwarded to litellm."""

  llm_request = LlmRequest(
      contents=[
          types.Content(
              role="user", parts=[types.Part.from_text(text="Test prompt")]
          )
      ],
      config=types.GenerateContentConfig(
          http_options=types.HttpOptions(
              extra_body={"custom_field": "custom_value", "priority": "high"}
          )
      ),
  )

  async for _ in lite_llm_instance.generate_content_async(llm_request):
    pass

  mock_acompletion.assert_called_once()
  _, kwargs = mock_acompletion.call_args
  assert "extra_body" in kwargs
  assert kwargs["extra_body"]["custom_field"] == "custom_value"
  assert kwargs["extra_body"]["priority"] == "high"


def _split_into_chunks(text, sizes):
  pieces = []
  pos = 0
  for size in sizes:
    pieces.append(text[pos : pos + size])
    pos += size
  if pos < len(text):
    pieces.append(text[pos:])
  return pieces


def test_brace_depth_tracker_simple_object():
  tracker = _BraceDepthTracker()
  assert tracker.feed('{"a": 1}') is True


def test_brace_depth_tracker_empty_object():
  tracker = _BraceDepthTracker()
  assert tracker.feed("{}") is True


def test_brace_depth_tracker_only_opens():
  tracker = _BraceDepthTracker()
  assert tracker.feed('{"a": ') is False


def test_brace_depth_tracker_completes_after_more_fragments():
  tracker = _BraceDepthTracker()
  assert tracker.feed('{"a": ') is False
  assert tracker.feed('"b"') is False
  assert tracker.feed("}") is True


def test_brace_depth_tracker_nested_objects():
  tracker = _BraceDepthTracker()
  assert tracker.feed('{"a": {"b": {"c": 1}}}') is True


def test_brace_depth_tracker_nested_split_across_fragments():
  tracker = _BraceDepthTracker()
  fragments = _split_into_chunks(
      '{"a": {"b": {"c": 1}, "d": [1, 2, 3]}, "e": "f"}', [3, 5, 4, 7, 9, 2, 1]
  )
  closes = [tracker.feed(f) for f in fragments]
  assert sum(closes) == 1
  assert closes[-1] is True


def test_brace_depth_tracker_string_with_braces_ignored():
  tracker = _BraceDepthTracker()
  # Braces inside strings must not affect depth.
  assert tracker.feed('{"x": "{}{{}}"}') is True


def test_brace_depth_tracker_string_with_braces_split_across_fragments():
  tracker = _BraceDepthTracker()
  fragments = ['{"x": "', "abc{def", "}ghi", '"}']
  closes = [tracker.feed(f) for f in fragments]
  assert closes == [False, False, False, True]


def test_brace_depth_tracker_escaped_quote_in_string():
  tracker = _BraceDepthTracker()
  # Escaped quote should not end the string; the trailing } closes the obj.
  assert tracker.feed(r'{"x": "a\"b}c"}') is True


def test_brace_depth_tracker_escaped_backslash_then_quote_ends_string():
  tracker = _BraceDepthTracker()
  # \\ is an escaped backslash; the next " ends the string. Then } closes.
  assert tracker.feed(r'{"x": "a\\"}') is True


def test_brace_depth_tracker_escape_split_across_fragments():
  tracker = _BraceDepthTracker()
  # Backslash arrives in one fragment, the escaped quote in the next.
  fragments = ['{"x": "a', "\\", '"', 'b"}']
  closes = [tracker.feed(f) for f in fragments]
  assert closes == [False, False, False, True]


def test_brace_depth_tracker_two_consecutive_objects():
  tracker = _BraceDepthTracker()
  assert tracker.feed('{"a": 1}{"b": 2}') is True


def test_brace_depth_tracker_one_char_at_a_time():
  tracker = _BraceDepthTracker()
  text = '{"key": {"nested": "v{}al"}, "n": 42}'
  closes = [tracker.feed(ch) for ch in text]
  assert sum(closes) == 1
  assert closes[-1] is True


def test_brace_depth_tracker_leading_whitespace_ignored():
  tracker = _BraceDepthTracker()
  assert tracker.feed('   \n  {"a": 1}') is True


def _function_chunks_for_args(arg_fragments):
  return [
      FunctionChunk(
          id="call_xyz" if i == 0 else None,
          name="my_func" if i == 0 else None,
          args=fragment,
          index=0,
      )
      for i, fragment in enumerate(arg_fragments)
  ]


def _stream_chunks_from_function_chunks(function_chunks):
  stream = []
  for chunk in function_chunks:
    delta_kwargs = {"role": "assistant"}
    if chunk.args is not None:
      delta_kwargs["tool_calls"] = [
          ChatCompletionDeltaToolCall(
              type="function",
              id=chunk.id,
              function=Function(name=chunk.name, arguments=chunk.args),
              index=chunk.index,
          )
      ]
    stream.append(
        ModelResponseStream(
            choices=[
                StreamingChoices(
                    finish_reason=None, delta=Delta(**delta_kwargs)
                )
            ]
        )
    )
  stream.append(
      ModelResponseStream(
          choices=[StreamingChoices(finish_reason="tool_calls", delta=Delta())]
      )
  )
  return stream


@pytest.mark.asyncio
async def test_streaming_tool_call_args_assembled_from_many_fragments(
    mock_completion, lite_llm_instance
):
  full_args = (
      '{"city": "San Francisco", "details": {"radius": 5, "tags": ["a{}",'
      ' "b\\"c"]}}'
  )
  fragments = _split_into_chunks(full_args, [4, 6, 1, 8, 3, 11, 2, 9, 7, 5, 1])
  mock_completion.return_value = iter(
      _stream_chunks_from_function_chunks(_function_chunks_for_args(fragments))
  )

  responses = [
      r
      async for r in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]

  assert len(responses) == 1
  function_call = responses[0].content.parts[0].function_call
  assert function_call.name == "my_func"
  assert function_call.id == "call_xyz"
  assert function_call.args == json.loads(full_args)


async def _count_full_buffer_loads(
    lite_llm_instance, mock_completion, full_args, fragments
):
  mock_completion.return_value = iter(
      _stream_chunks_from_function_chunks(_function_chunks_for_args(fragments))
  )
  real_loads = json.loads
  with patch(
      "google.adk.models.lite_llm.json.loads", side_effect=real_loads
  ) as patched_loads:
    responses = [
        r
        async for r in lite_llm_instance.generate_content_async(
            LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
        )
    ]
  full_buffer_calls = [
      c
      for c in patched_loads.call_args_list
      if c.args and c.args[0] == full_args
  ]
  return responses, len(full_buffer_calls)


@pytest.mark.asyncio
async def test_streaming_tool_call_json_loads_count_independent_of_fragment_count(
    mock_completion, lite_llm_instance
):
  # The previous implementation called json.loads(buffer) after every
  # fragment, so the count grew with the fragment count (O(N) calls and
  # O(N^2) total parse cost). The fix makes the count constant.
  full_args = '{"a": 1, "b": {"c": 2}}'
  one_chunk = [full_args]
  one_char_at_a_time = _split_into_chunks(full_args, [1] * len(full_args))

  _, count_one_chunk = await _count_full_buffer_loads(
      lite_llm_instance, mock_completion, full_args, one_chunk
  )
  _, count_many_chunks = await _count_full_buffer_loads(
      lite_llm_instance, mock_completion, full_args, one_char_at_a_time
  )
  assert count_one_chunk == count_many_chunks


@pytest.mark.asyncio
async def test_streaming_tool_call_brace_in_string_does_not_falsely_complete(
    mock_completion, lite_llm_instance
):
  # The closing brace inside the string must not advance fallback_index.
  # If it did, the second tool call would be merged into a single bucket
  # and the assembled args would be invalid.
  full_args_a = '{"text": "a{b}c"}'
  full_args_b = '{"x": 1}'
  fragments_a = _split_into_chunks(full_args_a, [1] * len(full_args_a))
  fragments_b = _split_into_chunks(full_args_b, [1] * len(full_args_b))

  function_chunks = _function_chunks_for_args(fragments_a)
  # Second tool call: provider emits no index, relies on fallback_index advance.
  for i, fragment in enumerate(fragments_b):
    function_chunks.append(
        FunctionChunk(
            id="call_2" if i == 0 else None,
            name="other_func" if i == 0 else None,
            args=fragment,
            index=0,
        )
    )

  mock_completion.return_value = iter(
      _stream_chunks_from_function_chunks(function_chunks)
  )

  responses = [
      r
      async for r in lite_llm_instance.generate_content_async(
          LLM_REQUEST_WITH_FUNCTION_DECLARATION, stream=True
      )
  ]

  assert len(responses) == 1
  parts = responses[0].content.parts
  assert len(parts) == 2
  args_by_name = {p.function_call.name: p.function_call.args for p in parts}
  assert args_by_name["my_func"] == json.loads(full_args_a)
  assert args_by_name["other_func"] == json.loads(full_args_b)
