scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-11 22:51:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 123

fix(main.py): only return mock model response stream if stream is true

Browse source
This commit is contained in:
Krrish Dholakia 2025-10-25 12:58:53 -07:00
parent c9bc6d593c
commit 42eb86245a
1 changed files with 130 additions and 174 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

304
litellm/main.py
View file

@ -87,7 +87,6 @@ from litellm.llms.base_llm.base_model_iterator import (
convert_model_response_to_streaming,
)
from litellm.llms.bedrock.common_utils import BedrockModelInfo
from litellm.llms.cohere.common_utils import CohereModelInfo
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler
from litellm.llms.vertex_ai.common_utils import (
VertexAIModelRoute,
@ -158,7 +157,6 @@ from .llms.bedrock.chat import BedrockConverseLLM, BedrockLLM
from .llms.bedrock.embed.embedding import BedrockEmbedding
from .llms.bedrock.image.image_handler import BedrockImageGeneration
from .llms.bytez.chat.transformation import BytezChatConfig
from .llms.clarifai.chat.transformation import ClarifaiConfig
from .llms.codestral.completion.handler import CodestralTextCompletion
from .llms.cohere.embed import handler as cohere_embed
from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler
@ -388,7 +386,6 @@ async def acompletion(
Literal["none", "minimal", "low", "medium", "high", "default"]
] = None,
safety_identifier: Optional[str] = None,
service_tier: Optional[str] = None,
# set api_base, api_version, api_key
base_url: Optional[str] = None,
api_version: Optional[str] = None,
@ -538,7 +535,6 @@ async def acompletion(
"model_list": model_list,
"reasoning_effort": reasoning_effort,
"safety_identifier": safety_identifier,
"service_tier": service_tier,
"extra_headers": extra_headers,
"acompletion": True, # assuming this is a required parameter
"thinking": thinking,
@ -824,9 +820,10 @@ def mock_completion(
# convert to ModelResponseStream
mock_response = convert_model_response_to_streaming(mock_response) # type: ignore
model_response = ModelResponseStream()
model_response: Union[ModelResponse, ModelResponseStream] = ModelResponse()
if stream is True:
model_response = ModelResponseStream()
# don't try to access stream object,
if kwargs.get("acompletion", False) is True:
return CustomStreamWrapper(
@ -968,7 +965,6 @@ def completion( # type: ignore # noqa: PLR0915
deployment_id=None,
extra_headers: Optional[dict] = None,
safety_identifier: Optional[str] = None,
service_tier: Optional[str] = None,
# soon to be deprecated params by OpenAI
functions: Optional[List] = None,
function_call: Optional[str] = None,
@ -1311,7 +1307,6 @@ def completion( # type: ignore # noqa: PLR0915
"thinking": thinking,
"web_search_options": web_search_options,
"safety_identifier": safety_identifier,
"service_tier": service_tier,
"allowed_openai_params": kwargs.get("allowed_openai_params"),
}
optional_params = get_optional_params(
@ -2047,7 +2042,6 @@ def completion( # type: ignore # noqa: PLR0915
or custom_llm_provider == "together_ai"
or custom_llm_provider == "nebius"
or custom_llm_provider == "wandb"
or custom_llm_provider == "clarifai"
or custom_llm_provider in litellm.openai_compatible_providers
or "ft:gpt-3.5-turbo" in model # finetune gpt-3.5-turbo
): # allow user to make an openai call with a custom base
@ -2242,7 +2236,40 @@ def completion( # type: ignore # noqa: PLR0915
or custom_llm_provider == "clarifai"
or model in litellm.clarifai_models
):
pass # Deprecated - handled in the openai compatible provider section above
clarifai_key = None
clarifai_key = (
api_key
or litellm.clarifai_key
or litellm.api_key
or get_secret("CLARIFAI_API_KEY")
or get_secret("CLARIFAI_API_TOKEN")
)
api_base = (
api_base
or litellm.api_base
or get_secret("CLARIFAI_API_BASE")
or "https://api.clarifai.com/v2"
)
api_base = litellm.ClarifaiConfig()._convert_model_to_url(model, api_base)
response = base_llm_http_handler.completion(
model=model,
stream=stream,
fake_stream=True, # clarifai does not support streaming, we fake it
messages=messages,
acompletion=acompletion,
api_base=api_base,
model_response=model_response,
optional_params=optional_params,
litellm_params=litellm_params,
shared_session=shared_session,
custom_llm_provider="clarifai",
timeout=timeout,
headers=headers,
encoding=encoding,
api_key=clarifai_key,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
)
elif custom_llm_provider == "anthropic_text":
api_key = (
api_key
@ -2435,7 +2462,7 @@ def completion( # type: ignore # noqa: PLR0915
)
return response
response = model_response
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
elif custom_llm_provider == "cohere_chat" or custom_llm_provider == "cohere":
cohere_key = (
api_key
or litellm.cohere_key
@ -2444,26 +2471,12 @@ def completion( # type: ignore # noqa: PLR0915
or litellm.api_key
)
cohere_route = CohereModelInfo.get_cohere_route(model)
verbose_logger.debug(f"Cohere route: {cohere_route}")
# Set API base based on route
if cohere_route == "v2":
api_base = (
api_base
or litellm.api_base
or get_secret_str("COHERE_API_BASE")
or "https://api.cohere.com/v2/chat"
)
# Remove v2/ prefix from model name for the actual API call
if "v2/" in model:
model = model.replace("v2/", "")
else:
api_base = (
api_base
or litellm.api_base
or get_secret_str("COHERE_API_BASE")
or "https://api.cohere.ai/v1/chat"
)
api_base = (
api_base
or litellm.api_base
or get_secret_str("COHERE_API_BASE")
or "https://api.cohere.ai/v1/chat"
)
headers = headers or litellm.headers or {}
if headers is None:
@ -2472,8 +2485,6 @@ def completion( # type: ignore # noqa: PLR0915
if extra_headers is not None:
headers.update(extra_headers)
verbose_logger.debug(f"Model: {model}, API Base: {api_base}")
verbose_logger.debug(f"Provider Config: {provider_config}")
response = base_llm_http_handler.completion(
model=model,
stream=stream,
@ -2489,7 +2500,6 @@ def completion( # type: ignore # noqa: PLR0915
headers=headers,
encoding=encoding,
api_key=cohere_key,
provider_config=provider_config,
logging_obj=logging, # model call logging done inside the class as we make need to modify I/O to fit aleph alpha's requirements
)
elif custom_llm_provider == "maritalk":
@ -4002,7 +4012,6 @@ def embedding( # noqa: PLR0915
"""
azure = kwargs.get("azure", None)
client = kwargs.pop("client", None)
shared_session = kwargs.get("shared_session", None)
max_retries = kwargs.get("max_retries", None)
litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore
mock_response: Optional[List[float]] = kwargs.get("mock_response", None) # type: ignore
@ -4192,7 +4201,6 @@ def embedding( # noqa: PLR0915
client=client,
aembedding=aembedding,
max_retries=max_retries,
shared_session=shared_session,
)
elif custom_llm_provider == "databricks":
api_base = api_base or litellm.api_base or get_secret("DATABRICKS_API_BASE") # type: ignore
@ -4760,33 +4768,6 @@ def embedding( # noqa: PLR0915
aembedding=aembedding,
litellm_params={},
)
elif custom_llm_provider == "cometapi":
api_key = (
api_key
or litellm.cometapi_key
or get_secret_str("COMETAPI_KEY")
or litellm.api_key
)
api_base = (
api_base
or litellm.api_base
or get_secret_str("COMETAPI_API_BASE")
or "https://api.cometapi.com/v1"
)
response = base_llm_http_handler.embedding(
model=model,
input=input,
custom_llm_provider=custom_llm_provider,
api_base=api_base,
api_key=api_key,
logging_obj=logging,
timeout=timeout,
model_response=EmbeddingResponse(),
optional_params=optional_params,
client=client,
aembedding=aembedding,
litellm_params={},
)
elif custom_llm_provider in litellm._custom_providers:
custom_handler: Optional[CustomLLM] = None
for item in litellm.custom_provider_map:
@ -5722,35 +5703,17 @@ def speech( # noqa: PLR0915
optional_params["speed"] = speed # type: ignore
if instructions is not None:
optional_params["instructions"] = instructions
if timeout is None:
timeout = litellm.request_timeout
if max_retries is None:
max_retries = litellm.num_retries or openai.DEFAULT_MAX_RETRIES
litellm_params_dict = get_litellm_params(**kwargs)
# Get provider-specific text-to-speech config and map parameters
text_to_speech_provider_config = ProviderConfigManager.get_provider_text_to_speech_config(
model=model,
provider=litellm.LlmProviders(custom_llm_provider),
)
# Map OpenAI params to provider-specific params if config exists
if text_to_speech_provider_config is not None:
voice, optional_params = text_to_speech_provider_config.map_openai_params(
model=model,
optional_params=optional_params,
voice=voice,
drop_params=False,
kwargs=kwargs,
)
logging_obj: Logging = cast(Logging, kwargs.get("litellm_logging_obj"))
logging_obj.update_environment_variables(
model=model,
user=user,
optional_params=optional_params,
optional_params={},
litellm_params={
"litellm_call_id": litellm_call_id,
"proxy_server_request": proxy_server_request,
@ -5819,85 +5782,52 @@ def speech( # noqa: PLR0915
aspeech=aspeech,
)
elif custom_llm_provider == "azure":
# Check if this is Azure Speech Service (Cognitive Services TTS)
if model.startswith("speech/"):
from litellm.llms.azure.text_to_speech.transformation import (
AzureAVATextToSpeechConfig,
)
# Azure AVA (Cognitive Services) Text-to-Speech
if text_to_speech_provider_config is None:
raise litellm.BadRequestError(
message="Azure Speech Service configuration not found",
model=model,
llm_provider=custom_llm_provider,
)
# Cast to specific Azure config type to access dispatch method
azure_config = cast(AzureAVATextToSpeechConfig, text_to_speech_provider_config)
response = azure_config.dispatch_text_to_speech( # type: ignore
# azure configs
if voice is None or not (isinstance(voice, str)):
raise litellm.BadRequestError(
message="'voice' is required to be passed as a string for Azure TTS",
model=model,
input=input,
voice=voice,
optional_params=optional_params,
litellm_params_dict=litellm_params_dict,
logging_obj=logging_obj,
timeout=timeout,
extra_headers=extra_headers,
base_llm_http_handler=base_llm_http_handler,
aspeech=aspeech or False,
api_base=api_base,
api_key=api_key,
**kwargs,
llm_provider=custom_llm_provider,
)
else:
# Azure OpenAI TTS
if voice is None or not (isinstance(voice, str)):
raise litellm.BadRequestError(
message="'voice' is required to be passed as a string for Azure TTS",
model=model,
llm_provider=custom_llm_provider,
)
api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_base = api_base or litellm.api_base or get_secret("AZURE_API_BASE") # type: ignore
api_version = api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_version = api_version or litellm.api_version or get_secret("AZURE_API_VERSION") # type: ignore
api_key = (
api_key
or litellm.api_key
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
api_key = (
api_key
or litellm.api_key
or litellm.azure_key
or get_secret("AZURE_OPENAI_API_KEY")
or get_secret("AZURE_API_KEY")
) # type: ignore
azure_ad_token: Optional[str] = optional_params.get("extra_body", {}).pop( # type: ignore
"azure_ad_token", None
) or get_secret(
"AZURE_AD_TOKEN"
)
azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None)
azure_ad_token: Optional[str] = optional_params.get("extra_body", {}).pop( # type: ignore
"azure_ad_token", None
) or get_secret(
"AZURE_AD_TOKEN"
)
azure_ad_token_provider = kwargs.get("azure_ad_token_provider", None)
if extra_headers:
optional_params["extra_headers"] = extra_headers
if extra_headers:
optional_params["extra_headers"] = extra_headers
response = azure_chat_completions.audio_speech(
model=model,
input=input,
voice=voice,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
api_version=api_version,
azure_ad_token=azure_ad_token,
azure_ad_token_provider=azure_ad_token_provider,
organization=organization,
max_retries=max_retries,
timeout=timeout,
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
litellm_params=litellm_params_dict,
)
response = azure_chat_completions.audio_speech(
model=model,
input=input,
voice=voice,
optional_params=optional_params,
api_key=api_key,
api_base=api_base,
api_version=api_version,
azure_ad_token=azure_ad_token,
azure_ad_token_provider=azure_ad_token_provider,
organization=organization,
max_retries=max_retries,
timeout=timeout,
client=client, # pass AsyncOpenAI, OpenAI client
aspeech=aspeech,
litellm_params=litellm_params_dict,
)
elif custom_llm_provider == "vertex_ai" or custom_llm_provider == "vertex_ai_beta":
generic_optional_params = GenericLiteLLMParams(**kwargs)
@ -5993,8 +5923,6 @@ async def ahealth_check(
"batch",
"rerank",
"realtime",
"responses",
"ocr",
]
] = "chat",
prompt: Optional[str] = None,
@ -6057,13 +5985,53 @@ async def ahealth_check(
litellm_logging_obj=litellm_logging_obj,
)
mode_handlers = HealthCheckHelpers.get_mode_handlers(
model=model,
custom_llm_provider=custom_llm_provider,
model_params=model_params,
prompt=prompt,
input=input,
)
mode_handlers = {
"chat": lambda: litellm.acompletion(
**model_params,
),
"completion": lambda: litellm.atext_completion(
**_filter_model_params(model_params),
prompt=prompt or "test",
),
"embedding": lambda: litellm.aembedding(
**_filter_model_params(model_params),
input=input or ["test"],
),
"audio_speech": lambda: litellm.aspeech(
**{
**_filter_model_params(model_params),
**(
{"voice": "alloy"}
if "voice" not in _filter_model_params(model_params)
else {}
),
},
input=prompt or "test",
),
"audio_transcription": lambda: litellm.atranscription(
**_filter_model_params(model_params),
file=get_audio_file_for_health_check(),
),
"image_generation": lambda: litellm.aimage_generation(
**_filter_model_params(model_params),
prompt=prompt,
),
"rerank": lambda: litellm.arerank(
**_filter_model_params(model_params),
query=prompt or "",
documents=["my sample text"],
),
"realtime": lambda: _realtime_health_check(
model=model,
custom_llm_provider=custom_llm_provider,
api_base=model_params.get("api_base", None),
api_key=model_params.get("api_key", None),
api_version=model_params.get("api_version", None),
),
"batch": lambda: litellm.alist_batches(
**_filter_model_params(model_params),
),
}
if mode in mode_handlers:
_response = await mode_handlers[mode]()
@ -6298,18 +6266,6 @@ def stream_chunk_builder( # noqa: PLR0915
processor.get_combined_reasoning_content(reasoning_chunks)
)
annotation_chunks = [
chunk
for chunk in chunks
if len(chunk["choices"]) > 0
and "annotations" in chunk["choices"][0]["delta"]
and chunk["choices"][0]["delta"]["annotations"] is not None
]
if len(annotation_chunks) > 0:
annotations = annotation_chunks[0]["choices"][0]["delta"]["annotations"]
response["choices"][0]["message"]["annotations"] = annotations
audio_chunks = [
chunk
for chunk in chunks