mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
b8af46e1a2
* 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
110 lines
3.5 KiB
Python
110 lines
3.5 KiB
Python
import json
|
|
import re
|
|
import threading
|
|
from datetime import datetime
|
|
from typing import 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 (
|
|
VertexLLM,
|
|
)
|
|
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
|
from litellm.types.utils import StandardPassThroughResponseObject
|
|
|
|
from .llm_provider_handlers.anthropic_passthrough_logging_handler import (
|
|
AnthropicPassthroughLoggingHandler,
|
|
)
|
|
from .llm_provider_handlers.vertex_passthrough_logging_handler import (
|
|
VertexPassthroughLoggingHandler,
|
|
)
|
|
|
|
|
|
class PassThroughEndpointLogging:
|
|
def __init__(self):
|
|
self.TRACKED_VERTEX_ROUTES = [
|
|
"generateContent",
|
|
"streamGenerateContent",
|
|
"predict",
|
|
]
|
|
|
|
# Anthropic
|
|
self.TRACKED_ANTHROPIC_ROUTES = ["/messages"]
|
|
|
|
async def pass_through_async_success_handler(
|
|
self,
|
|
httpx_response: httpx.Response,
|
|
response_body: Optional[dict],
|
|
logging_obj: LiteLLMLoggingObj,
|
|
url_route: str,
|
|
result: str,
|
|
start_time: datetime,
|
|
end_time: datetime,
|
|
cache_hit: bool,
|
|
**kwargs,
|
|
):
|
|
if self.is_vertex_route(url_route):
|
|
await VertexPassthroughLoggingHandler.vertex_passthrough_handler(
|
|
httpx_response=httpx_response,
|
|
logging_obj=logging_obj,
|
|
url_route=url_route,
|
|
result=result,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
**kwargs,
|
|
)
|
|
elif self.is_anthropic_route(url_route):
|
|
await AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler(
|
|
httpx_response=httpx_response,
|
|
response_body=response_body or {},
|
|
logging_obj=logging_obj,
|
|
url_route=url_route,
|
|
result=result,
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=cache_hit,
|
|
**kwargs,
|
|
)
|
|
else:
|
|
standard_logging_response_object = StandardPassThroughResponseObject(
|
|
response=httpx_response.text
|
|
)
|
|
threading.Thread(
|
|
target=logging_obj.success_handler,
|
|
args=(
|
|
standard_logging_response_object,
|
|
start_time,
|
|
end_time,
|
|
cache_hit,
|
|
),
|
|
).start()
|
|
await logging_obj.async_success_handler(
|
|
result=(
|
|
json.dumps(result)
|
|
if isinstance(result, dict)
|
|
else standard_logging_response_object
|
|
),
|
|
start_time=start_time,
|
|
end_time=end_time,
|
|
cache_hit=False,
|
|
**kwargs,
|
|
)
|
|
|
|
def is_vertex_route(self, url_route: str):
|
|
for route in self.TRACKED_VERTEX_ROUTES:
|
|
if route in url_route:
|
|
return True
|
|
return False
|
|
|
|
def is_anthropic_route(self, url_route: str):
|
|
for route in self.TRACKED_ANTHROPIC_ROUTES:
|
|
if route in url_route:
|
|
return True
|
|
return False
|