forked from phoenix/litellm-mirror
6005450c8f
* fix(utils.py): don't return 'none' response headers Fixes https://github.com/BerriAI/litellm/issues/6123 * fix(vertex_and_google_ai_studio_gemini.py): support parsing out additional properties and strict value for tool calls Fixes https://github.com/BerriAI/litellm/issues/6136 * fix(cost_calculator.py): set default character value to none Fixes https://github.com/BerriAI/litellm/issues/6133#issuecomment-2403290196 * fix(google.py): fix cost per token / cost per char conversion Fixes https://github.com/BerriAI/litellm/issues/6133#issuecomment-2403370287 * build(model_prices_and_context_window.json): update gemini pricing Fixes https://github.com/BerriAI/litellm/issues/6133 * build(model_prices_and_context_window.json): update gemini pricing * fix(litellm_logging.py): fix streaming caching logging when 'turn_off_message_logging' enabled Stores unredacted response in cache * build(model_prices_and_context_window.json): update gemini-1.5-flash pricing * fix(cost_calculator.py): fix default prompt_character count logic Fixes error in gemini cost calculation * fix(cost_calculator.py): fix cost calc for tts models
83 lines
2.3 KiB
Python
83 lines
2.3 KiB
Python
import json
|
|
import os
|
|
import sys
|
|
import traceback
|
|
|
|
from dotenv import load_dotenv
|
|
|
|
load_dotenv()
|
|
import io
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
sys.path.insert(
|
|
0, os.path.abspath("../..")
|
|
) # Adds the parent directory to the system path
|
|
|
|
import litellm
|
|
|
|
|
|
def test_completion_pydantic_obj_2():
|
|
from pydantic import BaseModel
|
|
from litellm.llms.custom_httpx.http_handler import HTTPHandler
|
|
|
|
litellm.set_verbose = True
|
|
|
|
class CalendarEvent(BaseModel):
|
|
name: str
|
|
date: str
|
|
participants: list[str]
|
|
|
|
class EventsList(BaseModel):
|
|
events: list[CalendarEvent]
|
|
|
|
messages = [
|
|
{"role": "user", "content": "List important events from the 20th century."}
|
|
]
|
|
expected_request_body = {
|
|
"contents": [
|
|
{
|
|
"role": "user",
|
|
"parts": [{"text": "List important events from the 20th century."}],
|
|
}
|
|
],
|
|
"generationConfig": {
|
|
"response_mime_type": "application/json",
|
|
"response_schema": {
|
|
"properties": {
|
|
"events": {
|
|
"items": {
|
|
"properties": {
|
|
"name": {"type": "string"},
|
|
"date": {"type": "string"},
|
|
"participants": {
|
|
"items": {"type": "string"},
|
|
"type": "array",
|
|
},
|
|
},
|
|
"type": "object",
|
|
},
|
|
"type": "array",
|
|
}
|
|
},
|
|
"type": "object",
|
|
},
|
|
},
|
|
}
|
|
client = HTTPHandler()
|
|
with patch.object(client, "post", new=MagicMock()) as mock_post:
|
|
mock_post.return_value = expected_request_body
|
|
try:
|
|
litellm.completion(
|
|
model="gemini/gemini-1.5-pro",
|
|
messages=messages,
|
|
response_format=EventsList,
|
|
client=client,
|
|
)
|
|
except Exception as e:
|
|
print(e)
|
|
|
|
mock_post.assert_called_once()
|
|
|
|
print(mock_post.call_args.kwargs)
|
|
|
|
assert mock_post.call_args.kwargs["json"] == expected_request_body
|