add cost tracking for rerank+ test

This commit is contained in:
Ishaan Jaff 2024-09-06 16:06:19 -07:00
parent d1342c5991
commit b659095f71
2 changed files with 12 additions and 1 deletions

View file

@ -30,6 +30,7 @@ from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.redact_messages import (
redact_message_input_output_from_logging,
)
from litellm.rerank_api.types import RerankResponse
from litellm.types.llms.openai import HttpxBinaryResponseContent
from litellm.types.router import SPECIAL_MODEL_INFO_PARAMS
from litellm.types.utils import (
@ -524,6 +525,7 @@ class Logging:
TranscriptionResponse,
TextCompletionResponse,
HttpxBinaryResponseContent,
RerankResponse,
],
cache_hit: Optional[bool] = None,
):
@ -585,6 +587,7 @@ class Logging:
or isinstance(result, TranscriptionResponse)
or isinstance(result, TextCompletionResponse)
or isinstance(result, HttpxBinaryResponseContent) # tts
or isinstance(result, RerankResponse)
):
## RESPONSE COST ##
self.model_call_details["response_cost"] = (