mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-27 03:34:10 +00:00
(feat) Add usage tracking for streaming /anthropic passthrough routes (#6842)
* use 1 file for AnthropicPassthroughLoggingHandler * add support for anthropic streaming usage tracking * ci/cd run again * fix - add real streaming for anthropic pass through * remove unused function stream_response * working anthropic streaming logging * fix code quality * fix use 1 file for vertex success handler * use helper for _handle_logging_vertex_collected_chunks * enforce vertex streaming to use sse for streaming * test test_basic_vertex_ai_pass_through_streaming_with_spendlog * fix type hints * add comment * fix linting * add pass through logging unit testing
This commit is contained in:
Ishaan Jaff
GitHub
parent
920f4c9f82
commit
b8af46e1a2
12 changed files with 688 additions and 295 deletions
|
|
@ -0,0 +1,206 @@
|
|||
import json
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.handler import (
|
||||
ModelResponseIterator as AnthropicModelResponseIterator,
|
||||
)
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..success_handler import PassThroughEndpointLogging
|
||||
from ..types import EndpointType
|
||||
else:
|
||||
PassThroughEndpointLogging = Any
|
||||
EndpointType = Any
|
||||
|
||||
|
||||
class AnthropicPassthroughLoggingHandler:
|
||||
|
||||
@staticmethod
|
||||
async def anthropic_passthrough_handler(
|
||||
httpx_response: httpx.Response,
|
||||
response_body: dict,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
url_route: str,
|
||||
result: str,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
cache_hit: bool,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Transforms Anthropic response to OpenAI response, generates a standard logging object so downstream logging can be handled
|
||||
"""
|
||||
model = response_body.get("model", "")
|
||||
litellm_model_response: litellm.ModelResponse = (
|
||||
AnthropicConfig._process_response(
|
||||
response=httpx_response,
|
||||
model_response=litellm.ModelResponse(),
|
||||
model=model,
|
||||
stream=False,
|
||||
messages=[],
|
||||
logging_obj=logging_obj,
|
||||
optional_params={},
|
||||
api_key="",
|
||||
data={},
|
||||
print_verbose=litellm.print_verbose,
|
||||
encoding=None,
|
||||
json_mode=False,
|
||||
)
|
||||
)
|
||||
|
||||
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
|
||||
litellm_model_response=litellm_model_response,
|
||||
model=model,
|
||||
kwargs=kwargs,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
await logging_obj.async_success_handler(
|
||||
result=litellm_model_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def _create_anthropic_response_logging_payload(
|
||||
litellm_model_response: Union[
|
||||
litellm.ModelResponse, litellm.TextCompletionResponse
|
||||
],
|
||||
model: str,
|
||||
kwargs: dict,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
):
|
||||
"""
|
||||
Create the standard logging object for Anthropic passthrough
|
||||
|
||||
handles streaming and non-streaming responses
|
||||
"""
|
||||
response_cost = litellm.completion_cost(
|
||||
completion_response=litellm_model_response,
|
||||
model=model,
|
||||
)
|
||||
kwargs["response_cost"] = response_cost
|
||||
kwargs["model"] = model
|
||||
|
||||
# Make standard logging object for Vertex AI
|
||||
standard_logging_object = get_standard_logging_object_payload(
|
||||
kwargs=kwargs,
|
||||
init_response_obj=litellm_model_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=logging_obj,
|
||||
status="success",
|
||||
)
|
||||
|
||||
# pretty print standard logging object
|
||||
verbose_proxy_logger.debug(
|
||||
"standard_logging_object= %s", json.dumps(standard_logging_object, indent=4)
|
||||
)
|
||||
kwargs["standard_logging_object"] = standard_logging_object
|
||||
return kwargs
|
||||
|
||||
@staticmethod
|
||||
async def _handle_logging_anthropic_collected_chunks(
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
passthrough_success_handler_obj: PassThroughEndpointLogging,
|
||||
url_route: str,
|
||||
request_body: dict,
|
||||
endpoint_type: EndpointType,
|
||||
start_time: datetime,
|
||||
all_chunks: List[str],
|
||||
end_time: datetime,
|
||||
):
|
||||
"""
|
||||
Takes raw chunks from Anthropic passthrough endpoint and logs them in litellm callbacks
|
||||
|
||||
- Builds complete response from chunks
|
||||
- Creates standard logging object
|
||||
- Logs in litellm callbacks
|
||||
"""
|
||||
model = request_body.get("model", "")
|
||||
complete_streaming_response = (
|
||||
AnthropicPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=all_chunks,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
model=model,
|
||||
)
|
||||
)
|
||||
if complete_streaming_response is None:
|
||||
verbose_proxy_logger.error(
|
||||
"Unable to build complete streaming response for Anthropic passthrough endpoint, not logging..."
|
||||
)
|
||||
return
|
||||
kwargs = AnthropicPassthroughLoggingHandler._create_anthropic_response_logging_payload(
|
||||
litellm_model_response=complete_streaming_response,
|
||||
model=model,
|
||||
kwargs={},
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
logging_obj=litellm_logging_obj,
|
||||
)
|
||||
await litellm_logging_obj.async_success_handler(
|
||||
result=complete_streaming_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=False,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_complete_streaming_response(
|
||||
all_chunks: List[str],
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
model: str,
|
||||
) -> Optional[Union[litellm.ModelResponse, litellm.TextCompletionResponse]]:
|
||||
"""
|
||||
Builds complete response from raw Anthropic chunks
|
||||
|
||||
- Converts str chunks to generic chunks
|
||||
- Converts generic chunks to litellm chunks (OpenAI format)
|
||||
- Builds complete response from litellm chunks
|
||||
"""
|
||||
anthropic_model_response_iterator = AnthropicModelResponseIterator(
|
||||
streaming_response=None,
|
||||
sync_stream=False,
|
||||
)
|
||||
litellm_custom_stream_wrapper = litellm.CustomStreamWrapper(
|
||||
completion_stream=anthropic_model_response_iterator,
|
||||
model=model,
|
||||
logging_obj=litellm_logging_obj,
|
||||
custom_llm_provider="anthropic",
|
||||
)
|
||||
all_openai_chunks = []
|
||||
for _chunk_str in all_chunks:
|
||||
try:
|
||||
generic_chunk = anthropic_model_response_iterator.convert_str_chunk_to_generic_chunk(
|
||||
chunk=_chunk_str
|
||||
)
|
||||
litellm_chunk = litellm_custom_stream_wrapper.chunk_creator(
|
||||
chunk=generic_chunk
|
||||
)
|
||||
if litellm_chunk is not None:
|
||||
all_openai_chunks.append(litellm_chunk)
|
||||
except (StopIteration, StopAsyncIteration):
|
||||
break
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=all_openai_chunks
|
||||
)
|
||||
return complete_streaming_response
|
||||
|
|
@ -0,0 +1,195 @@
|
|||
import json
|
||||
import re
|
||||
from datetime import datetime
|
||||
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
get_standard_logging_object_payload,
|
||||
)
|
||||
from litellm.llms.vertex_ai_and_google_ai_studio.gemini.vertex_and_google_ai_studio_gemini import (
|
||||
ModelResponseIterator as VertexModelResponseIterator,
|
||||
)
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ..success_handler import PassThroughEndpointLogging
|
||||
from ..types import EndpointType
|
||||
else:
|
||||
PassThroughEndpointLogging = Any
|
||||
EndpointType = Any
|
||||
|
||||
|
||||
class VertexPassthroughLoggingHandler:
|
||||
@staticmethod
|
||||
async def vertex_passthrough_handler(
|
||||
httpx_response: httpx.Response,
|
||||
logging_obj: LiteLLMLoggingObj,
|
||||
url_route: str,
|
||||
result: str,
|
||||
start_time: datetime,
|
||||
end_time: datetime,
|
||||
cache_hit: bool,
|
||||
**kwargs,
|
||||
):
|
||||
if "generateContent" in url_route:
|
||||
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
|
||||
instance_of_vertex_llm = litellm.VertexGeminiConfig()
|
||||
litellm_model_response: litellm.ModelResponse = (
|
||||
instance_of_vertex_llm._transform_response(
|
||||
model=model,
|
||||
messages=[
|
||||
{"role": "user", "content": "no-message-pass-through-endpoint"}
|
||||
],
|
||||
response=httpx_response,
|
||||
model_response=litellm.ModelResponse(),
|
||||
logging_obj=logging_obj,
|
||||
optional_params={},
|
||||
litellm_params={},
|
||||
api_key="",
|
||||
data={},
|
||||
print_verbose=litellm.print_verbose,
|
||||
encoding=None,
|
||||
)
|
||||
)
|
||||
logging_obj.model = litellm_model_response.model or model
|
||||
logging_obj.model_call_details["model"] = logging_obj.model
|
||||
|
||||
await logging_obj.async_success_handler(
|
||||
result=litellm_model_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
**kwargs,
|
||||
)
|
||||
elif "predict" in url_route:
|
||||
from litellm.llms.vertex_ai_and_google_ai_studio.image_generation.image_generation_handler import (
|
||||
VertexImageGeneration,
|
||||
)
|
||||
from litellm.types.utils import PassthroughCallTypes
|
||||
|
||||
vertex_image_generation_class = VertexImageGeneration()
|
||||
|
||||
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
_json_response = httpx_response.json()
|
||||
|
||||
litellm_prediction_response: Union[
|
||||
litellm.ModelResponse, litellm.EmbeddingResponse, litellm.ImageResponse
|
||||
] = litellm.ModelResponse()
|
||||
if vertex_image_generation_class.is_image_generation_response(
|
||||
_json_response
|
||||
):
|
||||
litellm_prediction_response = (
|
||||
vertex_image_generation_class.process_image_generation_response(
|
||||
_json_response,
|
||||
model_response=litellm.ImageResponse(),
|
||||
model=model,
|
||||
)
|
||||
)
|
||||
|
||||
logging_obj.call_type = (
|
||||
PassthroughCallTypes.passthrough_image_generation.value
|
||||
)
|
||||
else:
|
||||
litellm_prediction_response = litellm.vertexAITextEmbeddingConfig.transform_vertex_response_to_openai(
|
||||
response=_json_response,
|
||||
model=model,
|
||||
model_response=litellm.EmbeddingResponse(),
|
||||
)
|
||||
if isinstance(litellm_prediction_response, litellm.EmbeddingResponse):
|
||||
litellm_prediction_response.model = model
|
||||
|
||||
logging_obj.model = model
|
||||
logging_obj.model_call_details["model"] = logging_obj.model
|
||||
|
||||
await logging_obj.async_success_handler(
|
||||
result=litellm_prediction_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=cache_hit,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
async def _handle_logging_vertex_collected_chunks(
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
passthrough_success_handler_obj: PassThroughEndpointLogging,
|
||||
url_route: str,
|
||||
request_body: dict,
|
||||
endpoint_type: EndpointType,
|
||||
start_time: datetime,
|
||||
all_chunks: List[str],
|
||||
end_time: datetime,
|
||||
):
|
||||
"""
|
||||
Takes raw chunks from Vertex passthrough endpoint and logs them in litellm callbacks
|
||||
|
||||
- Builds complete response from chunks
|
||||
- Creates standard logging object
|
||||
- Logs in litellm callbacks
|
||||
"""
|
||||
kwargs: Dict[str, Any] = {}
|
||||
model = VertexPassthroughLoggingHandler.extract_model_from_url(url_route)
|
||||
complete_streaming_response = (
|
||||
VertexPassthroughLoggingHandler._build_complete_streaming_response(
|
||||
all_chunks=all_chunks,
|
||||
litellm_logging_obj=litellm_logging_obj,
|
||||
model=model,
|
||||
)
|
||||
)
|
||||
|
||||
if complete_streaming_response is None:
|
||||
verbose_proxy_logger.error(
|
||||
"Unable to build complete streaming response for Vertex passthrough endpoint, not logging..."
|
||||
)
|
||||
return
|
||||
await litellm_logging_obj.async_success_handler(
|
||||
result=complete_streaming_response,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
cache_hit=False,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _build_complete_streaming_response(
|
||||
all_chunks: List[str],
|
||||
litellm_logging_obj: LiteLLMLoggingObj,
|
||||
model: str,
|
||||
) -> Optional[Union[litellm.ModelResponse, litellm.TextCompletionResponse]]:
|
||||
vertex_iterator = VertexModelResponseIterator(
|
||||
streaming_response=None,
|
||||
sync_stream=False,
|
||||
)
|
||||
litellm_custom_stream_wrapper = litellm.CustomStreamWrapper(
|
||||
completion_stream=vertex_iterator,
|
||||
model=model,
|
||||
logging_obj=litellm_logging_obj,
|
||||
custom_llm_provider="vertex_ai",
|
||||
)
|
||||
all_openai_chunks = []
|
||||
for chunk in all_chunks:
|
||||
generic_chunk = vertex_iterator._common_chunk_parsing_logic(chunk)
|
||||
litellm_chunk = litellm_custom_stream_wrapper.chunk_creator(
|
||||
chunk=generic_chunk
|
||||
)
|
||||
if litellm_chunk is not None:
|
||||
all_openai_chunks.append(litellm_chunk)
|
||||
|
||||
complete_streaming_response = litellm.stream_chunk_builder(
|
||||
chunks=all_openai_chunks
|
||||
)
|
||||
|
||||
return complete_streaming_response
|
||||
|
||||
@staticmethod
|
||||
def extract_model_from_url(url: str) -> str:
|
||||
pattern = r"/models/([^:]+)"
|
||||
match = re.search(pattern, url)
|
||||
if match:
|
||||
return match.group(1)
|
||||
return "unknown"
|
||||
Loading…
Add table
Add a link
Reference in a new issue