From be66800a986baa3a58764fd68ad3864b7eba3a39 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 13:31:46 -0700 Subject: [PATCH 01/14] feat(main.py): initial commit - refactoring google ai studio to just use vertex httpx Uses the same calling logic for google ai studio/vertex ai. Simplifies logic, gives google ai studio integration all of vertex ai features. --- litellm/__init__.py | 3 + litellm/_logging.py | 6 +- litellm/llms/vertex_httpx.py | 137 +++++++++++++++++------- litellm/main.py | 173 +++++++++++++++---------------- litellm/tests/conftest.py | 6 +- litellm/tests/test_completion.py | 22 ++-- litellm/utils.py | 6 +- 7 files changed, 216 insertions(+), 137 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index e0378d2ed0d..65c304bd135 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -13,7 +13,10 @@ from litellm._logging import ( verbose_logger, json_logs, _turn_on_json, + log_level, ) + + from litellm.proxy._types import ( KeyManagementSystem, KeyManagementSettings, diff --git a/litellm/_logging.py b/litellm/_logging.py index 52a445b49e4..c4d7c035a0e 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -1,6 +1,8 @@ -import logging, os, json -from logging import Formatter +import json +import logging +import os import traceback +from logging import Formatter set_verbose = False diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index c9e48f3e175..79d79567007 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -1,41 +1,47 @@ # What is this? ## httpx client for vertex ai calls ## Initial implementation - covers gemini + image gen calls -from functools import partial -import os, types +import inspect import json -from enum import Enum -import requests # type: ignore +import os import time -from typing import Callable, Optional, Union, List, Any, Tuple +import types +import uuid +from enum import Enum +from functools import partial +from typing import Any, Callable, List, Literal, Optional, Tuple, Union + +import httpx # type: ignore +import requests # type: ignore + +import litellm import litellm.litellm_core_utils import litellm.litellm_core_utils.litellm_logging -from litellm.utils import ModelResponse, Usage, CustomStreamWrapper from litellm.litellm_core_utils.core_helpers import map_finish_reason -import litellm, uuid -import httpx, inspect # type: ignore from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler -from .base import BaseLLM -from litellm.types.llms.vertex_ai import ( - ContentType, - SystemInstructions, - PartType, - RequestBody, - GenerateContentResponseBody, - FunctionCallingConfig, - FunctionDeclaration, - Tools, - ToolConfig, - GenerationConfig, -) from litellm.llms.vertex_ai import _gemini_convert_messages_with_history -from litellm.types.utils import GenericStreamingChunk from litellm.types.llms.openai import ( - ChatCompletionUsageBlock, + ChatCompletionResponseMessage, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, - ChatCompletionResponseMessage, + ChatCompletionUsageBlock, ) +from litellm.types.llms.vertex_ai import ( + ContentType, + FunctionCallingConfig, + FunctionDeclaration, + GenerateContentResponseBody, + GenerationConfig, + PartType, + RequestBody, + SystemInstructions, + ToolConfig, + Tools, +) +from litellm.types.utils import GenericStreamingChunk +from litellm.utils import CustomStreamWrapper, ModelResponse, Usage + +from .base import BaseLLM class VertexGeminiConfig: @@ -414,9 +420,11 @@ class VertexLLM(BaseLLM): def load_auth( self, credentials: Optional[str], project_id: Optional[str] ) -> Tuple[Any, str]: - from google.auth.transport.requests import Request # type: ignore[import-untyped] - from google.auth.credentials import Credentials # type: ignore[import-untyped] import google.auth as google_auth + from google.auth.credentials import Credentials # type: ignore[import-untyped] + from google.auth.transport.requests import ( + Request, # type: ignore[import-untyped] + ) if credentials is not None and isinstance(credentials, str): import google.oauth2.service_account @@ -449,7 +457,9 @@ class VertexLLM(BaseLLM): return creds, project_id def refresh_auth(self, credentials: Any) -> None: - from google.auth.transport.requests import Request # type: ignore[import-untyped] + from google.auth.transport.requests import ( + Request, # type: ignore[import-untyped] + ) credentials.refresh(Request()) @@ -482,6 +492,50 @@ class VertexLLM(BaseLLM): return self._credentials.token, self.project_id + def _get_token_and_url( + self, + model: str, + gemini_api_key: Optional[str], + vertex_project: Optional[str], + vertex_location: Optional[str], + vertex_credentials: Optional[str], + stream: Optional[bool], + custom_llm_provider: Literal["vertex_ai", "vertex_ai_beta", "gemini"], + ) -> Tuple[Optional[str], str]: + """ + Internal function. Returns the token and url for the call. + + Handles logic if it's google ai studio vs. vertex ai. + + Returns + token, url + """ + if custom_llm_provider == "gemini": + _gemini_model_name = "models/{}".format(model) + auth_header = None + endpoint = "generateContent" + if stream is True: + endpoint = "streamGenerateContent" + + url = ( + "https://generativelanguage.googleapis.com/v1beta/{}:{}?key={}".format( + _gemini_model_name, endpoint, gemini_api_key + ) + ) + else: + auth_header, vertex_project = self._ensure_access_token( + credentials=vertex_credentials, project_id=vertex_project + ) + vertex_location = self.get_vertex_region(vertex_region=vertex_location) + + ### SET RUNTIME ENDPOINT ### + endpoint = "generateContent" + if stream is True: + endpoint = "streamGenerateContent" + url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:{endpoint}" + + return auth_header, url + async def async_streaming( self, model: str, @@ -574,6 +628,9 @@ class VertexLLM(BaseLLM): messages: list, model_response: ModelResponse, print_verbose: Callable, + custom_llm_provider: Literal[ + "vertex_ai", "vertex_ai_beta", "gemini" + ], # if it's vertex_ai or gemini (google ai studio) encoding, logging_obj, optional_params: dict, @@ -582,20 +639,23 @@ class VertexLLM(BaseLLM): vertex_project: Optional[str], vertex_location: Optional[str], vertex_credentials: Optional[str], + gemini_api_key: Optional[str], litellm_params=None, logger_fn=None, extra_headers: Optional[dict] = None, client: Optional[Union[AsyncHTTPHandler, HTTPHandler]] = None, ) -> Union[ModelResponse, CustomStreamWrapper]: - - auth_header, vertex_project = self._ensure_access_token( - credentials=vertex_credentials, project_id=vertex_project - ) - vertex_location = self.get_vertex_region(vertex_region=vertex_location) stream: Optional[bool] = optional_params.pop("stream", None) # type: ignore - ### SET RUNTIME ENDPOINT ### - url = f"https://{vertex_location}-aiplatform.googleapis.com/v1/projects/{vertex_project}/locations/{vertex_location}/publishers/google/models/{model}:generateContent" + auth_header, url = self._get_token_and_url( + model=model, + gemini_api_key=gemini_api_key, + vertex_project=vertex_project, + vertex_location=vertex_location, + vertex_credentials=vertex_credentials, + stream=stream, + custom_llm_provider=custom_llm_provider, + ) ## TRANSFORMATION ## # Separate system prompt from rest of message @@ -609,14 +669,16 @@ class VertexLLM(BaseLLM): if len(system_prompt_indices) > 0: for idx in reversed(system_prompt_indices): messages.pop(idx) - system_instructions = SystemInstructions(parts=system_content_blocks) content = _gemini_convert_messages_with_history(messages=messages) tools: Optional[Tools] = optional_params.pop("tools", None) tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None) generation_config: Optional[GenerationConfig] = GenerationConfig( **optional_params ) - data = RequestBody(system_instruction=system_instructions, contents=content) + data = RequestBody(contents=content) + if len(system_content_blocks) > 0: + system_instructions = SystemInstructions(parts=system_content_blocks) + data["system_instruction"] = system_instructions if tools is not None: data["tools"] = tools if tool_choice is not None: @@ -626,8 +688,9 @@ class VertexLLM(BaseLLM): headers = { "Content-Type": "application/json; charset=utf-8", - "Authorization": f"Bearer {auth_header}", } + if auth_header is not None: + headers["Authorization"] = f"Bearer {auth_header}" ## LOGGING logging_obj.pre_call( diff --git a/litellm/main.py b/litellm/main.py index 77fe38fd2d8..6921dabab6f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -7,107 +7,130 @@ # # Thank you ! We ❤️ you! - Krrish & Ishaan -import os, openai, sys, json, inspect, uuid, datetime, threading -from typing import Any, Literal, Union, BinaryIO -from typing_extensions import overload -from functools import partial - -import dotenv, traceback, random, asyncio, time, contextvars +import asyncio +import contextvars +import datetime +import inspect +import json +import os +import random +import sys +import threading +import time +import traceback +import uuid +from concurrent.futures import ThreadPoolExecutor from copy import deepcopy +from functools import partial +from typing import ( + Any, + BinaryIO, + Callable, + Dict, + List, + Literal, + Mapping, + Optional, + Union, +) + +import dotenv import httpx +import openai +import tiktoken +from typing_extensions import overload + import litellm -from ._logging import verbose_logger from litellm import ( # type: ignore + Logging, client, exception_type, - get_optional_params, get_litellm_params, - Logging, + get_optional_params, ) from litellm.utils import ( - get_secret, CustomStreamWrapper, - read_config_args, - completion_with_fallbacks, - get_llm_provider, - get_api_key, - mock_completion_streaming_obj, + Usage, async_mock_completion_streaming_obj, + completion_with_fallbacks, convert_to_model_response_object, - token_counter, create_pretrained_tokenizer, create_tokenizer, - Usage, + get_api_key, + get_llm_provider, get_optional_params_embeddings, get_optional_params_image_gen, + get_secret, + mock_completion_streaming_obj, + read_config_args, supports_httpx_timeout, + token_counter, ) + +from ._logging import verbose_logger +from .caching import disable_cache, enable_cache, update_cache from .llms import ( - anthropic_text, - together_ai, ai21, - sagemaker, - bedrock, - triton, - huggingface_restapi, - replicate, aleph_alpha, - nlp_cloud, + anthropic_text, baseten, - vllm, - ollama, - ollama_chat, - cloudflare, + bedrock, clarifai, + cloudflare, cohere, cohere_chat, - petals, + gemini, + huggingface_restapi, + maritalk, + nlp_cloud, + ollama, + ollama_chat, oobabooga, openrouter, palm, - gemini, + petals, + replicate, + sagemaker, + together_ai, + triton, vertex_ai, vertex_ai_anthropic, - maritalk, + vllm, watsonx, ) -from .llms.openai import OpenAIChatCompletion, OpenAITextCompletion -from .llms.azure import AzureChatCompletion -from .llms.databricks import DatabricksChatCompletion -from .llms.azure_text import AzureTextCompletion from .llms.anthropic import AnthropicChatCompletion from .llms.anthropic_text import AnthropicTextCompletion +from .llms.azure import AzureChatCompletion +from .llms.azure_text import AzureTextCompletion +from .llms.bedrock_httpx import BedrockConverseLLM, BedrockLLM +from .llms.databricks import DatabricksChatCompletion from .llms.huggingface_restapi import Huggingface +from .llms.openai import OpenAIChatCompletion, OpenAITextCompletion from .llms.predibase import PredibaseChatCompletion -from .llms.bedrock_httpx import BedrockLLM, BedrockConverseLLM -from .llms.vertex_httpx import VertexLLM -from .llms.triton import TritonChatCompletion from .llms.prompt_templates.factory import ( - prompt_factory, custom_prompt, function_call_prompt, map_system_message_pt, + prompt_factory, ) -import tiktoken -from concurrent.futures import ThreadPoolExecutor -from typing import Callable, List, Optional, Dict, Union, Mapping -from .caching import enable_cache, disable_cache, update_cache +from .llms.triton import TritonChatCompletion +from .llms.vertex_httpx import VertexLLM from .types.llms.openai import HttpxBinaryResponseContent encoding = tiktoken.get_encoding("cl100k_base") from litellm.utils import ( - get_secret, + Choices, CustomStreamWrapper, - TextCompletionStreamWrapper, - ModelResponse, - TextCompletionResponse, - TextChoices, EmbeddingResponse, ImageResponse, - read_config_args, - Choices, Message, + ModelResponse, + TextChoices, + TextCompletionResponse, + TextCompletionStreamWrapper, TranscriptionResponse, + get_secret, + read_config_args, ) ####### ENVIRONMENT VARIABLES ################### @@ -1845,43 +1868,7 @@ def completion( ) return response response = model_response - elif custom_llm_provider == "gemini": - gemini_api_key = ( - api_key - or get_secret("GEMINI_API_KEY") - or get_secret("PALM_API_KEY") # older palm api key should also work - or litellm.api_key - ) - - # palm does not support streaming as yet :( - model_response = gemini.completion( - model=model, - messages=messages, - model_response=model_response, - print_verbose=print_verbose, - optional_params=optional_params, - litellm_params=litellm_params, - logger_fn=logger_fn, - encoding=encoding, - api_key=gemini_api_key, - logging_obj=logging, - acompletion=acompletion, - custom_prompt_dict=custom_prompt_dict, - ) - if ( - "stream" in optional_params - and optional_params["stream"] == True - and acompletion == False - ): - response = CustomStreamWrapper( - iter(model_response), - model, - custom_llm_provider="gemini", - logging_obj=logging, - ) - return response - response = model_response - elif custom_llm_provider == "vertex_ai_beta": + elif custom_llm_provider == "vertex_ai_beta" or custom_llm_provider == "gemini": vertex_ai_project = ( optional_params.pop("vertex_project", None) or optional_params.pop("vertex_ai_project", None) @@ -1899,6 +1886,14 @@ def completion( or optional_params.pop("vertex_ai_credentials", None) or get_secret("VERTEXAI_CREDENTIALS") ) + + gemini_api_key = ( + api_key + or get_secret("GEMINI_API_KEY") + or get_secret("PALM_API_KEY") # older palm api key should also work + or litellm.api_key + ) + new_params = deepcopy(optional_params) response = vertex_chat_completion.completion( # type: ignore model=model, @@ -1912,9 +1907,11 @@ def completion( vertex_location=vertex_ai_location, vertex_project=vertex_ai_project, vertex_credentials=vertex_credentials, + gemini_api_key=gemini_api_key, logging_obj=logging, acompletion=acompletion, timeout=timeout, + custom_llm_provider=custom_llm_provider, ) elif custom_llm_provider == "vertex_ai": diff --git a/litellm/tests/conftest.py b/litellm/tests/conftest.py index 8c2ce781f82..244ea07544c 100644 --- a/litellm/tests/conftest.py +++ b/litellm/tests/conftest.py @@ -1,7 +1,10 @@ # conftest.py -import pytest, sys, os import importlib +import os +import sys + +import pytest sys.path.insert( 0, os.path.abspath("../..") @@ -18,6 +21,7 @@ def setup_and_teardown(): sys.path.insert( 0, os.path.abspath("../..") ) # Adds the project directory to the system path + print("LITELLM_LOG - {}".format(os.getenv("LITELLM_LOG"))) import litellm from litellm import Router diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 91144684660..0134b659258 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -1,20 +1,26 @@ -import sys, os, json +import json +import os +import sys import traceback + from dotenv import load_dotenv load_dotenv() -import os, io +import io +import os sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path +import os +from unittest.mock import MagicMock, patch + import pytest + import litellm -from litellm import embedding, completion, completion_cost, Timeout -from litellm import RateLimitError +from litellm import RateLimitError, Timeout, completion, completion_cost, embedding +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.prompt_templates.factory import anthropic_messages_pt -from unittest.mock import patch, MagicMock -from litellm.llms.custom_httpx.http_handler import HTTPHandler, AsyncHTTPHandler # litellm.num_retries =3 litellm.cache = None @@ -1472,7 +1478,9 @@ def test_ollama_image(): data is untouched. """ - import io, base64 + import base64 + import io + from PIL import Image def mock_post(url, **kwargs): diff --git a/litellm/utils.py b/litellm/utils.py index 0b898165dfb..dd8dac709d4 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6972,7 +6972,9 @@ def exception_type( exception_mapping_worked = True if hasattr(original_exception, "request"): raise APIConnectionError( - message=f"{str(original_exception)}", + message="{}\n{}".format( + str(original_exception), traceback.format_exc() + ), llm_provider=custom_llm_provider, model=model, request=original_exception.request, @@ -7186,7 +7188,7 @@ def get_secret( else: raise ValueError( f"Google KMS requires the encrypted secret to be encoded in base64" - )#fix for this vulnerability https://huntr.com/bounties/ae623c2f-b64b-4245-9ed4-f13a0a5824ce + ) # fix for this vulnerability https://huntr.com/bounties/ae623c2f-b64b-4245-9ed4-f13a0a5824ce response = client.decrypt( request={ "name": litellm._google_kms_resource_name, From e92570534cab8522c2a3981b168c81b77bd2810b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 14:59:30 -0700 Subject: [PATCH 02/14] fix(vertex_httpx.py): support async streaming for google ai studio gemini --- litellm/llms/vertex_httpx.py | 68 ++++++++++++++++++++++++++++++------ 1 file changed, 57 insertions(+), 11 deletions(-) diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index 79d79567007..f3640c27c9d 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -12,6 +12,7 @@ from functools import partial from typing import Any, Callable, List, Literal, Optional, Tuple, Union import httpx # type: ignore +import ijson import requests # type: ignore import litellm @@ -257,7 +258,7 @@ async def make_call( raise VertexAIError(status_code=response.status_code, message=response.text) completion_stream = ModelResponseIterator( - streaming_response=response.aiter_bytes(chunk_size=2056) + streaming_response=response.aiter_bytes(), sync_stream=False ) # LOGGING logging_obj.post_call( @@ -288,7 +289,7 @@ def make_sync_call( raise VertexAIError(status_code=response.status_code, message=response.read()) completion_stream = ModelResponseIterator( - streaming_response=response.iter_bytes(chunk_size=2056) + streaming_response=response.iter_bytes(chunk_size=2056), sync_stream=True ) # LOGGING @@ -705,6 +706,25 @@ class VertexLLM(BaseLLM): ### ROUTING (ASYNC, STREAMING, SYNC) if acompletion: + ### ASYNC STREAMING + if stream is True: + return self.async_streaming( + model=model, + messages=messages, + data=json.dumps(data), # type: ignore + api_base=url, + model_response=model_response, + print_verbose=print_verbose, + encoding=encoding, + logging_obj=logging_obj, + optional_params=optional_params, + stream=stream, + litellm_params=litellm_params, + logger_fn=logger_fn, + headers=headers, + timeout=timeout, + client=client, # type: ignore + ) ### ASYNC COMPLETION return self.async_completion( model=model, @@ -916,9 +936,13 @@ class VertexLLM(BaseLLM): class ModelResponseIterator: - def __init__(self, streaming_response): + def __init__(self, streaming_response, sync_stream: bool): self.streaming_response = streaming_response - self.response_iterator = iter(self.streaming_response) + if sync_stream: + self.response_iterator = iter(self.streaming_response) + + self.events = ijson.sendable_list() + self.coro = ijson.items_coro(self.events, "item") def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: try: @@ -970,10 +994,21 @@ class ModelResponseIterator: def __next__(self): try: - chunk = next(self.response_iterator) - chunk = chunk.decode() - json_chunk = json.loads(chunk) - return self.chunk_parser(chunk=json_chunk) + chunk = self.response_iterator.__next__() + self.coro.send(chunk) + if self.events: + event = self.events[0] + json_chunk = event + self.events.clear() + return self.chunk_parser(chunk=json_chunk) + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=0, + tool_use=None, + ) except StopIteration: raise StopIteration except ValueError as e: @@ -987,9 +1022,20 @@ class ModelResponseIterator: async def __anext__(self): try: chunk = await self.async_response_iterator.__anext__() - chunk = chunk.decode() - json_chunk = json.loads(chunk) - return self.chunk_parser(chunk=json_chunk) + self.coro.send(chunk) + if self.events: + event = self.events[0] + json_chunk = event + self.events.clear() + return self.chunk_parser(chunk=json_chunk) + return GenericStreamingChunk( + text="", + is_finished=False, + finish_reason="", + usage=None, + index=0, + tool_use=None, + ) except StopAsyncIteration: raise StopAsyncIteration except ValueError as e: From a7d39d6bdd2e58d81ee52e56d24b1879abab635b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 16:04:40 -0700 Subject: [PATCH 03/14] build(pyproject.toml): add ijson as a package dep for google ai studio http streaming --- pyproject.toml | 1 + requirements.txt | 1 + 2 files changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 877cef85032..6d014101855 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,6 +27,7 @@ jinja2 = "^3.1.2" aiohttp = "*" requests = "^2.31.0" pydantic = "^2.0.0" +isjon = "*" uvicorn = {version = "^0.22.0", optional = true} gunicorn = {version = "^22.0.0", optional = true} diff --git a/requirements.txt b/requirements.txt index ab755fec335..fbf2bfc1d12 100644 --- a/requirements.txt +++ b/requirements.txt @@ -44,4 +44,5 @@ aiohttp==3.9.0 # for network calls aioboto3==12.3.0 # for async sagemaker calls tenacity==8.2.3 # for retrying requests, when litellm.num_retries set pydantic==2.7.1 # proxy + openai req. +ijson==3.2.3 # for google ai studio streaming #### \ No newline at end of file From d48f9e258e406a867a0706768cbbcb84ad548b3b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 16:21:49 -0700 Subject: [PATCH 04/14] test(test_completion.py): skip watson tests (account removed) --- litellm/tests/test_completion.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 0134b659258..7010fa4f489 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3308,6 +3308,7 @@ def test_mistral_anyscale_stream(): #### Test A121 ################### +@pytest.mark.skip(reason="Local test") def test_completion_ai21(): print("running ai21 j2light test") litellm.set_verbose = True @@ -3545,6 +3546,7 @@ def test_unified_auth_params(provider, model, project, region_name, token): assert value in translated_optional_params +@pytest.mark.skip(reason="Local test") @pytest.mark.asyncio async def test_acompletion_watsonx(): litellm.set_verbose = True @@ -3565,6 +3567,7 @@ async def test_acompletion_watsonx(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.skip(reason="Local test") @pytest.mark.asyncio async def test_acompletion_stream_watsonx(): litellm.set_verbose = True From 57d94792528d071184f413b04ddaf65c7188c869 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 16:26:10 -0700 Subject: [PATCH 05/14] build(config.yml): add ijson to ci/cd --- .circleci/config.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index 736bb8e8a14..f68ea2c5cdd 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -65,6 +65,7 @@ jobs: pip install "pydantic==2.7.1" pip install "diskcache==5.6.1" pip install "Pillow==10.3.0" + pip install "Pillow==3.2.3" - save_cache: paths: - ./venv From 1b215d704db2acec5a388652d35fbf207863d4bc Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 16:33:58 -0700 Subject: [PATCH 06/14] test: cleanup tests --- litellm/tests/test_completion.py | 2 ++ litellm/tests/test_streaming.py | 24 +++++++++++++++--------- 2 files changed, 17 insertions(+), 9 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index 7010fa4f489..79e859f5cfa 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3459,6 +3459,7 @@ def test_completion_palm_stream(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.skip(reason="Account deleted by IBM.") def test_completion_watsonx(): litellm.set_verbose = True model_name = "watsonx/ibm/granite-13b-chat-v2" @@ -3479,6 +3480,7 @@ def test_completion_watsonx(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.skip(reason="Skip test. account deleted.") def test_completion_stream_watsonx(): litellm.set_verbose = True model_name = "watsonx/ibm/granite-13b-chat-v2" diff --git a/litellm/tests/test_streaming.py b/litellm/tests/test_streaming.py index 622b2efc84d..ecb21b9f2b4 100644 --- a/litellm/tests/test_streaming.py +++ b/litellm/tests/test_streaming.py @@ -1,12 +1,17 @@ #### What this tests #### # This tests streaming for the completion endpoint -import sys, os, asyncio +import asyncio +import os +import sys +import time import traceback -import time, pytest, uuid -from pydantic import BaseModel +import uuid from typing import Tuple +import pytest +from pydantic import BaseModel + sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path @@ -15,12 +20,12 @@ from dotenv import load_dotenv load_dotenv() import litellm from litellm import ( - completion, - acompletion, AuthenticationError, BadRequestError, - RateLimitError, ModelResponse, + RateLimitError, + acompletion, + completion, ) litellm.logging = False @@ -1644,9 +1649,8 @@ def test_sagemaker_weird_response(): When the stream ends, flush any remaining holding chunks. """ try: - from litellm.llms.sagemaker import TokenIterator - import json import json + from litellm.llms.sagemaker import TokenIterator chunk = """[INST] Hey, how's it going? [/INST], @@ -1772,6 +1776,7 @@ def test_completion_sagemaker_stream(): pytest.fail(f"Error occurred: {e}") +@pytest.mark.skip(reason="Account deleted by IBM.") def test_completion_watsonx_stream(): litellm.set_verbose = True try: @@ -2631,9 +2636,10 @@ def test_success_callback_streaming(): # test_success_callback_streaming() +from typing import List, Optional + #### STREAMING + FUNCTION CALLING ### from pydantic import BaseModel -from typing import List, Optional class Function(BaseModel): From 377f8b77c0693baa5573e5d1e3f2da8452b78dd4 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 16:51:19 -0700 Subject: [PATCH 07/14] docs(gemini.py): add refactor note to code \ --- litellm/llms/gemini.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/litellm/llms/gemini.py b/litellm/llms/gemini.py index cfdf39eca2b..f48c4e29ea6 100644 --- a/litellm/llms/gemini.py +++ b/litellm/llms/gemini.py @@ -1,14 +1,22 @@ -import types -import traceback +#################################### +######### DEPRECATED FILE ########## +#################################### +# logic moved to `vertex_httpx.py` # + import copy import time +import traceback +import types from typing import Callable, Optional -from litellm.utils import ModelResponse, Choices, Message, Usage -import litellm + import httpx -from .prompt_templates.factory import prompt_factory, custom_prompt, get_system_prompt from packaging.version import Version + +import litellm from litellm import verbose_logger +from litellm.utils import Choices, Message, ModelResponse, Usage + +from .prompt_templates.factory import custom_prompt, get_system_prompt, prompt_factory class GeminiError(Exception): @@ -186,8 +194,8 @@ def completion( if _system_instruction and len(system_prompt) > 0: _params["system_instruction"] = system_prompt _model = genai.GenerativeModel(**_params) - if stream == True: - if acompletion == True: + if stream is True: + if acompletion is True: async def async_streaming(): try: From 3d9ef689e70a0afa0ed994ccc3753596e8b305bd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 17:30:38 -0700 Subject: [PATCH 08/14] fix(vertex_httpx.py): check if model supports system messages before sending separately --- litellm/__init__.py | 1 + litellm/_logging.py | 2 +- litellm/llms/vertex_httpx.py | 29 +++++--- litellm/router.py | 120 ++++++++++++++++++------------- litellm/tests/test_completion.py | 19 ++++- litellm/types/utils.py | 10 +++ litellm/utils.py | 82 +++++++++++++++++---- 7 files changed, 190 insertions(+), 73 deletions(-) diff --git a/litellm/__init__.py b/litellm/__init__.py index 2c4845c6ec7..6aee920c50e 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -739,6 +739,7 @@ from .utils import ( supports_function_calling, supports_parallel_function_calling, supports_vision, + supports_system_messages, get_litellm_params, acreate, get_model_list, diff --git a/litellm/_logging.py b/litellm/_logging.py index c4d7c035a0e..a98d85e1c47 100644 --- a/litellm/_logging.py +++ b/litellm/_logging.py @@ -12,7 +12,7 @@ if set_verbose is True: ) json_logs = bool(os.getenv("JSON_LOGS", False)) # Create a handler for the logger (you may need to adapt this based on your needs) -log_level = os.getenv("LITELLM_LOG", "ERROR") +log_level = os.getenv("LITELLM_LOG", "DEBUG") numeric_level: str = getattr(logging, log_level.upper()) handler = logging.StreamHandler() handler.setLevel(numeric_level) diff --git a/litellm/llms/vertex_httpx.py b/litellm/llms/vertex_httpx.py index f3640c27c9d..479e9bf3e2e 100644 --- a/litellm/llms/vertex_httpx.py +++ b/litellm/llms/vertex_httpx.py @@ -18,6 +18,7 @@ import requests # type: ignore import litellm import litellm.litellm_core_utils import litellm.litellm_core_utils.litellm_logging +from litellm import verbose_logger from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai import _gemini_convert_messages_with_history @@ -659,17 +660,29 @@ class VertexLLM(BaseLLM): ) ## TRANSFORMATION ## + try: + supports_system_message = litellm.supports_system_messages( + model=model, custom_llm_provider=custom_llm_provider + ) + except Exception as e: + verbose_logger.error( + "Unable to identify if system message supported. Defaulting to 'False'. Received error message - {}\nAdd it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json".format( + str(e) + ) + ) + supports_system_message = False # Separate system prompt from rest of message system_prompt_indices = [] system_content_blocks: List[PartType] = [] - for idx, message in enumerate(messages): - if message["role"] == "system": - _system_content_block = PartType(text=message["content"]) - system_content_blocks.append(_system_content_block) - system_prompt_indices.append(idx) - if len(system_prompt_indices) > 0: - for idx in reversed(system_prompt_indices): - messages.pop(idx) + if supports_system_message is True: + for idx, message in enumerate(messages): + if message["role"] == "system": + _system_content_block = PartType(text=message["content"]) + system_content_blocks.append(_system_content_block) + system_prompt_indices.append(idx) + if len(system_prompt_indices) > 0: + for idx in reversed(system_prompt_indices): + messages.pop(idx) content = _gemini_convert_messages_with_history(messages=messages) tools: Optional[Tools] = optional_params.pop("tools", None) tool_choice: Optional[ToolConfig] = optional_params.pop("tool_choice", None) diff --git a/litellm/router.py b/litellm/router.py index cd6c9c16eb1..db38df29f0f 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -7,66 +7,86 @@ # # Thank you ! We ❤️ you! - Krrish & Ishaan -import copy, httpx -from datetime import datetime -from typing import Dict, List, Optional, Union, Literal, Any, BinaryIO, Tuple, TypedDict -from typing_extensions import overload -import random, threading, time, traceback, uuid -import litellm, openai, hashlib, json -from litellm.caching import RedisCache, InMemoryCache, DualCache -import datetime as datetime_og -import logging, asyncio -import inspect, concurrent -from openai import AsyncOpenAI -from collections import defaultdict -from litellm.router_strategy.least_busy import LeastBusyLoggingHandler -from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler -from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler -from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler -from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 -from litellm.llms.custom_httpx.azure_dall_e_2 import ( - CustomHTTPTransport, - AsyncCustomHTTPTransport, -) -from litellm.utils import ( - ModelResponse, - CustomStreamWrapper, - get_utc_datetime, - calculate_max_parallel_requests, - _is_region_eu, -) +import asyncio +import concurrent import copy -from litellm._logging import verbose_router_logger +import datetime as datetime_og +import hashlib +import inspect +import json import logging -from litellm.types.utils import ModelInfo as ModelMapInfo -from litellm.types.router import ( - Deployment, - ModelInfo, - LiteLLM_Params, - RouterErrors, - updateDeployment, - updateLiteLLMParams, - RetryPolicy, - AllowedFailsPolicy, - AlertingConfig, - DeploymentTypedDict, - ModelGroupInfo, - AssistantsTypedDict, +import random +import threading +import time +import traceback +import uuid +from collections import defaultdict +from datetime import datetime +from typing import ( + Any, + BinaryIO, + Dict, + Iterable, + List, + Literal, + Optional, + Tuple, + TypedDict, + Union, ) + +import httpx +import openai +from openai import AsyncOpenAI +from typing_extensions import overload + +import litellm +from litellm._logging import verbose_router_logger +from litellm.caching import DualCache, InMemoryCache, RedisCache from litellm.integrations.custom_logger import CustomLogger from litellm.llms.azure import get_azure_ad_token_from_oidc +from litellm.llms.custom_httpx.azure_dall_e_2 import ( + AsyncCustomHTTPTransport, + CustomHTTPTransport, +) +from litellm.router_strategy.least_busy import LeastBusyLoggingHandler +from litellm.router_strategy.lowest_cost import LowestCostLoggingHandler +from litellm.router_strategy.lowest_latency import LowestLatencyLoggingHandler +from litellm.router_strategy.lowest_tpm_rpm import LowestTPMLoggingHandler +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 +from litellm.router_utils.handle_error import send_llm_exception_alert +from litellm.scheduler import FlowItem, Scheduler from litellm.types.llms.openai import ( - AsyncCursorPage, Assistant, - Thread, + AssistantToolParam, + AsyncCursorPage, Attachment, OpenAIMessage, Run, - AssistantToolParam, + Thread, +) +from litellm.types.router import ( + AlertingConfig, + AllowedFailsPolicy, + AssistantsTypedDict, + Deployment, + DeploymentTypedDict, + LiteLLM_Params, + ModelGroupInfo, + ModelInfo, + RetryPolicy, + RouterErrors, + updateDeployment, + updateLiteLLMParams, +) +from litellm.types.utils import ModelInfo as ModelMapInfo +from litellm.utils import ( + CustomStreamWrapper, + ModelResponse, + _is_region_eu, + calculate_max_parallel_requests, + get_utc_datetime, ) -from litellm.scheduler import Scheduler, FlowItem -from typing import Iterable -from litellm.router_utils.handle_error import send_llm_exception_alert class Router: @@ -3114,6 +3134,7 @@ class Router: # proxy support import os + import httpx # Check if the HTTP_PROXY and HTTPS_PROXY environment variables are set and use them accordingly. @@ -3800,6 +3821,7 @@ class Router: litellm_provider=llm_provider, mode="chat", supported_openai_params=supported_openai_params, + supports_system_messages=None, ) if model_group_info is None: diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index a60dd85070a..f1ee63564c2 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3392,15 +3392,28 @@ def test_completion_deep_infra_mistral(): # Gemini tests -def test_completion_gemini(): +@pytest.mark.parametrize( + "model", + [ + # "gemini-1.0-pro", + "gemini-1.5-pro", + # "gemini-1.5-flash", + ], +) +def test_completion_gemini(model): litellm.set_verbose = True - model_name = "gemini/gemini-1.5-pro-latest" - messages = [{"role": "user", "content": "Hey, how's it going?"}] + model_name = "gemini/{}".format(model) + messages = [ + {"role": "system", "content": "Be a good bot!"}, + {"role": "user", "content": "Hey, how's it going?"}, + ] try: response = completion(model=model_name, messages=messages) # Add any assertions,here to check the response print(response) assert response.choices[0]["index"] == 0 + + assert False except litellm.APIError as e: pass except Exception as e: diff --git a/litellm/types/utils.py b/litellm/types/utils.py index b7c0e318e44..f021fcd3452 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -1,13 +1,22 @@ import json import time import uuid +import json +import time +import uuid from enum import Enum from typing import Dict, List, Literal, Optional, Tuple, Union +from typing import Dict, List, Literal, Optional, Tuple, Union + from openai._models import BaseModel as OpenAIObject from pydantic import ConfigDict from typing_extensions import Dict, Required, TypedDict, override +from ..litellm_core_utils.core_helpers import map_finish_reason +from .llms.openai import ChatCompletionToolCallChunk, ChatCompletionUsageBlock +from typing_extensions import Dict, Required, TypedDict, override + from ..litellm_core_utils.core_helpers import map_finish_reason from .llms.openai import ChatCompletionToolCallChunk, ChatCompletionUsageBlock @@ -60,6 +69,7 @@ class ModelInfo(TypedDict, total=False): ] ] supported_openai_params: Required[Optional[List[str]]] + supports_system_messages: Optional[bool] class GenericStreamingChunk(TypedDict): diff --git a/litellm/utils.py b/litellm/utils.py index 2a0b5689186..8b640b16d37 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -1823,6 +1823,32 @@ def supports_httpx_timeout(custom_llm_provider: str) -> bool: return False +def supports_system_messages(model: str, custom_llm_provider: Optional[str]) -> bool: + """ + Check if the given model supports function calling and return a boolean value. + + Parameters: + model (str): The model name to be checked. + + Returns: + bool: True if the model supports function calling, False otherwise. + + Raises: + Exception: If the given model is not found in model_prices_and_context_window.json. + """ + try: + model_info = litellm.get_model_info( + model=model, custom_llm_provider=custom_llm_provider + ) + if model_info.get("supports_system_messages", False) is True: + return True + return False + except Exception: + raise Exception( + f"Model not in model_prices_and_context_window.json. You passed model={model}, custom_llm_provider={custom_llm_provider}." + ) + + def supports_function_calling(model: str) -> bool: """ Check if the given model supports function calling and return a boolean value. @@ -1838,7 +1864,7 @@ def supports_function_calling(model: str) -> bool: """ if model in litellm.model_cost: model_info = litellm.model_cost[model] - if model_info.get("supports_function_calling", False): + if model_info.get("supports_function_calling", False) is True: return True return False else: @@ -1862,7 +1888,7 @@ def supports_vision(model: str): """ if model in litellm.model_cost: model_info = litellm.model_cost[model] - if model_info.get("supports_vision", False): + if model_info.get("supports_vision", False) is True: return True return False else: @@ -1884,7 +1910,7 @@ def supports_parallel_function_calling(model: str): """ if model in litellm.model_cost: model_info = litellm.model_cost[model] - if model_info.get("supports_parallel_function_calling", False): + if model_info.get("supports_parallel_function_calling", False) is True: return True return False else: @@ -4319,14 +4345,17 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod ) if custom_llm_provider == "huggingface": max_tokens = _get_max_position_embeddings(model_name=model) - return { - "max_tokens": max_tokens, # type: ignore - "input_cost_per_token": 0, - "output_cost_per_token": 0, - "litellm_provider": "huggingface", - "mode": "chat", - "supported_openai_params": supported_openai_params, - } + return ModelInfo( + max_tokens=max_tokens, # type: ignore + max_input_tokens=None, + max_output_tokens=None, + input_cost_per_token=0, + output_cost_per_token=0, + litellm_provider="huggingface", + mode="chat", + supported_openai_params=supported_openai_params, + supports_system_messages=None, + ) else: """ Check if: (in order of specificity) @@ -4361,6 +4390,21 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod pass else: raise Exception + return ModelInfo( + max_tokens=_model_info.get("max_tokens", None), + max_input_tokens=_model_info.get("max_input_tokens", None), + max_output_tokens=_model_info.get("max_output_tokens", None), + input_cost_per_token=_model_info.get("input_cost_per_token", 0), + output_cost_per_token=_model_info.get("output_cost_per_token", 0), + litellm_provider=_model_info.get( + "litellm_provider", custom_llm_provider + ), + mode=_model_info.get("mode"), + supported_openai_params=supported_openai_params, + supports_system_messages=_model_info.get( + "supports_system_messages", None + ), + ) return _model_info elif split_model in litellm.model_cost: _model_info = litellm.model_cost[split_model] @@ -4375,7 +4419,21 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod pass else: raise Exception - return _model_info + return ModelInfo( + max_tokens=_model_info.get("max_tokens", None), + max_input_tokens=_model_info.get("max_input_tokens", None), + max_output_tokens=_model_info.get("max_output_tokens", None), + input_cost_per_token=_model_info.get("input_cost_per_token", 0), + output_cost_per_token=_model_info.get("output_cost_per_token", 0), + litellm_provider=_model_info.get( + "litellm_provider", custom_llm_provider + ), + mode=_model_info.get("mode"), + supported_openai_params=supported_openai_params, + supports_system_messages=_model_info.get( + "supports_system_messages", None + ), + ) else: raise ValueError( "This model isn't mapped yet. Add it here - https://github.com/BerriAI/litellm/blob/main/model_prices_and_context_window.json" From 8f09876486852990eb23c8170e52f0e0547aecc5 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 17:37:15 -0700 Subject: [PATCH 09/14] test(test_completion.py): cleanup test --- litellm/tests/test_completion.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/litellm/tests/test_completion.py b/litellm/tests/test_completion.py index f1ee63564c2..3425d6ac13d 100644 --- a/litellm/tests/test_completion.py +++ b/litellm/tests/test_completion.py @@ -3412,8 +3412,6 @@ def test_completion_gemini(model): # Add any assertions,here to check the response print(response) assert response.choices[0]["index"] == 0 - - assert False except litellm.APIError as e: pass except Exception as e: From 109e0d2f4c6ad0ecec52f29144560b796d80ea0d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 17:40:21 -0700 Subject: [PATCH 10/14] ci(config.yml): fix config.yml --- .circleci/config.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index f68ea2c5cdd..c0c83ade886 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -65,7 +65,7 @@ jobs: pip install "pydantic==2.7.1" pip install "diskcache==5.6.1" pip install "Pillow==10.3.0" - pip install "Pillow==3.2.3" + pip install "ijson==3.2.3" - save_cache: paths: - ./venv From 1d2f1b0bb9826def9dc386b66f488a8e58f9bc92 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 17:41:56 -0700 Subject: [PATCH 11/14] ci(config.yml): update config.yml --- .circleci/config.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index c0c83ade886..46234c3db6c 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -127,6 +127,7 @@ jobs: pip install jinja2 pip install tokenizers pip install openai + pip install ijson - run: name: Run tests command: | From 63a6ae9d5581ee4efbbe9fbd3329297c964e3440 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 18:48:34 -0700 Subject: [PATCH 12/14] build(pyproject.toml): fix ijson --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 6d014101855..9a4e14eb7c6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,7 +27,7 @@ jinja2 = "^3.1.2" aiohttp = "*" requests = "^2.31.0" pydantic = "^2.0.0" -isjon = "*" +ijson = "*" uvicorn = {version = "^0.22.0", optional = true} gunicorn = {version = "^22.0.0", optional = true} From b1775068e8fc44b75f40b1830ad9611d487eb909 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 19:09:29 -0700 Subject: [PATCH 13/14] build(config.yml): fix config --- .circleci/config.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index 46234c3db6c..d070c55dca7 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -182,6 +182,7 @@ jobs: pip install numpydoc pip install prisma pip install fastapi + pip install ijson pip install "httpx==0.24.1" pip install "gunicorn==21.2.0" pip install "anyio==3.7.1" From fca2ffb4801fd41bb44bcc2871697f4169c7b01e Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 17 Jun 2024 19:15:02 -0700 Subject: [PATCH 14/14] fix(utils.py): return cost above 128k from get_model_info --- litellm/utils.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/litellm/utils.py b/litellm/utils.py index 8b640b16d37..79269a4b83f 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -4376,7 +4376,27 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod pass else: raise Exception - return _model_info + return ModelInfo( + max_tokens=_model_info.get("max_tokens", None), + max_input_tokens=_model_info.get("max_input_tokens", None), + max_output_tokens=_model_info.get("max_output_tokens", None), + input_cost_per_token=_model_info.get("input_cost_per_token", 0), + input_cost_per_token_above_128k_tokens=_model_info.get( + "input_cost_per_token_above_128k_tokens", None + ), + output_cost_per_token=_model_info.get("output_cost_per_token", 0), + output_cost_per_token_above_128k_tokens=_model_info.get( + "output_cost_per_token_above_128k_tokens", None + ), + litellm_provider=_model_info.get( + "litellm_provider", custom_llm_provider + ), + mode=_model_info.get("mode"), + supported_openai_params=supported_openai_params, + supports_system_messages=_model_info.get( + "supports_system_messages", None + ), + ) elif model in litellm.model_cost: _model_info = litellm.model_cost[model] _model_info["supported_openai_params"] = supported_openai_params @@ -4395,7 +4415,13 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod max_input_tokens=_model_info.get("max_input_tokens", None), max_output_tokens=_model_info.get("max_output_tokens", None), input_cost_per_token=_model_info.get("input_cost_per_token", 0), + input_cost_per_token_above_128k_tokens=_model_info.get( + "input_cost_per_token_above_128k_tokens", None + ), output_cost_per_token=_model_info.get("output_cost_per_token", 0), + output_cost_per_token_above_128k_tokens=_model_info.get( + "output_cost_per_token_above_128k_tokens", None + ), litellm_provider=_model_info.get( "litellm_provider", custom_llm_provider ), @@ -4405,7 +4431,6 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod "supports_system_messages", None ), ) - return _model_info elif split_model in litellm.model_cost: _model_info = litellm.model_cost[split_model] _model_info["supported_openai_params"] = supported_openai_params @@ -4424,7 +4449,13 @@ def get_model_info(model: str, custom_llm_provider: Optional[str] = None) -> Mod max_input_tokens=_model_info.get("max_input_tokens", None), max_output_tokens=_model_info.get("max_output_tokens", None), input_cost_per_token=_model_info.get("input_cost_per_token", 0), + input_cost_per_token_above_128k_tokens=_model_info.get( + "input_cost_per_token_above_128k_tokens", None + ), output_cost_per_token=_model_info.get("output_cost_per_token", 0), + output_cost_per_token_above_128k_tokens=_model_info.get( + "output_cost_per_token_above_128k_tokens", None + ), litellm_provider=_model_info.get( "litellm_provider", custom_llm_provider ),