mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-17 08:18:09 +00:00
Merge pull request #4536 from BerriAI/litellm_anthropic_tool_calling_streaming_fix
*real* Anthropic tool calling + streaming support
This commit is contained in:
Krish Dholakia
+323
-153
@@ -18,7 +18,20 @@ from litellm.llms.custom_httpx.http_handler import (
|
||||
_get_async_httpx_client,
|
||||
_get_httpx_client,
|
||||
)
|
||||
from litellm.types.llms.anthropic import AnthropicMessagesToolChoice
|
||||
from litellm.types.llms.anthropic import (
|
||||
AnthropicMessagesToolChoice,
|
||||
ContentBlockDelta,
|
||||
ContentBlockStart,
|
||||
MessageBlockDelta,
|
||||
MessageStartBlock,
|
||||
)
|
||||
from litellm.types.llms.openai import (
|
||||
ChatCompletionResponseMessage,
|
||||
ChatCompletionToolCallChunk,
|
||||
ChatCompletionToolCallFunctionChunk,
|
||||
ChatCompletionUsageBlock,
|
||||
)
|
||||
from litellm.types.utils import GenericStreamingChunk
|
||||
from litellm.utils import CustomStreamWrapper, ModelResponse, Usage
|
||||
|
||||
from .base import BaseLLM
|
||||
@@ -198,7 +211,9 @@ async def make_call(
|
||||
status_code=response.status_code, message=await response.aread()
|
||||
)
|
||||
|
||||
completion_stream = response.aiter_lines()
|
||||
completion_stream = ModelResponseIterator(
|
||||
streaming_response=response.aiter_lines(), sync_stream=False
|
||||
)
|
||||
|
||||
# LOGGING
|
||||
logging_obj.post_call(
|
||||
@@ -215,120 +230,120 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
|
||||
def process_streaming_response(
|
||||
self,
|
||||
model: str,
|
||||
response: Union[requests.Response, httpx.Response],
|
||||
model_response: ModelResponse,
|
||||
stream: bool,
|
||||
logging_obj: litellm.litellm_core_utils.litellm_logging.Logging,
|
||||
optional_params: dict,
|
||||
api_key: str,
|
||||
data: Union[dict, str],
|
||||
messages: List,
|
||||
print_verbose,
|
||||
encoding,
|
||||
) -> CustomStreamWrapper:
|
||||
"""
|
||||
Return stream object for tool-calling + streaming
|
||||
"""
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
original_response=response.text,
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
print_verbose(f"raw model_response: {response.text}")
|
||||
## RESPONSE OBJECT
|
||||
try:
|
||||
completion_response = response.json()
|
||||
except:
|
||||
raise AnthropicError(
|
||||
message=response.text, status_code=response.status_code
|
||||
)
|
||||
text_content = ""
|
||||
tool_calls = []
|
||||
for content in completion_response["content"]:
|
||||
if content["type"] == "text":
|
||||
text_content += content["text"]
|
||||
## TOOL CALLING
|
||||
elif content["type"] == "tool_use":
|
||||
tool_calls.append(
|
||||
{
|
||||
"id": content["id"],
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": content["name"],
|
||||
"arguments": json.dumps(content["input"]),
|
||||
},
|
||||
}
|
||||
)
|
||||
if "error" in completion_response:
|
||||
raise AnthropicError(
|
||||
message=str(completion_response["error"]),
|
||||
status_code=response.status_code,
|
||||
)
|
||||
_message = litellm.Message(
|
||||
tool_calls=tool_calls,
|
||||
content=text_content or None,
|
||||
)
|
||||
model_response.choices[0].message = _message # type: ignore
|
||||
model_response._hidden_params["original_response"] = completion_response[
|
||||
"content"
|
||||
] # allow user to access raw anthropic tool calling response
|
||||
# def process_streaming_response(
|
||||
# self,
|
||||
# model: str,
|
||||
# response: Union[requests.Response, httpx.Response],
|
||||
# model_response: ModelResponse,
|
||||
# stream: bool,
|
||||
# logging_obj: litellm.litellm_core_utils.litellm_logging.Logging,
|
||||
# optional_params: dict,
|
||||
# api_key: str,
|
||||
# data: Union[dict, str],
|
||||
# messages: List,
|
||||
# print_verbose,
|
||||
# encoding,
|
||||
# ) -> CustomStreamWrapper:
|
||||
# """
|
||||
# Return stream object for tool-calling + streaming
|
||||
# """
|
||||
# ## LOGGING
|
||||
# logging_obj.post_call(
|
||||
# input=messages,
|
||||
# api_key=api_key,
|
||||
# original_response=response.text,
|
||||
# additional_args={"complete_input_dict": data},
|
||||
# )
|
||||
# print_verbose(f"raw model_response: {response.text}")
|
||||
# ## RESPONSE OBJECT
|
||||
# try:
|
||||
# completion_response = response.json()
|
||||
# except:
|
||||
# raise AnthropicError(
|
||||
# message=response.text, status_code=response.status_code
|
||||
# )
|
||||
# text_content = ""
|
||||
# tool_calls = []
|
||||
# for content in completion_response["content"]:
|
||||
# if content["type"] == "text":
|
||||
# text_content += content["text"]
|
||||
# ## TOOL CALLING
|
||||
# elif content["type"] == "tool_use":
|
||||
# tool_calls.append(
|
||||
# {
|
||||
# "id": content["id"],
|
||||
# "type": "function",
|
||||
# "function": {
|
||||
# "name": content["name"],
|
||||
# "arguments": json.dumps(content["input"]),
|
||||
# },
|
||||
# }
|
||||
# )
|
||||
# if "error" in completion_response:
|
||||
# raise AnthropicError(
|
||||
# message=str(completion_response["error"]),
|
||||
# status_code=response.status_code,
|
||||
# )
|
||||
# _message = litellm.Message(
|
||||
# tool_calls=tool_calls,
|
||||
# content=text_content or None,
|
||||
# )
|
||||
# model_response.choices[0].message = _message # type: ignore
|
||||
# model_response._hidden_params["original_response"] = completion_response[
|
||||
# "content"
|
||||
# ] # allow user to access raw anthropic tool calling response
|
||||
|
||||
model_response.choices[0].finish_reason = map_finish_reason(
|
||||
completion_response["stop_reason"]
|
||||
)
|
||||
# model_response.choices[0].finish_reason = map_finish_reason(
|
||||
# completion_response["stop_reason"]
|
||||
# )
|
||||
|
||||
print_verbose("INSIDE ANTHROPIC STREAMING TOOL CALLING CONDITION BLOCK")
|
||||
# return an iterator
|
||||
streaming_model_response = ModelResponse(stream=True)
|
||||
streaming_model_response.choices[0].finish_reason = model_response.choices[ # type: ignore
|
||||
0
|
||||
].finish_reason
|
||||
# streaming_model_response.choices = [litellm.utils.StreamingChoices()]
|
||||
streaming_choice = litellm.utils.StreamingChoices()
|
||||
streaming_choice.index = model_response.choices[0].index
|
||||
_tool_calls = []
|
||||
print_verbose(
|
||||
f"type of model_response.choices[0]: {type(model_response.choices[0])}"
|
||||
)
|
||||
print_verbose(f"type of streaming_choice: {type(streaming_choice)}")
|
||||
if isinstance(model_response.choices[0], litellm.Choices):
|
||||
if getattr(
|
||||
model_response.choices[0].message, "tool_calls", None
|
||||
) is not None and isinstance(
|
||||
model_response.choices[0].message.tool_calls, list
|
||||
):
|
||||
for tool_call in model_response.choices[0].message.tool_calls:
|
||||
_tool_call = {**tool_call.dict(), "index": 0}
|
||||
_tool_calls.append(_tool_call)
|
||||
delta_obj = litellm.utils.Delta(
|
||||
content=getattr(model_response.choices[0].message, "content", None),
|
||||
role=model_response.choices[0].message.role,
|
||||
tool_calls=_tool_calls,
|
||||
)
|
||||
streaming_choice.delta = delta_obj
|
||||
streaming_model_response.choices = [streaming_choice]
|
||||
completion_stream = ModelResponseIterator(
|
||||
model_response=streaming_model_response
|
||||
)
|
||||
print_verbose(
|
||||
"Returns anthropic CustomStreamWrapper with 'cached_response' streaming object"
|
||||
)
|
||||
return CustomStreamWrapper(
|
||||
completion_stream=completion_stream,
|
||||
model=model,
|
||||
custom_llm_provider="cached_response",
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
else:
|
||||
raise AnthropicError(
|
||||
status_code=422,
|
||||
message="Unprocessable response object - {}".format(response.text),
|
||||
)
|
||||
# print_verbose("INSIDE ANTHROPIC STREAMING TOOL CALLING CONDITION BLOCK")
|
||||
# # return an iterator
|
||||
# streaming_model_response = ModelResponse(stream=True)
|
||||
# streaming_model_response.choices[0].finish_reason = model_response.choices[ # type: ignore
|
||||
# 0
|
||||
# ].finish_reason
|
||||
# # streaming_model_response.choices = [litellm.utils.StreamingChoices()]
|
||||
# streaming_choice = litellm.utils.StreamingChoices()
|
||||
# streaming_choice.index = model_response.choices[0].index
|
||||
# _tool_calls = []
|
||||
# print_verbose(
|
||||
# f"type of model_response.choices[0]: {type(model_response.choices[0])}"
|
||||
# )
|
||||
# print_verbose(f"type of streaming_choice: {type(streaming_choice)}")
|
||||
# if isinstance(model_response.choices[0], litellm.Choices):
|
||||
# if getattr(
|
||||
# model_response.choices[0].message, "tool_calls", None
|
||||
# ) is not None and isinstance(
|
||||
# model_response.choices[0].message.tool_calls, list
|
||||
# ):
|
||||
# for tool_call in model_response.choices[0].message.tool_calls:
|
||||
# _tool_call = {**tool_call.dict(), "index": 0}
|
||||
# _tool_calls.append(_tool_call)
|
||||
# delta_obj = litellm.utils.Delta(
|
||||
# content=getattr(model_response.choices[0].message, "content", None),
|
||||
# role=model_response.choices[0].message.role,
|
||||
# tool_calls=_tool_calls,
|
||||
# )
|
||||
# streaming_choice.delta = delta_obj
|
||||
# streaming_model_response.choices = [streaming_choice]
|
||||
# completion_stream = ModelResponseIterator(
|
||||
# model_response=streaming_model_response
|
||||
# )
|
||||
# print_verbose(
|
||||
# "Returns anthropic CustomStreamWrapper with 'cached_response' streaming object"
|
||||
# )
|
||||
# return CustomStreamWrapper(
|
||||
# completion_stream=completion_stream,
|
||||
# model=model,
|
||||
# custom_llm_provider="cached_response",
|
||||
# logging_obj=logging_obj,
|
||||
# )
|
||||
# else:
|
||||
# raise AnthropicError(
|
||||
# status_code=422,
|
||||
# message="Unprocessable response object - {}".format(response.text),
|
||||
# )
|
||||
|
||||
def process_response(
|
||||
self,
|
||||
@@ -484,21 +499,19 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
headers={},
|
||||
) -> Union[ModelResponse, CustomStreamWrapper]:
|
||||
async_handler = _get_async_httpx_client()
|
||||
response = await async_handler.post(api_base, headers=headers, json=data)
|
||||
if stream and _is_function_call:
|
||||
return self.process_streaming_response(
|
||||
model=model,
|
||||
response=response,
|
||||
model_response=model_response,
|
||||
stream=stream,
|
||||
logging_obj=logging_obj,
|
||||
|
||||
try:
|
||||
response = await async_handler.post(api_base, headers=headers, json=data)
|
||||
except Exception as e:
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
api_key=api_key,
|
||||
data=data,
|
||||
messages=messages,
|
||||
print_verbose=print_verbose,
|
||||
optional_params=optional_params,
|
||||
encoding=encoding,
|
||||
original_response=str(e),
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
raise e
|
||||
|
||||
return self.process_response(
|
||||
model=model,
|
||||
response=response,
|
||||
@@ -608,7 +621,7 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
print_verbose(f"_is_function_call: {_is_function_call}")
|
||||
if acompletion == True:
|
||||
if (
|
||||
stream and not _is_function_call
|
||||
stream is True
|
||||
): # if function call - fake the streaming (need complete blocks for output parsing in openai format)
|
||||
print_verbose("makes async anthropic streaming POST request")
|
||||
data["stream"] = stream
|
||||
@@ -652,7 +665,7 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
else:
|
||||
## COMPLETION CALL
|
||||
if (
|
||||
stream and not _is_function_call
|
||||
stream is True
|
||||
): # if function call - fake the streaming (need complete blocks for output parsing in openai format)
|
||||
print_verbose("makes anthropic streaming POST request")
|
||||
data["stream"] = stream
|
||||
@@ -668,7 +681,9 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
status_code=response.status_code, message=response.text
|
||||
)
|
||||
|
||||
completion_stream = response.iter_lines()
|
||||
completion_stream = ModelResponseIterator(
|
||||
streaming_response=response.iter_lines(), sync_stream=True
|
||||
)
|
||||
streaming_response = CustomStreamWrapper(
|
||||
completion_stream=completion_stream,
|
||||
model=model,
|
||||
@@ -686,20 +701,6 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
status_code=response.status_code, message=response.text
|
||||
)
|
||||
|
||||
if stream and _is_function_call:
|
||||
return self.process_streaming_response(
|
||||
model=model,
|
||||
response=response,
|
||||
model_response=model_response,
|
||||
stream=stream,
|
||||
logging_obj=logging_obj,
|
||||
api_key=api_key,
|
||||
data=data,
|
||||
messages=messages,
|
||||
print_verbose=print_verbose,
|
||||
optional_params=optional_params,
|
||||
encoding=encoding,
|
||||
)
|
||||
return self.process_response(
|
||||
model=model,
|
||||
response=response,
|
||||
@@ -720,26 +721,195 @@ class AnthropicChatCompletion(BaseLLM):
|
||||
|
||||
|
||||
class ModelResponseIterator:
|
||||
def __init__(self, model_response):
|
||||
self.model_response = model_response
|
||||
self.is_done = False
|
||||
def __init__(self, streaming_response, sync_stream: bool):
|
||||
self.streaming_response = streaming_response
|
||||
self.response_iterator = self.streaming_response
|
||||
|
||||
def chunk_parser(self, chunk: dict) -> GenericStreamingChunk:
|
||||
try:
|
||||
type_chunk = chunk.get("type", "") or ""
|
||||
|
||||
text = ""
|
||||
tool_use: Optional[ChatCompletionToolCallChunk] = None
|
||||
is_finished = False
|
||||
finish_reason = ""
|
||||
usage: Optional[ChatCompletionUsageBlock] = None
|
||||
|
||||
index = int(chunk.get("index", 0))
|
||||
if type_chunk == "content_block_delta":
|
||||
"""
|
||||
Anthropic content chunk
|
||||
chunk = {'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': 'Hello'}}
|
||||
"""
|
||||
content_block = ContentBlockDelta(**chunk) # type: ignore
|
||||
if "text" in content_block["delta"]:
|
||||
text = content_block["delta"]["text"]
|
||||
elif "partial_json" in content_block["delta"]:
|
||||
tool_use = {
|
||||
"id": None,
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": None,
|
||||
"arguments": content_block["delta"]["partial_json"],
|
||||
},
|
||||
}
|
||||
elif type_chunk == "content_block_start":
|
||||
"""
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_01T1x1fJ34qAmk2tNTrN7Up6","name":"get_weather","input":{}}}
|
||||
"""
|
||||
content_block_start = ContentBlockStart(**chunk) # type: ignore
|
||||
if content_block_start["content_block"]["type"] == "text":
|
||||
text = content_block_start["content_block"]["text"]
|
||||
elif content_block_start["content_block"]["type"] == "tool_use":
|
||||
tool_use = {
|
||||
"id": content_block_start["content_block"]["id"],
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": content_block_start["content_block"]["name"],
|
||||
"arguments": json.dumps(
|
||||
content_block_start["content_block"]["input"]
|
||||
),
|
||||
},
|
||||
}
|
||||
elif type_chunk == "message_delta":
|
||||
"""
|
||||
Anthropic
|
||||
chunk = {'type': 'message_delta', 'delta': {'stop_reason': 'max_tokens', 'stop_sequence': None}, 'usage': {'output_tokens': 10}}
|
||||
"""
|
||||
# TODO - get usage from this chunk, set in response
|
||||
message_delta = MessageBlockDelta(**chunk) # type: ignore
|
||||
finish_reason = map_finish_reason(
|
||||
finish_reason=message_delta["delta"].get("stop_reason", "stop")
|
||||
or "stop"
|
||||
)
|
||||
usage = ChatCompletionUsageBlock(
|
||||
prompt_tokens=message_delta["usage"].get("input_tokens", 0),
|
||||
completion_tokens=message_delta["usage"].get("output_tokens", 0),
|
||||
total_tokens=message_delta["usage"].get("input_tokens", 0)
|
||||
+ message_delta["usage"].get("output_tokens", 0),
|
||||
)
|
||||
is_finished = True
|
||||
elif type_chunk == "message_start":
|
||||
"""
|
||||
Anthropic
|
||||
chunk = {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_vrtx_011PqREFEMzd3REdCoUFAmdG",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-3-sonnet-20240229",
|
||||
"content": [],
|
||||
"stop_reason": null,
|
||||
"stop_sequence": null,
|
||||
"usage": {
|
||||
"input_tokens": 270,
|
||||
"output_tokens": 1
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
message_start_block = MessageStartBlock(**chunk) # type: ignore
|
||||
usage = ChatCompletionUsageBlock(
|
||||
prompt_tokens=message_start_block["message"]
|
||||
.get("usage", {})
|
||||
.get("input_tokens", 0),
|
||||
completion_tokens=message_start_block["message"]
|
||||
.get("usage", {})
|
||||
.get("output_tokens", 0),
|
||||
total_tokens=message_start_block["message"]
|
||||
.get("usage", {})
|
||||
.get("input_tokens", 0)
|
||||
+ message_start_block["message"]
|
||||
.get("usage", {})
|
||||
.get("output_tokens", 0),
|
||||
)
|
||||
returned_chunk = GenericStreamingChunk(
|
||||
text=text,
|
||||
tool_use=tool_use,
|
||||
is_finished=is_finished,
|
||||
finish_reason=finish_reason,
|
||||
usage=usage,
|
||||
index=index,
|
||||
)
|
||||
|
||||
return returned_chunk
|
||||
|
||||
except json.JSONDecodeError:
|
||||
raise ValueError(f"Failed to decode JSON from chunk: {chunk}")
|
||||
|
||||
# Sync iterator
|
||||
def __iter__(self):
|
||||
return self
|
||||
|
||||
def __next__(self):
|
||||
if self.is_done:
|
||||
try:
|
||||
chunk = self.response_iterator.__next__()
|
||||
except StopIteration:
|
||||
raise StopIteration
|
||||
self.is_done = True
|
||||
return self.model_response
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
||||
try:
|
||||
str_line = chunk
|
||||
if isinstance(chunk, bytes): # Handle binary data
|
||||
str_line = chunk.decode("utf-8") # Convert bytes to string
|
||||
index = str_line.find("data:")
|
||||
if index != -1:
|
||||
str_line = str_line[index:]
|
||||
|
||||
if str_line.startswith("data:"):
|
||||
data_json = json.loads(str_line[5:])
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
else:
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=0,
|
||||
tool_use=None,
|
||||
)
|
||||
except StopIteration:
|
||||
raise StopIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
|
||||
|
||||
# Async iterator
|
||||
def __aiter__(self):
|
||||
self.async_response_iterator = self.streaming_response.__aiter__()
|
||||
return self
|
||||
|
||||
async def __anext__(self):
|
||||
if self.is_done:
|
||||
try:
|
||||
chunk = await self.async_response_iterator.__anext__()
|
||||
except StopAsyncIteration:
|
||||
raise StopAsyncIteration
|
||||
self.is_done = True
|
||||
return self.model_response
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error receiving chunk from stream: {e}")
|
||||
|
||||
try:
|
||||
str_line = chunk
|
||||
if isinstance(chunk, bytes): # Handle binary data
|
||||
str_line = chunk.decode("utf-8") # Convert bytes to string
|
||||
index = str_line.find("data:")
|
||||
if index != -1:
|
||||
str_line = str_line[index:]
|
||||
|
||||
if str_line.startswith("data:"):
|
||||
data_json = json.loads(str_line[5:])
|
||||
return self.chunk_parser(chunk=data_json)
|
||||
else:
|
||||
return GenericStreamingChunk(
|
||||
text="",
|
||||
is_finished=False,
|
||||
finish_reason="",
|
||||
usage=None,
|
||||
index=0,
|
||||
tool_use=None,
|
||||
)
|
||||
except StopAsyncIteration:
|
||||
raise StopAsyncIteration
|
||||
except ValueError as e:
|
||||
raise RuntimeError(f"Error parsing chunk: {e},\nReceived chunk: {chunk}")
|
||||
|
||||
@@ -1283,7 +1283,9 @@ def anthropic_messages_pt(messages: list):
|
||||
)
|
||||
else:
|
||||
raise Exception(
|
||||
"Invalid first message. Should always start with 'role'='user' for Anthropic. System prompt is sent separately for Anthropic. set 'litellm.modify_params = True' or 'litellm_settings:modify_params = True' on proxy, to insert a placeholder user message - '.' as the first message, "
|
||||
"Invalid first message={}. Should always start with 'role'='user' for Anthropic. System prompt is sent separately for Anthropic. set 'litellm.modify_params = True' or 'litellm_settings:modify_params = True' on proxy, to insert a placeholder user message - '.' as the first message, ".format(
|
||||
new_messages
|
||||
)
|
||||
)
|
||||
|
||||
if new_messages[-1]["role"] == "assistant":
|
||||
|
||||
@@ -235,190 +235,44 @@ def completion(
|
||||
if k not in optional_params:
|
||||
optional_params[k] = v
|
||||
|
||||
## Format Prompt
|
||||
_is_function_call = False
|
||||
_is_json_schema = False
|
||||
messages = copy.deepcopy(messages)
|
||||
optional_params = copy.deepcopy(optional_params)
|
||||
# Separate system prompt from rest of message
|
||||
system_prompt_indices = []
|
||||
system_prompt = ""
|
||||
for idx, message in enumerate(messages):
|
||||
if message["role"] == "system":
|
||||
system_prompt += message["content"]
|
||||
system_prompt_indices.append(idx)
|
||||
if len(system_prompt_indices) > 0:
|
||||
for idx in reversed(system_prompt_indices):
|
||||
messages.pop(idx)
|
||||
if len(system_prompt) > 0:
|
||||
optional_params["system"] = system_prompt
|
||||
# Checks for 'response_schema' support - if passed in
|
||||
if "response_format" in optional_params:
|
||||
response_format_chunk = ResponseFormatChunk(
|
||||
**optional_params["response_format"] # type: ignore
|
||||
)
|
||||
supports_response_schema = litellm.supports_response_schema(
|
||||
model=model, custom_llm_provider="vertex_ai"
|
||||
)
|
||||
if (
|
||||
supports_response_schema is False
|
||||
and response_format_chunk["type"] == "json_object"
|
||||
and "response_schema" in response_format_chunk
|
||||
):
|
||||
_is_json_schema = True
|
||||
user_response_schema_message = response_schema_prompt(
|
||||
model=model,
|
||||
response_schema=response_format_chunk["response_schema"],
|
||||
)
|
||||
messages.append(
|
||||
{"role": "user", "content": user_response_schema_message}
|
||||
)
|
||||
messages.append({"role": "assistant", "content": "{"})
|
||||
optional_params.pop("response_format")
|
||||
# Format rest of message according to anthropic guidelines
|
||||
try:
|
||||
messages = prompt_factory(
|
||||
model=model, messages=messages, custom_llm_provider="anthropic_xml"
|
||||
)
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=400, message=str(e))
|
||||
## CONSTRUCT API BASE
|
||||
stream = optional_params.get("stream", False)
|
||||
|
||||
## Handle Tool Calling
|
||||
if "tools" in optional_params:
|
||||
_is_function_call = True
|
||||
tool_calling_system_prompt = construct_tool_use_system_prompt(
|
||||
tools=optional_params["tools"]
|
||||
)
|
||||
optional_params["system"] = (
|
||||
optional_params.get("system", "\n") + tool_calling_system_prompt
|
||||
) # add the anthropic tool calling prompt to the system prompt
|
||||
optional_params.pop("tools")
|
||||
|
||||
stream = optional_params.pop("stream", None)
|
||||
|
||||
data = {
|
||||
"model": model,
|
||||
"messages": messages,
|
||||
**optional_params,
|
||||
}
|
||||
print_verbose(f"_is_function_call: {_is_function_call}")
|
||||
|
||||
## Completion Call
|
||||
|
||||
print_verbose(
|
||||
f"VERTEX AI: vertex_project={vertex_project}; vertex_location={vertex_location}; vertex_credentials={vertex_credentials}"
|
||||
api_base = create_vertex_anthropic_url(
|
||||
vertex_location=vertex_location or "us-central1",
|
||||
vertex_project=vertex_project or project_id,
|
||||
model=model,
|
||||
stream=stream,
|
||||
)
|
||||
|
||||
if acompletion == True:
|
||||
"""
|
||||
- async streaming
|
||||
- async completion
|
||||
"""
|
||||
if stream is not None and stream == True:
|
||||
return async_streaming(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=data,
|
||||
print_verbose=print_verbose,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
optional_params=optional_params,
|
||||
client=client,
|
||||
access_token=access_token,
|
||||
)
|
||||
else:
|
||||
return async_completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
data=data,
|
||||
print_verbose=print_verbose,
|
||||
model_response=model_response,
|
||||
logging_obj=logging_obj,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
optional_params=optional_params,
|
||||
client=client,
|
||||
access_token=access_token,
|
||||
)
|
||||
if stream is not None and stream == True:
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=None,
|
||||
additional_args={
|
||||
"complete_input_dict": optional_params,
|
||||
},
|
||||
)
|
||||
response = vertex_ai_client.messages.create(**data, stream=True) # type: ignore
|
||||
return response
|
||||
|
||||
## LOGGING
|
||||
logging_obj.pre_call(
|
||||
input=messages,
|
||||
api_key=None,
|
||||
additional_args={
|
||||
"complete_input_dict": optional_params,
|
||||
},
|
||||
)
|
||||
|
||||
message = vertex_ai_client.messages.create(**data) # type: ignore
|
||||
|
||||
## LOGGING
|
||||
logging_obj.post_call(
|
||||
input=messages,
|
||||
api_key="",
|
||||
original_response=message,
|
||||
additional_args={"complete_input_dict": data},
|
||||
)
|
||||
|
||||
text_content: str = message.content[0].text
|
||||
## TOOL CALLING - OUTPUT PARSE
|
||||
if text_content is not None and contains_tag("invoke", text_content):
|
||||
function_name = extract_between_tags("tool_name", text_content)[0]
|
||||
function_arguments_str = extract_between_tags("invoke", text_content)[
|
||||
0
|
||||
].strip()
|
||||
function_arguments_str = f"<invoke>{function_arguments_str}</invoke>"
|
||||
function_arguments = parse_xml_params(function_arguments_str)
|
||||
_message = litellm.Message(
|
||||
tool_calls=[
|
||||
{
|
||||
"id": f"call_{uuid.uuid4()}",
|
||||
"type": "function",
|
||||
"function": {
|
||||
"name": function_name,
|
||||
"arguments": json.dumps(function_arguments),
|
||||
},
|
||||
}
|
||||
],
|
||||
content=None,
|
||||
)
|
||||
model_response.choices[0].message = _message # type: ignore
|
||||
if headers is not None:
|
||||
vertex_headers = headers
|
||||
else:
|
||||
if (
|
||||
_is_json_schema
|
||||
): # follows https://github.com/anthropics/anthropic-cookbook/blob/main/misc/how_to_enable_json_mode.ipynb
|
||||
json_response = "{" + text_content[: text_content.rfind("}") + 1]
|
||||
model_response.choices[0].message.content = json_response # type: ignore
|
||||
else:
|
||||
model_response.choices[0].message.content = text_content # type: ignore
|
||||
model_response.choices[0].finish_reason = map_finish_reason(message.stop_reason)
|
||||
vertex_headers = {}
|
||||
|
||||
## CALCULATING USAGE
|
||||
prompt_tokens = message.usage.input_tokens
|
||||
completion_tokens = message.usage.output_tokens
|
||||
vertex_headers.update({"Authorization": "Bearer {}".format(access_token)})
|
||||
|
||||
model_response["created"] = int(time.time())
|
||||
model_response["model"] = model
|
||||
usage = Usage(
|
||||
prompt_tokens=prompt_tokens,
|
||||
completion_tokens=completion_tokens,
|
||||
total_tokens=prompt_tokens + completion_tokens,
|
||||
optional_params.update(
|
||||
{"anthropic_version": "vertex-2023-10-16", "is_vertex_request": True}
|
||||
)
|
||||
setattr(model_response, "usage", usage)
|
||||
return model_response
|
||||
|
||||
return anthropic_chat_completions.completion(
|
||||
model=model,
|
||||
messages=messages,
|
||||
api_base=api_base,
|
||||
custom_prompt_dict=custom_prompt_dict,
|
||||
model_response=model_response,
|
||||
print_verbose=print_verbose,
|
||||
encoding=encoding,
|
||||
api_key=access_token,
|
||||
logging_obj=logging_obj,
|
||||
optional_params=optional_params,
|
||||
acompletion=acompletion,
|
||||
litellm_params=litellm_params,
|
||||
logger_fn=logger_fn,
|
||||
headers=vertex_headers,
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
raise VertexAIError(status_code=500, message=str(e))
|
||||
|
||||
|
||||
+23
-14
@@ -2026,18 +2026,18 @@ def completion(
|
||||
acompletion=acompletion,
|
||||
)
|
||||
|
||||
if (
|
||||
"stream" in optional_params
|
||||
and optional_params["stream"] == True
|
||||
and acompletion == False
|
||||
):
|
||||
response = CustomStreamWrapper(
|
||||
model_response,
|
||||
model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
logging_obj=logging,
|
||||
)
|
||||
return response
|
||||
if (
|
||||
"stream" in optional_params
|
||||
and optional_params["stream"] == True
|
||||
and acompletion == False
|
||||
):
|
||||
response = CustomStreamWrapper(
|
||||
model_response,
|
||||
model,
|
||||
custom_llm_provider="vertex_ai",
|
||||
logging_obj=logging,
|
||||
)
|
||||
return response
|
||||
response = model_response
|
||||
elif custom_llm_provider == "predibase":
|
||||
tenant_id = (
|
||||
@@ -4944,14 +4944,23 @@ def stream_chunk_builder(
|
||||
else:
|
||||
completion_output = ""
|
||||
# # Update usage information if needed
|
||||
prompt_tokens = 0
|
||||
completion_tokens = 0
|
||||
for chunk in chunks:
|
||||
if "usage" in chunk:
|
||||
if "prompt_tokens" in chunk["usage"]:
|
||||
prompt_tokens += chunk["usage"].get("prompt_tokens", 0) or 0
|
||||
if "completion_tokens" in chunk["usage"]:
|
||||
completion_tokens += chunk["usage"].get("completion_tokens", 0) or 0
|
||||
|
||||
try:
|
||||
response["usage"]["prompt_tokens"] = token_counter(
|
||||
response["usage"]["prompt_tokens"] = prompt_tokens or token_counter(
|
||||
model=model, messages=messages
|
||||
)
|
||||
except: # don't allow this failing to block a complete streaming response from being returned
|
||||
print_verbose(f"token_counter failed, assuming prompt tokens is 0")
|
||||
response["usage"]["prompt_tokens"] = 0
|
||||
response["usage"]["completion_tokens"] = token_counter(
|
||||
response["usage"]["completion_tokens"] = completion_tokens or token_counter(
|
||||
model=model,
|
||||
text=completion_output,
|
||||
count_response_tokens=True, # count_response_tokens is a Flag to tell token counter this is a response, No need to add extra tokens we do for input messages
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
model_list:
|
||||
- model_name: claude-3-5-sonnet
|
||||
litellm_params:
|
||||
model: anthropic/claude-3-5-sonnet
|
||||
model: claude-3-haiku-20240307
|
||||
# - model_name: gemini-1.5-flash-gemini
|
||||
# litellm_params:
|
||||
# model: vertex_ai_beta/gemini-1.5-flash
|
||||
@@ -18,7 +18,6 @@ model_list:
|
||||
- model_name: fake-openai-endpoint
|
||||
litellm_params:
|
||||
model: predibase/llama-3-8b-instruct
|
||||
api_base: "http://0.0.0.0:8081"
|
||||
api_key: os.environ/PREDIBASE_API_KEY
|
||||
tenant_id: os.environ/PREDIBASE_TENANT_ID
|
||||
max_new_tokens: 256
|
||||
|
||||
@@ -203,7 +203,7 @@ def test_vertex_ai_anthropic():
|
||||
# )
|
||||
def test_vertex_ai_anthropic_streaming():
|
||||
try:
|
||||
# load_vertex_ai_credentials()
|
||||
load_vertex_ai_credentials()
|
||||
|
||||
# litellm.set_verbose = True
|
||||
|
||||
@@ -223,8 +223,9 @@ def test_vertex_ai_anthropic_streaming():
|
||||
stream=True,
|
||||
)
|
||||
# print("\nModel Response", response)
|
||||
for chunk in response:
|
||||
for idx, chunk in enumerate(response):
|
||||
print(f"chunk: {chunk}")
|
||||
streaming_format_tests(idx=idx, chunk=chunk)
|
||||
|
||||
# raise Exception("it worked!")
|
||||
except litellm.RateLimitError as e:
|
||||
@@ -294,8 +295,10 @@ async def test_vertex_ai_anthropic_async_streaming():
|
||||
stream=True,
|
||||
)
|
||||
|
||||
idx = 0
|
||||
async for chunk in response:
|
||||
print(f"chunk: {chunk}")
|
||||
streaming_format_tests(idx=idx, chunk=chunk)
|
||||
idx += 1
|
||||
except litellm.RateLimitError as e:
|
||||
pass
|
||||
except Exception as e:
|
||||
|
||||
@@ -2020,7 +2020,7 @@ def test_openai_chat_completion_complete_response_call():
|
||||
# test_openai_chat_completion_complete_response_call()
|
||||
@pytest.mark.parametrize(
|
||||
"model",
|
||||
["gpt-3.5-turbo", "azure/chatgpt-v-2"],
|
||||
["gpt-3.5-turbo", "azure/chatgpt-v-2", "claude-3-haiku-20240307"], #
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"sync",
|
||||
@@ -2028,14 +2028,14 @@ def test_openai_chat_completion_complete_response_call():
|
||||
)
|
||||
@pytest.mark.asyncio
|
||||
async def test_openai_stream_options_call(model, sync):
|
||||
litellm.set_verbose = False
|
||||
litellm.set_verbose = True
|
||||
usage = None
|
||||
chunks = []
|
||||
if sync:
|
||||
response = litellm.completion(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "say GM - we're going to make it "}
|
||||
{"role": "user", "content": "say GM - we're going to make it "},
|
||||
],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
@@ -2047,9 +2047,7 @@ async def test_openai_stream_options_call(model, sync):
|
||||
else:
|
||||
response = await litellm.acompletion(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "system", "content": "say GM - we're going to make it "}
|
||||
],
|
||||
messages=[{"role": "user", "content": "say GM - we're going to make it "}],
|
||||
stream=True,
|
||||
stream_options={"include_usage": True},
|
||||
max_tokens=10,
|
||||
@@ -2746,7 +2744,7 @@ class Chunk2(BaseModel):
|
||||
object: str
|
||||
created: int
|
||||
model: str
|
||||
system_fingerprint: str
|
||||
system_fingerprint: Optional[str]
|
||||
choices: List[Choices2]
|
||||
|
||||
|
||||
@@ -3001,7 +2999,7 @@ def test_completion_claude_3_function_call_with_streaming():
|
||||
model="claude-3-opus-20240229",
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
tool_choice="required",
|
||||
stream=True,
|
||||
)
|
||||
idx = 0
|
||||
@@ -3060,7 +3058,7 @@ async def test_acompletion_claude_3_function_call_with_streaming():
|
||||
model="claude-3-opus-20240229",
|
||||
messages=messages,
|
||||
tools=tools,
|
||||
tool_choice="auto",
|
||||
tool_choice="required",
|
||||
stream=True,
|
||||
)
|
||||
idx = 0
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
from typing import List, Optional, Union, Iterable
|
||||
from typing import Iterable, List, Optional, Union
|
||||
|
||||
from pydantic import BaseModel, validator
|
||||
|
||||
from typing_extensions import Literal, Required, TypedDict
|
||||
|
||||
|
||||
@@ -45,3 +44,114 @@ class AnthopicMessagesAssistantMessageParam(TypedDict, total=False):
|
||||
Provides the model information to differentiate between participants of the same
|
||||
role.
|
||||
"""
|
||||
|
||||
|
||||
class ContentTextBlockDelta(TypedDict):
|
||||
"""
|
||||
'delta': {'type': 'text_delta', 'text': 'Hello'}
|
||||
"""
|
||||
|
||||
type: str
|
||||
text: str
|
||||
|
||||
|
||||
class ContentJsonBlockDelta(TypedDict):
|
||||
"""
|
||||
"delta": {"type": "input_json_delta","partial_json": "{\"location\": \"San Fra"}}
|
||||
"""
|
||||
|
||||
type: str
|
||||
partial_json: str
|
||||
|
||||
|
||||
class ContentBlockDelta(TypedDict):
|
||||
type: str
|
||||
index: int
|
||||
delta: Union[ContentTextBlockDelta, ContentJsonBlockDelta]
|
||||
|
||||
|
||||
class ToolUseBlock(TypedDict):
|
||||
"""
|
||||
"content_block":{"type":"tool_use","id":"toolu_01T1x1fJ34qAmk2tNTrN7Up6","name":"get_weather","input":{}}
|
||||
"""
|
||||
|
||||
id: str
|
||||
|
||||
input: dict
|
||||
|
||||
name: str
|
||||
|
||||
type: Literal["tool_use"]
|
||||
|
||||
|
||||
class TextBlock(TypedDict):
|
||||
text: str
|
||||
|
||||
type: Literal["text"]
|
||||
|
||||
|
||||
class ContentBlockStart(TypedDict):
|
||||
"""
|
||||
event: content_block_start
|
||||
data: {"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"toolu_01T1x1fJ34qAmk2tNTrN7Up6","name":"get_weather","input":{}}}
|
||||
"""
|
||||
|
||||
type: str
|
||||
index: int
|
||||
content_block: Union[ToolUseBlock, TextBlock]
|
||||
|
||||
|
||||
class MessageDelta(TypedDict, total=False):
|
||||
stop_reason: Optional[str]
|
||||
|
||||
|
||||
class UsageDelta(TypedDict, total=False):
|
||||
input_tokens: int
|
||||
output_tokens: int
|
||||
|
||||
|
||||
class MessageBlockDelta(TypedDict):
|
||||
"""
|
||||
Anthropic
|
||||
chunk = {'type': 'message_delta', 'delta': {'stop_reason': 'max_tokens', 'stop_sequence': None}, 'usage': {'output_tokens': 10}}
|
||||
"""
|
||||
|
||||
type: Literal["message_delta"]
|
||||
delta: MessageDelta
|
||||
usage: UsageDelta
|
||||
|
||||
|
||||
class MessageChunk(TypedDict, total=False):
|
||||
id: str
|
||||
type: str
|
||||
role: str
|
||||
model: str
|
||||
content: List
|
||||
stop_reason: Optional[str]
|
||||
stop_sequence: Optional[str]
|
||||
usage: UsageDelta
|
||||
|
||||
|
||||
class MessageStartBlock(TypedDict):
|
||||
"""
|
||||
Anthropic
|
||||
chunk = {
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"id": "msg_vrtx_011PqREFEMzd3REdCoUFAmdG",
|
||||
"type": "message",
|
||||
"role": "assistant",
|
||||
"model": "claude-3-sonnet-20240229",
|
||||
"content": [],
|
||||
"stop_reason": null,
|
||||
"stop_sequence": null,
|
||||
"usage": {
|
||||
"input_tokens": 270,
|
||||
"output_tokens": 1
|
||||
}
|
||||
}
|
||||
}
|
||||
"""
|
||||
|
||||
type: Literal["message_start"]
|
||||
message: MessageChunk
|
||||
|
||||
+98
-56
@@ -8003,6 +8003,11 @@ class CustomStreamWrapper:
|
||||
return hold, curr_chunk
|
||||
|
||||
def handle_anthropic_text_chunk(self, chunk):
|
||||
"""
|
||||
For old anthropic models - claude-1, claude-2.
|
||||
|
||||
Claude-3 is handled from within Anthropic.py VIA ModelResponseIterator()
|
||||
"""
|
||||
str_line = chunk
|
||||
if isinstance(chunk, bytes): # Handle binary data
|
||||
str_line = chunk.decode("utf-8") # Convert bytes to string
|
||||
@@ -8031,44 +8036,6 @@ class CustomStreamWrapper:
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
|
||||
def handle_anthropic_chunk(self, chunk):
|
||||
str_line = chunk
|
||||
if isinstance(chunk, bytes): # Handle binary data
|
||||
str_line = chunk.decode("utf-8") # Convert bytes to string
|
||||
text = ""
|
||||
is_finished = False
|
||||
finish_reason = None
|
||||
if str_line.startswith("data:"):
|
||||
data_json = json.loads(str_line[5:])
|
||||
type_chunk = data_json.get("type", None)
|
||||
if type_chunk == "content_block_delta":
|
||||
"""
|
||||
Anthropic content chunk
|
||||
chunk = {'type': 'content_block_delta', 'index': 0, 'delta': {'type': 'text_delta', 'text': 'Hello'}}
|
||||
"""
|
||||
text = data_json.get("delta", {}).get("text", "")
|
||||
elif type_chunk == "message_delta":
|
||||
"""
|
||||
Anthropic
|
||||
chunk = {'type': 'message_delta', 'delta': {'stop_reason': 'max_tokens', 'stop_sequence': None}, 'usage': {'output_tokens': 10}}
|
||||
"""
|
||||
# TODO - get usage from this chunk, set in response
|
||||
finish_reason = data_json.get("delta", {}).get("stop_reason", None)
|
||||
is_finished = True
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
elif "error" in str_line:
|
||||
raise ValueError(f"Unable to parse response. Original response: {str_line}")
|
||||
else:
|
||||
return {
|
||||
"text": text,
|
||||
"is_finished": is_finished,
|
||||
"finish_reason": finish_reason,
|
||||
}
|
||||
|
||||
def handle_vertexai_anthropic_chunk(self, chunk):
|
||||
"""
|
||||
- MessageStartEvent(message=Message(id='msg_01LeRRgvX4gwkX3ryBVgtuYZ', content=[], model='claude-3-sonnet-20240229', role='assistant', stop_reason=None, stop_sequence=None, type='message', usage=Usage(input_tokens=8, output_tokens=1)), type='message_start'); custom_llm_provider: vertex_ai
|
||||
@@ -8779,7 +8746,7 @@ class CustomStreamWrapper:
|
||||
verbose_logger.debug(traceback.format_exc())
|
||||
return ""
|
||||
|
||||
def model_response_creator(self):
|
||||
def model_response_creator(self, chunk: Optional[dict] = None):
|
||||
_model = self.model
|
||||
_received_llm_provider = self.custom_llm_provider
|
||||
_logging_obj_llm_provider = self.logging_obj.model_call_details.get("custom_llm_provider", None) # type: ignore
|
||||
@@ -8788,13 +8755,18 @@ class CustomStreamWrapper:
|
||||
and _received_llm_provider != _logging_obj_llm_provider
|
||||
):
|
||||
_model = "{}/{}".format(_logging_obj_llm_provider, _model)
|
||||
if chunk is None:
|
||||
chunk = {}
|
||||
else:
|
||||
# pop model keyword
|
||||
chunk.pop("model", None)
|
||||
model_response = ModelResponse(
|
||||
stream=True, model=_model, stream_options=self.stream_options
|
||||
stream=True, model=_model, stream_options=self.stream_options, **chunk
|
||||
)
|
||||
if self.response_id is not None:
|
||||
model_response.id = self.response_id
|
||||
else:
|
||||
self.response_id = model_response.id
|
||||
self.response_id = model_response.id # type: ignore
|
||||
if self.system_fingerprint is not None:
|
||||
model_response.system_fingerprint = self.system_fingerprint
|
||||
model_response._hidden_params["custom_llm_provider"] = _logging_obj_llm_provider
|
||||
@@ -8819,10 +8791,37 @@ class CustomStreamWrapper:
|
||||
# return this for all models
|
||||
completion_obj = {"content": ""}
|
||||
if self.custom_llm_provider and self.custom_llm_provider == "anthropic":
|
||||
response_obj = self.handle_anthropic_chunk(chunk)
|
||||
completion_obj["content"] = response_obj["text"]
|
||||
if response_obj["is_finished"]:
|
||||
self.received_finish_reason = response_obj["finish_reason"]
|
||||
from litellm.types.utils import GenericStreamingChunk as GChunk
|
||||
|
||||
if self.received_finish_reason is not None:
|
||||
raise StopIteration
|
||||
anthropic_response_obj: GChunk = chunk
|
||||
completion_obj["content"] = anthropic_response_obj["text"]
|
||||
if anthropic_response_obj["is_finished"]:
|
||||
self.received_finish_reason = anthropic_response_obj[
|
||||
"finish_reason"
|
||||
]
|
||||
|
||||
if (
|
||||
self.stream_options
|
||||
and self.stream_options.get("include_usage", False) is True
|
||||
and anthropic_response_obj["usage"] is not None
|
||||
):
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=anthropic_response_obj["usage"]["prompt_tokens"],
|
||||
completion_tokens=anthropic_response_obj["usage"][
|
||||
"completion_tokens"
|
||||
],
|
||||
total_tokens=anthropic_response_obj["usage"]["total_tokens"],
|
||||
)
|
||||
|
||||
if (
|
||||
"tool_use" in anthropic_response_obj
|
||||
and anthropic_response_obj["tool_use"] is not None
|
||||
):
|
||||
completion_obj["tool_calls"] = [anthropic_response_obj["tool_use"]]
|
||||
|
||||
response_obj = anthropic_response_obj
|
||||
elif (
|
||||
self.custom_llm_provider
|
||||
and self.custom_llm_provider == "anthropic_text"
|
||||
@@ -8931,7 +8930,6 @@ class CustomStreamWrapper:
|
||||
and self.stream_options.get("include_usage", False) is True
|
||||
and response_obj["usage"] is not None
|
||||
):
|
||||
self.sent_stream_usage = True
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"]["prompt_tokens"],
|
||||
completion_tokens=response_obj["usage"]["completion_tokens"],
|
||||
@@ -9059,7 +9057,6 @@ class CustomStreamWrapper:
|
||||
and self.stream_options.get("include_usage", False) is True
|
||||
and response_obj["usage"] is not None
|
||||
):
|
||||
self.sent_stream_usage = True
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"]["inputTokens"],
|
||||
completion_tokens=response_obj["usage"]["outputTokens"],
|
||||
@@ -9131,7 +9128,6 @@ class CustomStreamWrapper:
|
||||
and self.stream_options.get("include_usage", False) == True
|
||||
and response_obj["usage"] is not None
|
||||
):
|
||||
self.sent_stream_usage = True
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
@@ -9150,7 +9146,6 @@ class CustomStreamWrapper:
|
||||
and self.stream_options.get("include_usage", False) == True
|
||||
and response_obj["usage"] is not None
|
||||
):
|
||||
self.sent_stream_usage = True
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
@@ -9167,7 +9162,6 @@ class CustomStreamWrapper:
|
||||
and self.stream_options.get("include_usage", False) == True
|
||||
and response_obj["usage"] is not None
|
||||
):
|
||||
self.sent_stream_usage = True
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
@@ -9231,7 +9225,6 @@ class CustomStreamWrapper:
|
||||
and self.stream_options["include_usage"] == True
|
||||
and response_obj["usage"] is not None
|
||||
):
|
||||
self.sent_stream_usage = True
|
||||
model_response.usage = litellm.Usage(
|
||||
prompt_tokens=response_obj["usage"].prompt_tokens,
|
||||
completion_tokens=response_obj["usage"].completion_tokens,
|
||||
@@ -9556,9 +9549,24 @@ class CustomStreamWrapper:
|
||||
self.rules.post_call_rules(
|
||||
input=self.response_uptil_now, model=self.model
|
||||
)
|
||||
# RETURN RESULT
|
||||
# HANDLE STREAM OPTIONS
|
||||
self.chunks.append(response)
|
||||
if hasattr(
|
||||
response, "usage"
|
||||
): # remove usage from chunk, only send on final chunk
|
||||
# Convert the object to a dictionary
|
||||
obj_dict = response.dict()
|
||||
|
||||
# Remove an attribute (e.g., 'attr2')
|
||||
if "usage" in obj_dict:
|
||||
del obj_dict["usage"]
|
||||
|
||||
# Create a new object without the removed attribute
|
||||
response = self.model_response_creator(chunk=obj_dict)
|
||||
|
||||
# RETURN RESULT
|
||||
return response
|
||||
|
||||
except StopIteration:
|
||||
if self.sent_last_chunk == True:
|
||||
if (
|
||||
@@ -9686,6 +9694,18 @@ class CustomStreamWrapper:
|
||||
)
|
||||
print_verbose(f"final returned processed chunk: {processed_chunk}")
|
||||
self.chunks.append(processed_chunk)
|
||||
if hasattr(
|
||||
processed_chunk, "usage"
|
||||
): # remove usage from chunk, only send on final chunk
|
||||
# Convert the object to a dictionary
|
||||
obj_dict = processed_chunk.dict()
|
||||
|
||||
# Remove an attribute (e.g., 'attr2')
|
||||
if "usage" in obj_dict:
|
||||
del obj_dict["usage"]
|
||||
|
||||
# Create a new object without the removed attribute
|
||||
processed_chunk = self.model_response_creator(chunk=obj_dict)
|
||||
return processed_chunk
|
||||
raise StopAsyncIteration
|
||||
else: # temporary patch for non-aiohttp async calls
|
||||
@@ -9728,11 +9748,11 @@ class CustomStreamWrapper:
|
||||
self.chunks.append(processed_chunk)
|
||||
return processed_chunk
|
||||
except StopAsyncIteration:
|
||||
if self.sent_last_chunk == True:
|
||||
if self.sent_last_chunk is True:
|
||||
if (
|
||||
self.sent_stream_usage == False
|
||||
self.sent_stream_usage is False
|
||||
and self.stream_options is not None
|
||||
and self.stream_options.get("include_usage", False) == True
|
||||
and self.stream_options.get("include_usage", False) is True
|
||||
):
|
||||
# send the final chunk with stream options
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
@@ -9766,7 +9786,29 @@ class CustomStreamWrapper:
|
||||
)
|
||||
return processed_chunk
|
||||
except StopIteration:
|
||||
if self.sent_last_chunk == True:
|
||||
if self.sent_last_chunk is True:
|
||||
if (
|
||||
self.sent_stream_usage is False
|
||||
and self.stream_options is not None
|
||||
and self.stream_options.get("include_usage", False) is True
|
||||
):
|
||||
# send the final chunk with stream options
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=self.chunks, messages=self.messages
|
||||
)
|
||||
response = self.model_response_creator()
|
||||
response.usage = complete_streaming_response.usage
|
||||
## LOGGING
|
||||
threading.Thread(
|
||||
target=self.logging_obj.success_handler, args=(response,)
|
||||
).start() # log response
|
||||
asyncio.create_task(
|
||||
self.logging_obj.async_success_handler(
|
||||
response,
|
||||
)
|
||||
)
|
||||
self.sent_stream_usage = True
|
||||
return response
|
||||
raise StopAsyncIteration
|
||||
else:
|
||||
self.sent_last_chunk = True
|
||||
|
||||
Reference in New Issue
Block a user