mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
Merge pull request #41662 from BerriAI/litellm_remove_retired_provider_streaming_handlers
chore(streaming): remove retired ai21/maritalk/baseten/azure raw-bytes handlers and dead palm completion code
This commit is contained in:
Mateo Wang
GitHub
commit
151a92a230
4 changed files with 1 additions and 316 deletions
|
|
@ -113,14 +113,6 @@ class _PredibaseStreamData(TypedDict):
|
|||
error: str | None
|
||||
|
||||
|
||||
class _Ai21StreamData(TypedDict):
|
||||
completions: Sequence[Mapping[str, Mapping[str, str]]]
|
||||
|
||||
|
||||
class _MaritalkStreamData(TypedDict):
|
||||
answer: str
|
||||
|
||||
|
||||
class _NlpCloudStreamData(TypedDict):
|
||||
generated_text: str
|
||||
|
||||
|
|
@ -129,25 +121,6 @@ class _AlephAlphaStreamData(TypedDict):
|
|||
completions: Sequence[Mapping[str, str]]
|
||||
|
||||
|
||||
class _AzureStreamChoice(TypedDict):
|
||||
delta: Mapping[str, str] | None
|
||||
finish_reason: str | None
|
||||
|
||||
|
||||
class _AzureStreamData(TypedDict):
|
||||
choices: Sequence[_AzureStreamChoice]
|
||||
|
||||
|
||||
class _BasetenModelOutput(TypedDict):
|
||||
data: NotRequired[Sequence[str]]
|
||||
|
||||
|
||||
class _BasetenStreamData(TypedDict):
|
||||
token: NotRequired[Mapping[str, str]]
|
||||
model_output: NotRequired["_BasetenModelOutput | str"]
|
||||
completion: NotRequired[object]
|
||||
|
||||
|
||||
class _DeltaDumpDict(TypedDict):
|
||||
role: NotRequired[str | None]
|
||||
tool_calls: NotRequired[Sequence[Mapping[str, object]]]
|
||||
|
|
@ -572,36 +545,6 @@ class CustomStreamWrapper:
|
|||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def handle_ai21_chunk(self, chunk): # fake streaming
|
||||
chunk = chunk.decode("utf-8")
|
||||
data_json: Final[_Ai21StreamData] = json.loads(chunk)
|
||||
try:
|
||||
text: Final = data_json["completions"][0]["data"]["text"]
|
||||
is_finished: Final = True
|
||||
finish_reason: Final = "stop"
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
except Exception:
|
||||
raise ValueError(f"Unable to parse response. Original response: {chunk}")
|
||||
|
||||
def handle_maritalk_chunk(self, chunk): # fake streaming
|
||||
chunk = chunk.decode("utf-8")
|
||||
data_json: Final[_MaritalkStreamData] = json.loads(chunk)
|
||||
try:
|
||||
text: Final = data_json["answer"]
|
||||
is_finished: Final = True
|
||||
finish_reason: Final = "stop"
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
except Exception:
|
||||
raise ValueError(f"Unable to parse response. Original response: {chunk}")
|
||||
|
||||
def handle_nlp_cloud_chunk(self, chunk):
|
||||
text = ""
|
||||
is_finished = False
|
||||
|
|
@ -640,46 +583,6 @@ class CustomStreamWrapper:
|
|||
except Exception:
|
||||
raise ValueError(f"Unable to parse response. Original response: {chunk}")
|
||||
|
||||
def handle_azure_chunk(self, chunk):
|
||||
is_finished = False
|
||||
finish_reason = ""
|
||||
text = ""
|
||||
print_verbose(f"chunk: {chunk}")
|
||||
if "data: [DONE]" in chunk:
|
||||
text = ""
|
||||
is_finished = True
|
||||
finish_reason = "stop"
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
elif chunk.startswith("data:"):
|
||||
data_json: Final[_AzureStreamData] = json.loads(chunk[5:]) # chunk.startswith("data:"):
|
||||
try:
|
||||
if len(data_json["choices"]) > 0:
|
||||
delta: Final = data_json["choices"][0]["delta"]
|
||||
text = "" if delta is None else delta.get("content", "")
|
||||
if data_json["choices"][0].get("finish_reason", None):
|
||||
is_finished = True
|
||||
finish_reason = data_json["choices"][0]["finish_reason"]
|
||||
print_verbose(f"text: {text}; is_finished: {is_finished}; finish_reason: {finish_reason}")
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
except Exception:
|
||||
raise ValueError(f"Unable to parse response. Original response: {chunk}")
|
||||
elif "error" in chunk:
|
||||
raise ValueError(f"Unable to parse response. Original response: {chunk}")
|
||||
else:
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
|
||||
def handle_replicate_chunk(self, chunk):
|
||||
try:
|
||||
text = ""
|
||||
|
|
@ -782,38 +685,6 @@ class CustomStreamWrapper:
|
|||
except Exception as e:
|
||||
raise e
|
||||
|
||||
def handle_baseten_chunk(self, chunk) -> str:
|
||||
try:
|
||||
chunk = chunk.decode("utf-8")
|
||||
if len(chunk) > 0:
|
||||
if chunk.startswith("data:"):
|
||||
data_json: _BasetenStreamData = json.loads(chunk[5:])
|
||||
if "token" in data_json and "text" in data_json["token"]:
|
||||
return data_json["token"]["text"]
|
||||
else:
|
||||
return ""
|
||||
data_json = json.loads(chunk)
|
||||
if "model_output" in data_json:
|
||||
if (
|
||||
isinstance(data_json["model_output"], dict)
|
||||
and "data" in data_json["model_output"]
|
||||
and isinstance(data_json["model_output"]["data"], list)
|
||||
):
|
||||
return data_json["model_output"]["data"][0]
|
||||
elif isinstance(data_json["model_output"], str):
|
||||
return data_json["model_output"]
|
||||
elif "completion" in data_json and isinstance(data_json["completion"], str):
|
||||
return data_json["completion"]
|
||||
else:
|
||||
raise ValueError(f"Unable to parse response. Original response: {chunk}")
|
||||
else:
|
||||
return ""
|
||||
else:
|
||||
return ""
|
||||
except Exception as e:
|
||||
verbose_logger.exception("litellm.CustomStreamWrapper.handle_baseten_chunk(): Exception occured - %s", e)
|
||||
return ""
|
||||
|
||||
def handle_triton_stream(self, chunk):
|
||||
try:
|
||||
if isinstance(chunk, dict):
|
||||
|
|
@ -1305,18 +1176,6 @@ class CustomStreamWrapper:
|
|||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "baseten": # baseten doesn't provide streaming
|
||||
completion_obj["content"] = self.handle_baseten_chunk(chunk)
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "ai21": # ai21 doesn't provide streaming
|
||||
response_obj = self.handle_ai21_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "maritalk":
|
||||
response_obj = self.handle_maritalk_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
elif self.custom_llm_provider and self.custom_llm_provider == "vllm":
|
||||
completion_obj["content"] = chunk[0].outputs[0].text
|
||||
elif (
|
||||
|
|
@ -1410,19 +1269,6 @@ class CustomStreamWrapper:
|
|||
new_chunk = stream[:chunk_size]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = stream[chunk_size:]
|
||||
elif self.custom_llm_provider == "palm":
|
||||
# fake streaming
|
||||
response_obj = {}
|
||||
if self.completion_stream is None or len(self.completion_stream) == 0:
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
else:
|
||||
self.received_finish_reason = "stop"
|
||||
chunk_size = 30
|
||||
stream = cast(Any, self.completion_stream)
|
||||
new_chunk = stream[:chunk_size]
|
||||
completion_obj["content"] = new_chunk
|
||||
self.completion_stream = stream[chunk_size:]
|
||||
elif self.custom_llm_provider == "triton":
|
||||
response_obj = self.handle_triton_stream(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
|
|
|
|||
|
|
@ -1,27 +1,6 @@
|
|||
import copy
|
||||
import time
|
||||
import traceback
|
||||
import types
|
||||
from collections.abc import Callable
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm.utils import Choices, Message, ModelResponse, Usage
|
||||
|
||||
|
||||
class PalmError(Exception):
|
||||
def __init__(self, status_code, message):
|
||||
self.status_code = status_code
|
||||
self.message = message
|
||||
self.request = httpx.Request(
|
||||
method="POST",
|
||||
url="https://developers.generativeai.google/api/python/google/generativeai/chat",
|
||||
)
|
||||
self.response = httpx.Response(status_code=status_code, request=self.request)
|
||||
super().__init__(self.message) # Call the base class constructor with the parameters it needs
|
||||
|
||||
|
||||
class PalmConfig:
|
||||
"""
|
||||
|
|
@ -84,111 +63,3 @@ class PalmConfig:
|
|||
)
|
||||
and v is not None
|
||||
}
|
||||
|
||||
|
||||
def completion(
|
||||
model: str,
|
||||
messages: list,
|
||||
model_response: ModelResponse,
|
||||
print_verbose: Callable,
|
||||
api_key,
|
||||
encoding,
|
||||
logging_obj,
|
||||
optional_params: dict,
|
||||
litellm_params=None,
|
||||
logger_fn=None,
|
||||
):
|
||||
try:
|
||||
import google.generativeai as palm
|
||||
except Exception:
|
||||
raise Exception("Importing google.generativeai failed, please run 'pip install -q google-generativeai")
|
||||
palm.configure(api_key=api_key)
|
||||
|
||||
model = model
|
||||
|
||||
## Load Config
|
||||
inference_params: Final = copy.deepcopy(optional_params)
|
||||
inference_params.pop(
|
||||
"stream", None
|
||||
) # palm does not support streaming, so we handle this by fake streaming in main.py
|
||||
config: Final = litellm.PalmConfig.get_config()
|
||||
for k, v in config.items():
|
||||
if (
|
||||
k not in inference_params
|
||||
): # completion(top_k=3) > palm_config(top_k=3) <- allows for dynamic variables to be passed in
|
||||
inference_params[k] = v
|
||||
|
||||
prompt = ""
|
||||
for message in messages:
|
||||
if "role" in message:
|
||||
if message["role"] == "user":
|
||||
prompt += f"{message['content']}"
|
||||
else:
|
||||
prompt += f"{message['content']}"
|
||||
else:
|
||||
prompt += f"{message['content']}"
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=prompt,
|
||||
api_key="",
|
||||
additional_args={"complete_input_dict": {"inference_params": inference_params}},
|
||||
)
|
||||
## COMPLETION CALL
|
||||
try:
|
||||
response: Final = palm.generate_text(prompt=prompt, **inference_params)
|
||||
except Exception as e:
|
||||
raise PalmError(
|
||||
message=str(e),
|
||||
status_code=500,
|
||||
)
|
||||
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=prompt,
|
||||
api_key="",
|
||||
original_response=response,
|
||||
additional_args={"complete_input_dict": {}},
|
||||
)
|
||||
print_verbose(f"raw model_response: {response}")
|
||||
## RESPONSE OBJECT
|
||||
completion_response = response
|
||||
try:
|
||||
choices_list: Final = []
|
||||
for idx, item in enumerate(completion_response.candidates):
|
||||
if len(item["output"]) > 0:
|
||||
message_obj = Message(content=item["output"])
|
||||
else:
|
||||
message_obj = Message(content=None)
|
||||
choice_obj = Choices(index=idx + 1, message=message_obj)
|
||||
choices_list.append(choice_obj)
|
||||
model_response.choices = choices_list
|
||||
except Exception:
|
||||
raise PalmError(message=traceback.format_exc(), status_code=response.status_code)
|
||||
|
||||
try:
|
||||
completion_response = model_response["choices"][0]["message"].get("content")
|
||||
except Exception:
|
||||
raise PalmError(
|
||||
status_code=400,
|
||||
message=f"No response received. Original response - {response}",
|
||||
)
|
||||
|
||||
## CALCULATING USAGE - baseten charges on time, not tokens - have some mapping of cost here.
|
||||
prompt_tokens: Final = len(encoding.encode(prompt))
|
||||
completion_tokens: Final = len(encoding.encode(model_response["choices"][0]["message"].get("content", "")))
|
||||
|
||||
model_response.created = int(time.time())
|
||||
model_response.model = "palm/" + model
|
||||
usage: Final = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
)
|
||||
setattr(model_response, "usage", usage)
|
||||
return model_response
|
||||
|
||||
|
||||
def embedding():
|
||||
# logic for parsing in - calling - parsing out model embedding calls
|
||||
pass
|
||||
|
|
|
|||
|
|
@ -206,7 +206,7 @@ from .llms.custom_httpx.aiohttp_handler import BaseLLMAIOHTTPHandler
|
|||
from .llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler
|
||||
from .llms.custom_llm import CustomLLM, custom_chat_llm_router
|
||||
from .llms.databricks.embed.handler import DatabricksEmbeddingHandler
|
||||
from .llms.deprecated_providers import aleph_alpha, palm
|
||||
from .llms.deprecated_providers import aleph_alpha
|
||||
from .llms.gdc.chat.transformation import GDCGeminiConfig
|
||||
from .llms.gemini.common_utils import get_api_key_from_env
|
||||
from .llms.groq.chat.handler import GroqChatCompletion
|
||||
|
|
|
|||
|
|
@ -2589,22 +2589,6 @@ def test_dispatch_petals_empty_stream_after_finish_raises(
|
|||
_run_dispatch(initialized_custom_stream_wrapper, chunk=None)
|
||||
|
||||
|
||||
def test_dispatch_palm_slices_completion_stream(
|
||||
initialized_custom_stream_wrapper: CustomStreamWrapper,
|
||||
):
|
||||
"""palm uses the same fake-streaming slice strategy as petals."""
|
||||
initialized_custom_stream_wrapper.custom_llm_provider = "palm"
|
||||
initialized_custom_stream_wrapper.completion_stream = "B" * 40
|
||||
|
||||
result, _, completion_obj = _run_dispatch(
|
||||
initialized_custom_stream_wrapper, chunk=None
|
||||
)
|
||||
|
||||
assert isinstance(result, _ProviderChunkParsed)
|
||||
assert completion_obj["content"] == "B" * 30
|
||||
assert initialized_custom_stream_wrapper.completion_stream == "B" * 10
|
||||
|
||||
|
||||
def test_dispatch_cached_response_extracts_delta(
|
||||
initialized_custom_stream_wrapper: CustomStreamWrapper,
|
||||
):
|
||||
|
|
@ -2844,22 +2828,6 @@ def test_dispatch_triton_stream(
|
|||
assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
|
||||
|
||||
|
||||
def test_dispatch_ai21_decodes_completion(
|
||||
initialized_custom_stream_wrapper: CustomStreamWrapper,
|
||||
):
|
||||
"""ai21 does fake streaming over a single byte-encoded JSON completion."""
|
||||
initialized_custom_stream_wrapper.custom_llm_provider = "ai21"
|
||||
chunk = json.dumps({"completions": [{"data": {"text": "ai21 text"}}]}).encode(
|
||||
"utf-8"
|
||||
)
|
||||
|
||||
result, _, completion_obj = _run_dispatch(initialized_custom_stream_wrapper, chunk)
|
||||
|
||||
assert isinstance(result, _ProviderChunkParsed)
|
||||
assert completion_obj["content"] == "ai21 text"
|
||||
assert initialized_custom_stream_wrapper.received_finish_reason == "stop"
|
||||
|
||||
|
||||
def test_dispatch_text_completion_openai_with_usage(
|
||||
initialized_custom_stream_wrapper: CustomStreamWrapper,
|
||||
):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue