diff --git a/.circleci/config.yml b/.circleci/config.yml index 39925ab3e0fe..1116ad74101f 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -199,7 +199,7 @@ jobs: command: | pwd ls - python -m pytest tests/local_testing --cov=litellm --cov-report=xml -vv -k "router" -x -s -v --junitxml=test-results/junit.xml --durations=5 + python -m pytest tests/local_testing tests/router_unit_tests --cov=litellm --cov-report=xml -vv -k "router" -x -s -v --junitxml=test-results/junit.xml --durations=5 no_output_timeout: 120m - run: name: Rename the coverage files @@ -328,6 +328,48 @@ jobs: paths: - llm_translation_coverage.xml - llm_translation_coverage + logging_testing: + docker: + - image: cimg/python:3.11 + auth: + username: ${DOCKERHUB_USERNAME} + password: ${DOCKERHUB_PASSWORD} + working_directory: ~/project + + steps: + - checkout + - run: + name: Install Dependencies + command: | + python -m pip install --upgrade pip + python -m pip install -r requirements.txt + pip install "pytest==7.3.1" + pip install "pytest-retry==1.6.3" + pip install "pytest-cov==5.0.0" + pip install "pytest-asyncio==0.21.1" + pip install "respx==0.21.1" + # Run pytest and generate JUnit XML report + - run: + name: Run tests + command: | + pwd + ls + python -m pytest -vv tests/logging_callback_tests --cov=litellm --cov-report=xml -x -s -v --junitxml=test-results/junit.xml --durations=5 + no_output_timeout: 120m + - run: + name: Rename the coverage files + command: | + mv coverage.xml logging_coverage.xml + mv .coverage logging_coverage + + # Store test results + - store_test_results: + path: test-results + - persist_to_workspace: + root: . + paths: + - logging_coverage.xml + - logging_coverage installing_litellm_on_python: docker: - image: circleci/python:3.8 @@ -380,6 +422,7 @@ jobs: - run: python -c "from litellm import *" || (echo '🚨 import failed, this means you introduced unprotected imports! 🚨'; exit 1) - run: ruff check ./litellm - run: python ./tests/documentation_tests/test_general_setting_keys.py + - run: python ./tests/code_coverage_tests/router_code_coverage.py - run: python ./tests/documentation_tests/test_env_keys.py db_migration_disable_update_check: @@ -768,7 +811,7 @@ jobs: python -m venv venv . venv/bin/activate pip install coverage - coverage combine llm_translation_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage ui_endpoint_testing_coverage + coverage combine llm_translation_coverage logging_coverage litellm_router_coverage local_testing_coverage litellm_assistants_api_coverage ui_endpoint_testing_coverage coverage xml - codecov/upload: file: ./coverage.xml @@ -1004,9 +1047,16 @@ workflows: only: - main - /litellm_.*/ + - logging_testing: + filters: + branches: + only: + - main + - /litellm_.*/ - upload-coverage: requires: - llm_translation_testing + - logging_testing - litellm_router_testing - local_testing - litellm_assistants_api_testing @@ -1035,6 +1085,7 @@ workflows: - build_and_test - load_testing - llm_translation_testing + - logging_testing - litellm_router_testing - litellm_assistants_api_testing - ui_endpoint_testing diff --git a/litellm/caching/caching_handler.py b/litellm/caching/caching_handler.py index 264fb405baa9..479b3bd1f8a2 100644 --- a/litellm/caching/caching_handler.py +++ b/litellm/caching/caching_handler.py @@ -26,6 +26,9 @@ RedisSemanticCache, S3Cache, ) +from litellm.litellm_core_utils.logging_utils import ( + _assemble_complete_response_from_streaming_chunks, +) from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, @@ -517,28 +520,14 @@ async def _add_streaming_response_to_cache(self, processed_chunk: ModelResponse) """ complete_streaming_response: Optional[ Union[ModelResponse, TextCompletionResponse] - ] = None - if ( - processed_chunk.choices[0].finish_reason is not None - ): # if it's the last chunk - self.async_streaming_chunks.append(processed_chunk) - try: - end_time: datetime.datetime = datetime.datetime.now() - complete_streaming_response = litellm.stream_chunk_builder( - self.async_streaming_chunks, - messages=self.request_kwargs.get("messages", None), - start_time=self.start_time, - end_time=end_time, - ) - except Exception as e: - verbose_logger.exception( - "Error occurred building stream chunk in success logging: {}".format( - str(e) - ) - ) - complete_streaming_response = None - else: - self.async_streaming_chunks.append(processed_chunk) + ] = _assemble_complete_response_from_streaming_chunks( + result=processed_chunk, + start_time=self.start_time, + end_time=datetime.datetime.now(), + request_kwargs=self.request_kwargs, + streaming_chunks=self.async_streaming_chunks, + is_async=True, + ) # if a complete_streaming_response is assembled, add it to the cache if complete_streaming_response is not None: diff --git a/litellm/integrations/gcs_bucket/gcs_bucket.py b/litellm/integrations/gcs_bucket/gcs_bucket.py index 744361453009..3d99c025711c 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket.py @@ -12,9 +12,6 @@ from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.gcs_bucket.gcs_bucket_base import GCSBucketBase -from litellm.litellm_core_utils.logging_utils import ( - convert_litellm_response_object_to_dict, -) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.proxy._types import CommonProxyErrors, SpendLogsMetadata, SpendLogsPayload from litellm.types.utils import ( diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_base.py b/litellm/integrations/gcs_bucket/gcs_bucket_base.py index b9e0930be7df..dac0790b65bd 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_base.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_base.py @@ -10,9 +10,6 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.logging_utils import ( - convert_litellm_response_object_to_dict, -) from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, httpxSpecialProvider, diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index fada59ab2cce..f767ea64e042 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -7,8 +7,8 @@ import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.types.services import ServiceLoggerPayload +from litellm.types.utils import StandardLoggingPayload if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -152,9 +152,15 @@ async def async_service_success_hook( context=trace.set_span_in_context(parent_otel_span), start_time=_start_time_ns, ) - service_logging_span.set_attribute(key="call_type", value=payload.call_type) - service_logging_span.set_attribute( - key="service", value=payload.service.value + self.safe_set_attribute( + span=service_logging_span, + key="call_type", + value=payload.call_type, + ) + self.safe_set_attribute( + span=service_logging_span, + key="service", + value=payload.service.value, ) if event_metadata: @@ -166,7 +172,11 @@ async def async_service_success_hook( value = str(value) except Exception: value = "litllm logging error - could_not_json_serialize" - service_logging_span.set_attribute(key, value) + self.safe_set_attribute( + span=service_logging_span, + key=key, + value=value, + ) service_logging_span.set_status(Status(StatusCode.OK)) service_logging_span.end(end_time=_end_time_ns) @@ -204,12 +214,22 @@ async def async_service_failure_hook( context=trace.set_span_in_context(parent_otel_span), start_time=_start_time_ns, ) - service_logging_span.set_attribute(key="call_type", value=payload.call_type) - service_logging_span.set_attribute( - key="service", value=payload.service.value + self.safe_set_attribute( + span=service_logging_span, + key="call_type", + value=payload.call_type, + ) + self.safe_set_attribute( + span=service_logging_span, + key="service", + value=payload.service.value, ) if error: - service_logging_span.set_attribute(key="error", value=error) + self.safe_set_attribute( + span=service_logging_span, + key="error", + value=error, + ) if event_metadata: for key, value in event_metadata.items(): if isinstance(value, dict): @@ -217,7 +237,11 @@ async def async_service_failure_hook( value = str(value) except Exception: value = "litllm logging error - could_not_json_serialize" - service_logging_span.set_attribute(key, value) + self.safe_set_attribute( + span=service_logging_span, + key=key, + value=value, + ) service_logging_span.set_status(Status(StatusCode.ERROR)) service_logging_span.end(end_time=_end_time_ns) @@ -241,8 +265,10 @@ async def async_post_call_failure_hook( name=_span_name, context=trace.set_span_in_context(parent_otel_span), ) - exception_logging_span.set_attribute( - key="exception", value=str(original_exception) + self.safe_set_attribute( + span=exception_logging_span, + key="exception", + value=str(original_exception), ) exception_logging_span.set_status(Status(StatusCode.ERROR)) exception_logging_span.end(end_time=self._to_ns(datetime.now())) @@ -329,10 +355,20 @@ def set_tools_attributes(self, span: Span, tools): continue prefix = f"{SpanAttributes.LLM_REQUEST_FUNCTIONS}.{i}" - span.set_attribute(f"{prefix}.name", function.get("name")) - span.set_attribute(f"{prefix}.description", function.get("description")) - span.set_attribute( - f"{prefix}.parameters", json.dumps(function.get("parameters")) + self.safe_set_attribute( + span=span, + key=f"{prefix}.name", + value=function.get("name"), + ) + self.safe_set_attribute( + span=span, + key=f"{prefix}.description", + value=function.get("description"), + ) + self.safe_set_attribute( + span=span, + key=f"{prefix}.parameters", + value=json.dumps(function.get("parameters")), ) except Exception as e: verbose_logger.error( @@ -340,10 +376,22 @@ def set_tools_attributes(self, span: Span, tools): ) pass - def is_primitive(self, value): + def cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: + """ + Casts the value to a primitive OTEL type if it is not already a primitive type. + + OTEL supports - str, bool, int, float + + If it's not a primitive type, then it's converted to a string + """ if value is None: - return False - return isinstance(value, (str, bool, int, float)) + return "" + if isinstance(value, (str, bool, int, float)): + return value + try: + return str(value) + except Exception: + return "" def set_attributes(self, span: Span, kwargs, response_obj): try: @@ -363,19 +411,22 @@ def set_attributes(self, span: Span, kwargs, response_obj): optional_params = kwargs.get("optional_params", {}) litellm_params = kwargs.get("litellm_params", {}) or {} + standard_logging_payload: Optional[StandardLoggingPayload] = kwargs.get( + "standard_logging_object" + ) + if standard_logging_payload is None: + raise ValueError("standard_logging_object not found in kwargs") # https://github.com/open-telemetry/semantic-conventions/blob/main/model/registry/gen-ai.yaml # Following Conventions here: https://github.com/open-telemetry/semantic-conventions/blob/main/docs/gen-ai/llm-spans.md ############################################# ############ LLM CALL METADATA ############## ############################################# - metadata = litellm_params.get("metadata", {}) or {} - - clean_metadata = redact_user_api_key_info(metadata=metadata) - - for key, value in clean_metadata.items(): - if self.is_primitive(value): - span.set_attribute("metadata.{}".format(key), value) + metadata = standard_logging_payload["metadata"] + for key, value in metadata.items(): + self.safe_set_attribute( + span=span, key="metadata.{}".format(key), value=value + ) ############################################# ########## LLM Request Attributes ########### @@ -383,71 +434,90 @@ def set_attributes(self, span: Span, kwargs, response_obj): # The name of the LLM a request is being made to if kwargs.get("model"): - span.set_attribute( - SpanAttributes.LLM_REQUEST_MODEL, kwargs.get("model") + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_REQUEST_MODEL, + value=kwargs.get("model"), ) # The Generative AI Provider: Azure, OpenAI, etc. - span.set_attribute( - SpanAttributes.LLM_SYSTEM, - litellm_params.get("custom_llm_provider", "Unknown"), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_SYSTEM, + value=litellm_params.get("custom_llm_provider", "Unknown"), ) # The maximum number of tokens the LLM generates for a request. if optional_params.get("max_tokens"): - span.set_attribute( - SpanAttributes.LLM_REQUEST_MAX_TOKENS, - optional_params.get("max_tokens"), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_REQUEST_MAX_TOKENS, + value=optional_params.get("max_tokens"), ) # The temperature setting for the LLM request. if optional_params.get("temperature"): - span.set_attribute( - SpanAttributes.LLM_REQUEST_TEMPERATURE, - optional_params.get("temperature"), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_REQUEST_TEMPERATURE, + value=optional_params.get("temperature"), ) # The top_p sampling setting for the LLM request. if optional_params.get("top_p"): - span.set_attribute( - SpanAttributes.LLM_REQUEST_TOP_P, optional_params.get("top_p") + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_REQUEST_TOP_P, + value=optional_params.get("top_p"), ) - span.set_attribute( - SpanAttributes.LLM_IS_STREAMING, - str(optional_params.get("stream", False)), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_IS_STREAMING, + value=str(optional_params.get("stream", False)), ) if optional_params.get("user"): - span.set_attribute(SpanAttributes.LLM_USER, optional_params.get("user")) + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_USER, + value=optional_params.get("user"), + ) # The unique identifier for the completion. if response_obj.get("id"): - span.set_attribute("gen_ai.response.id", response_obj.get("id")) + self.safe_set_attribute( + span=span, key="gen_ai.response.id", value=response_obj.get("id") + ) # The model used to generate the response. if response_obj.get("model"): - span.set_attribute( - SpanAttributes.LLM_RESPONSE_MODEL, response_obj.get("model") + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_RESPONSE_MODEL, + value=response_obj.get("model"), ) usage = response_obj.get("usage") if usage: - span.set_attribute( - SpanAttributes.LLM_USAGE_TOTAL_TOKENS, - usage.get("total_tokens"), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_USAGE_TOTAL_TOKENS, + value=usage.get("total_tokens"), ) # The number of tokens used in the LLM response (completion). - span.set_attribute( - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, - usage.get("completion_tokens"), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, + value=usage.get("completion_tokens"), ) # The number of tokens used in the LLM prompt. - span.set_attribute( - SpanAttributes.LLM_USAGE_PROMPT_TOKENS, - usage.get("prompt_tokens"), + self.safe_set_attribute( + span=span, + key=SpanAttributes.LLM_USAGE_PROMPT_TOKENS, + value=usage.get("prompt_tokens"), ) ######################################################################## @@ -466,17 +536,19 @@ def set_attributes(self, span: Span, kwargs, response_obj): if kwargs.get("messages"): for idx, prompt in enumerate(kwargs.get("messages")): if prompt.get("role"): - span.set_attribute( - f"{SpanAttributes.LLM_PROMPTS}.{idx}.role", - prompt.get("role"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_PROMPTS}.{idx}.role", + value=prompt.get("role"), ) if prompt.get("content"): if not isinstance(prompt.get("content"), str): prompt["content"] = str(prompt.get("content")) - span.set_attribute( - f"{SpanAttributes.LLM_PROMPTS}.{idx}.content", - prompt.get("content"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_PROMPTS}.{idx}.content", + value=prompt.get("content"), ) ############################################# ########## LLM Response Attributes ########## @@ -485,15 +557,17 @@ def set_attributes(self, span: Span, kwargs, response_obj): if response_obj.get("choices"): for idx, choice in enumerate(response_obj.get("choices")): if choice.get("finish_reason"): - span.set_attribute( - f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.finish_reason", - choice.get("finish_reason"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.finish_reason", + value=choice.get("finish_reason"), ) if choice.get("message"): if choice.get("message").get("role"): - span.set_attribute( - f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.role", - choice.get("message").get("role"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.role", + value=choice.get("message").get("role"), ) if choice.get("message").get("content"): if not isinstance( @@ -502,21 +576,26 @@ def set_attributes(self, span: Span, kwargs, response_obj): choice["message"]["content"] = str( choice.get("message").get("content") ) - span.set_attribute( - f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.content", - choice.get("message").get("content"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.content", + value=choice.get("message").get("content"), ) message = choice.get("message") tool_calls = message.get("tool_calls") if tool_calls: - span.set_attribute( - f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.name", - tool_calls[0].get("function").get("name"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.name", + value=tool_calls[0].get("function").get("name"), ) - span.set_attribute( - f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.arguments", - tool_calls[0].get("function").get("arguments"), + self.safe_set_attribute( + span=span, + key=f"{SpanAttributes.LLM_COMPLETIONS}.{idx}.function_call.arguments", + value=tool_calls[0] + .get("function") + .get("arguments"), ) except Exception as e: @@ -524,6 +603,30 @@ def set_attributes(self, span: Span, kwargs, response_obj): "OpenTelemetry logging error in set_attributes %s", str(e) ) + def _cast_as_primitive_value_type(self, value) -> Union[str, bool, int, float]: + """ + Casts the value to a primitive OTEL type if it is not already a primitive type. + + OTEL supports - str, bool, int, float + + If it's not a primitive type, then it's converted to a string + """ + if value is None: + return "" + if isinstance(value, (str, bool, int, float)): + return value + try: + return str(value) + except Exception: + return "" + + def safe_set_attribute(self, span: Span, key: str, value: Any): + """ + Safely sets an attribute on the span, ensuring the value is a primitive type. + """ + primitive_value = self._cast_as_primitive_value_type(value) + span.set_attribute(key, primitive_value) + def set_raw_request_attributes(self, span: Span, kwargs, response_obj): from litellm.proxy._types import SpanAttributes @@ -541,11 +644,8 @@ def set_raw_request_attributes(self, span: Span, kwargs, response_obj): # OTEL Attributes for the RAW Request to https://docs.anthropic.com/en/api/messages if complete_input_dict and isinstance(complete_input_dict, dict): for param, val in complete_input_dict.items(): - if not isinstance(val, str): - val = str(val) - span.set_attribute( - f"llm.{custom_llm_provider}.{param}", - val, + self.safe_set_attribute( + span=span, key=f"llm.{custom_llm_provider}.{param}", value=val ) ############################################# @@ -558,11 +658,10 @@ def set_raw_request_attributes(self, span: Span, kwargs, response_obj): try: _raw_response = json.loads(_raw_response) for param, val in _raw_response.items(): - if not isinstance(val, str): - val = str(val) - span.set_attribute( - f"llm.{custom_llm_provider}.{param}", - val, + self.safe_set_attribute( + span=span, + key=f"llm.{custom_llm_provider}.{param}", + value=val, ) except json.JSONDecodeError: verbose_logger.debug( @@ -570,9 +669,11 @@ def set_raw_request_attributes(self, span: Span, kwargs, response_obj): _raw_response ) ) - span.set_attribute( - f"llm.{custom_llm_provider}.stringified_raw_response", - _raw_response, + + self.safe_set_attribute( + span=span, + key=f"llm.{custom_llm_provider}.stringified_raw_response", + value=_raw_response, ) def _to_ns(self, dt): @@ -727,12 +828,21 @@ async def async_management_endpoint_success_hook( _request_data = logging_payload.request_data if _request_data is not None: for key, value in _request_data.items(): - management_endpoint_span.set_attribute(f"request.{key}", value) + self.safe_set_attribute( + span=management_endpoint_span, + key=f"request.{key}", + value=value, + ) _response = logging_payload.response if _response is not None: for key, value in _response.items(): - management_endpoint_span.set_attribute(f"response.{key}", value) + self.safe_set_attribute( + span=management_endpoint_span, + key=f"response.{key}", + value=value, + ) + management_endpoint_span.set_status(Status(StatusCode.OK)) management_endpoint_span.end(end_time=_end_time_ns) @@ -773,9 +883,17 @@ async def async_management_endpoint_failure_hook( _request_data = logging_payload.request_data if _request_data is not None: for key, value in _request_data.items(): - management_endpoint_span.set_attribute(f"request.{key}", value) + self.safe_set_attribute( + span=management_endpoint_span, + key=f"request.{key}", + value=value, + ) _exception = logging_payload.exception - management_endpoint_span.set_attribute("exception", str(_exception)) + self.safe_set_attribute( + span=management_endpoint_span, + key="exception", + value=str(_exception), + ) management_endpoint_span.set_status(Status(StatusCode.ERROR)) management_endpoint_span.end(end_time=_end_time_ns) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 1d4a3e3cc550..f7afbb7b8150 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -86,6 +86,7 @@ from ..integrations.traceloop import TraceloopLogger from ..integrations.weights_biases import WeightsBiasesLogger from .exception_mapping_utils import _get_response_headers +from .logging_utils import _assemble_complete_response_from_streaming_chunks try: from ..proxy.enterprise.enterprise_callbacks.generic_api_callback import ( @@ -878,32 +879,24 @@ def success_handler( # print(f"original response in success handler: {self.model_call_details['original_response']}") try: verbose_logger.debug(f"success callbacks: {litellm.success_callback}") + ## BUILD COMPLETE STREAMED RESPONSE - complete_streaming_response = None + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = None if "complete_streaming_response" in self.model_call_details: return # break out of this. - if self.stream and isinstance(result, ModelResponse): - if ( - result.choices[0].finish_reason is not None - ): # if it's the last chunk - self.sync_streaming_chunks.append(result) - # print_verbose(f"final set of received chunks: {self.sync_streaming_chunks}") - try: - complete_streaming_response = litellm.stream_chunk_builder( - self.sync_streaming_chunks, - messages=self.model_call_details.get("messages", None), - start_time=start_time, - end_time=end_time, - ) - except Exception as e: - verbose_logger.exception( - "LiteLLM.LoggingError: [Non-Blocking] Exception occurred while building complete streaming response in success logging {}".format( - str(e) - ) - ) - complete_streaming_response = None - else: - self.sync_streaming_chunks.append(result) + if self.stream: + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = _assemble_complete_response_from_streaming_chunks( + result=result, + start_time=start_time, + end_time=end_time, + request_kwargs=self.model_call_details, + streaming_chunks=self.sync_streaming_chunks, + is_async=False, + ) _caching_complete_streaming_response: Optional[ Union[ModelResponse, TextCompletionResponse] ] = None @@ -1495,29 +1488,23 @@ async def async_success_handler( start_time=start_time, end_time=end_time, result=result, cache_hit=cache_hit ) ## BUILD COMPLETE STREAMED RESPONSE - complete_streaming_response = None if "async_complete_streaming_response" in self.model_call_details: return # break out of this. - if self.stream: - if result.choices[0].finish_reason is not None: # if it's the last chunk - self.streaming_chunks.append(result) - # verbose_logger.debug(f"final set of received chunks: {self.streaming_chunks}") - try: - complete_streaming_response = litellm.stream_chunk_builder( - self.streaming_chunks, - messages=self.model_call_details.get("messages", None), - start_time=start_time, - end_time=end_time, - ) - except Exception as e: - verbose_logger.exception( - "Error occurred building stream chunk in success logging: {}".format( - str(e) - ) - ) - complete_streaming_response = None - else: - self.streaming_chunks.append(result) + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = None + if self.stream is True: + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = _assemble_complete_response_from_streaming_chunks( + result=result, + start_time=start_time, + end_time=end_time, + request_kwargs=self.model_call_details, + streaming_chunks=self.streaming_chunks, + is_async=True, + ) + if complete_streaming_response is not None: print_verbose("Async success callbacks: Got a complete streaming response") diff --git a/litellm/litellm_core_utils/logging_utils.py b/litellm/litellm_core_utils/logging_utils.py index 7fa1be9d8b60..0a9fc8342a4a 100644 --- a/litellm/litellm_core_utils/logging_utils.py +++ b/litellm/litellm_core_utils/logging_utils.py @@ -1,4 +1,8 @@ -from typing import TYPE_CHECKING, Any, Optional, Union +from datetime import datetime +from typing import TYPE_CHECKING, Any, List, Optional, Union + +from litellm._logging import verbose_logger +from litellm.types.utils import ModelResponse, TextCompletionResponse if TYPE_CHECKING: from litellm import ModelResponse as _ModelResponse @@ -15,21 +19,6 @@ """ -def convert_litellm_response_object_to_dict(response_obj: Any) -> dict: - """ - Convert a LiteLLM response object to a dictionary - - """ - if isinstance(response_obj, dict): - return response_obj - for _type in litellm.ALL_LITELLM_RESPONSE_TYPES: - if isinstance(response_obj, _type): - return response_obj.model_dump() - - # If it's not a LiteLLM type, return the object as is - return dict(response_obj) - - def convert_litellm_response_object_to_str( response_obj: Union[Any, LiteLLMModelResponse] ) -> Optional[str]: @@ -46,3 +35,55 @@ def convert_litellm_response_object_to_str( return response_str return None + + +def _assemble_complete_response_from_streaming_chunks( + result: Union[ModelResponse, TextCompletionResponse], + start_time: datetime, + end_time: datetime, + request_kwargs: dict, + streaming_chunks: List[Any], + is_async: bool, +): + """ + Assemble a complete response from a streaming chunks + + - assemble a complete streaming response if result.choices[0].finish_reason is not None + - else append the chunk to the streaming_chunks + + + Args: + result: ModelResponse + start_time: datetime + end_time: datetime + request_kwargs: dict + streaming_chunks: List[Any] + is_async: bool + + Returns: + Optional[Union[ModelResponse, TextCompletionResponse]]: Complete streaming response + + """ + complete_streaming_response: Optional[ + Union[ModelResponse, TextCompletionResponse] + ] = None + if result.choices[0].finish_reason is not None: # if it's the last chunk + streaming_chunks.append(result) + try: + complete_streaming_response = litellm.stream_chunk_builder( + chunks=streaming_chunks, + messages=request_kwargs.get("messages", None), + start_time=start_time, + end_time=end_time, + ) + except Exception as e: + log_message = ( + "Error occurred building stream chunk in {} success logging: {}".format( + "async" if is_async else "sync", str(e) + ) + ) + verbose_logger.exception(log_message) + complete_streaming_response = None + else: + streaming_chunks.append(result) + return complete_streaming_response diff --git a/litellm/litellm_core_utils/mock_functions.py b/litellm/litellm_core_utils/mock_functions.py new file mode 100644 index 000000000000..76425651aa19 --- /dev/null +++ b/litellm/litellm_core_utils/mock_functions.py @@ -0,0 +1,28 @@ +from typing import List, Optional + +from ..types.utils import ( + Categories, + CategoryAppliedInputTypes, + CategoryScores, + Embedding, + EmbeddingResponse, + ImageObject, + ImageResponse, + Moderation, + ModerationCreateResponse, +) + + +def mock_embedding(model: str, mock_response: Optional[List[float]]): + if mock_response is None: + mock_response = [0.0] * 1536 + return EmbeddingResponse( + model=model, + data=[Embedding(embedding=mock_response, index=0, object="embedding")], + ) + + +def mock_image_generation(model: str, mock_response: str): + return ImageResponse( + data=[ImageObject(url=mock_response)], + ) diff --git a/litellm/main.py b/litellm/main.py index fe1453836d95..c1cf3ae2330a 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -42,6 +42,10 @@ ) from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.mock_functions import ( + mock_embedding, + mock_image_generation, +) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.secret_managers.main import get_secret_str from litellm.utils import ( @@ -3163,6 +3167,7 @@ def embedding( tpm = kwargs.pop("tpm", None) litellm_logging_obj: LiteLLMLoggingObj = kwargs.get("litellm_logging_obj") # type: ignore cooldown_time = kwargs.get("cooldown_time", None) + mock_response: Optional[List[float]] = kwargs.get("mock_response", None) # type: ignore max_parallel_requests = kwargs.pop("max_parallel_requests", None) model_info = kwargs.get("model_info", None) metadata = kwargs.get("metadata", None) @@ -3268,6 +3273,9 @@ def embedding( custom_llm_provider=custom_llm_provider, **non_default_params, ) + + if mock_response is not None: + return mock_embedding(model=model, mock_response=mock_response) ### REGISTER CUSTOM MODEL PRICING -- IF GIVEN ### if input_cost_per_token is not None and output_cost_per_token is not None: litellm.register_model( @@ -4377,6 +4385,7 @@ def image_generation( aimg_generation = kwargs.get("aimg_generation", False) litellm_call_id = kwargs.get("litellm_call_id", None) logger_fn = kwargs.get("logger_fn", None) + mock_response: Optional[str] = kwargs.get("mock_response", None) # type: ignore proxy_server_request = kwargs.get("proxy_server_request", None) model_info = kwargs.get("model_info", None) metadata = kwargs.get("metadata", {}) @@ -4486,6 +4495,8 @@ def image_generation( }, custom_llm_provider=custom_llm_provider, ) + if mock_response is not None: + return mock_image_generation(model=model, mock_response=mock_response) if custom_llm_provider == "azure": # azure configs diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index efc3542b198d..6732c631c7d2 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -777,7 +777,7 @@ class NewCustomerRequest(LiteLLMBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[Literal["eu"]] = ( + allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] = ( None # require all user requests to use models in this specific region ) default_model: Optional[str] = ( @@ -804,7 +804,7 @@ class UpdateCustomerRequest(LiteLLMBase): blocked: bool = False # allow/disallow requests for this end-user max_budget: Optional[float] = None budget_id: Optional[str] = None # give either a budget_id or max_budget - allowed_model_region: Optional[Literal["eu"]] = ( + allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] = ( None # require all user requests to use models in this specific region ) default_model: Optional[str] = ( @@ -1384,7 +1384,7 @@ class UserAPIKeyAuth( api_key: Optional[str] = None user_role: Optional[LitellmUserRoles] = None - allowed_model_region: Optional[Literal["eu"]] = None + allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] = None parent_otel_span: Optional[Span] = None rpm_limit_per_model: Optional[Dict[str, int]] = None tpm_limit_per_model: Optional[Dict[str, int]] = None @@ -1466,7 +1466,7 @@ class LiteLLM_EndUserTable(LiteLLMBase): blocked: bool alias: Optional[str] = None spend: float = 0.0 - allowed_model_region: Optional[Literal["eu"]] = None + allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] = None default_model: Optional[str] = None litellm_budget_table: Optional[LiteLLM_BudgetTable] = None diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index c71e3fd40eb5..50d1a91bf72c 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/litellm/proxy/proxy_config.yaml @@ -6,12 +6,5 @@ model_list: api_base: https://exampleopenaiendpoint-production.up.railway.app/ litellm_settings: - callbacks: ["prometheus"] - turn_off_message_logging: true - s3_callback_params: - s3_bucket_name: load-testing-oct # AWS Bucket Name for S3 - s3_region_name: us-west-2 # AWS Region Name for S3 - s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID # us os.environ/ to pass environment variables. This is AWS Access Key ID for S3 - s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY # AWS Secret Access Key for S3 - + callbacks: ["otel"] diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 58a4ff346d60..a9fbaff70883 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5894,7 +5894,7 @@ async def new_end_user( - blocked: bool - Flag to allow or disallow requests for this end-user. Default is False. - max_budget: Optional[float] - The maximum budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. - budget_id: Optional[str] - The identifier for an existing budget allocated to the user. Either 'max_budget' or 'budget_id' should be provided, not both. - - allowed_model_region: Optional[Literal["eu"]] - Require all user requests to use models in this specific region. + - allowed_model_region: Optional[Union[Literal["eu"], Literal["us"]]] - Require all user requests to use models in this specific region. - default_model: Optional[str] - If no equivalent model in the allowed region, default all requests to this model. - metadata: Optional[dict] = Metadata for customer, store information for customer. Example metadata = {"data_training_opt_out": True} diff --git a/litellm/router.py b/litellm/router.py index c31536bd62f9..26730402fa23 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -111,6 +111,7 @@ RouterModelGroupAliasItem, RouterRateLimitError, RouterRateLimitErrorBasic, + RoutingStrategy, updateDeployment, updateLiteLLMParams, ) @@ -519,6 +520,9 @@ def __init__( self._initialize_alerting() def validate_fallbacks(self, fallback_param: Optional[List]): + """ + Validate the fallbacks parameter. + """ if fallback_param is None: return @@ -530,8 +534,13 @@ def validate_fallbacks(self, fallback_param: Optional[List]): f"Dictionary '{fallback_dict}' must have exactly one key, but has {len(fallback_dict)} keys." ) - def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: dict): - if routing_strategy == "least-busy": + def routing_strategy_init( + self, routing_strategy: Union[RoutingStrategy, str], routing_strategy_args: dict + ): + if ( + routing_strategy == RoutingStrategy.LEAST_BUSY.value + or routing_strategy == RoutingStrategy.LEAST_BUSY + ): self.leastbusy_logger = LeastBusyLoggingHandler( router_cache=self.cache, model_list=self.model_list ) @@ -542,7 +551,10 @@ def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: di litellm.input_callback = [self.leastbusy_logger] # type: ignore if isinstance(litellm.callbacks, list): litellm.callbacks.append(self.leastbusy_logger) # type: ignore - elif routing_strategy == "usage-based-routing": + elif ( + routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING.value + or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING + ): self.lowesttpm_logger = LowestTPMLoggingHandler( router_cache=self.cache, model_list=self.model_list, @@ -550,7 +562,10 @@ def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: di ) if isinstance(litellm.callbacks, list): litellm.callbacks.append(self.lowesttpm_logger) # type: ignore - elif routing_strategy == "usage-based-routing-v2": + elif ( + routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2.value + or routing_strategy == RoutingStrategy.USAGE_BASED_ROUTING_V2 + ): self.lowesttpm_logger_v2 = LowestTPMLoggingHandler_v2( router_cache=self.cache, model_list=self.model_list, @@ -558,7 +573,10 @@ def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: di ) if isinstance(litellm.callbacks, list): litellm.callbacks.append(self.lowesttpm_logger_v2) # type: ignore - elif routing_strategy == "latency-based-routing": + elif ( + routing_strategy == RoutingStrategy.LATENCY_BASED.value + or routing_strategy == RoutingStrategy.LATENCY_BASED + ): self.lowestlatency_logger = LowestLatencyLoggingHandler( router_cache=self.cache, model_list=self.model_list, @@ -566,7 +584,10 @@ def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: di ) if isinstance(litellm.callbacks, list): litellm.callbacks.append(self.lowestlatency_logger) # type: ignore - elif routing_strategy == "cost-based-routing": + elif ( + routing_strategy == RoutingStrategy.COST_BASED.value + or routing_strategy == RoutingStrategy.COST_BASED + ): self.lowestcost_logger = LowestCostLoggingHandler( router_cache=self.cache, model_list=self.model_list, @@ -574,10 +595,14 @@ def routing_strategy_init(self, routing_strategy: str, routing_strategy_args: di ) if isinstance(litellm.callbacks, list): litellm.callbacks.append(self.lowestcost_logger) # type: ignore + else: + pass def print_deployment(self, deployment: dict): """ returns a copy of the deployment with the api key masked + + Only returns 2 characters of the api key and masks the rest with * (10 *). """ try: _deployment_copy = copy.deepcopy(deployment) @@ -1746,7 +1771,6 @@ def text_completion( try: kwargs["model"] = model kwargs["prompt"] = prompt - kwargs["original_function"] = self.text_completion kwargs["num_retries"] = kwargs.get("num_retries", self.num_retries) kwargs.get("request_timeout", self.timeout) kwargs.setdefault("metadata", {}).update({"model_group": model}) @@ -1770,13 +1794,7 @@ def text_completion( # call via litellm.completion() return litellm.text_completion(**{**data, "prompt": prompt, "caching": self.cache_responses, **kwargs}) # type: ignore except Exception as e: - if self.num_retries > 0: - kwargs["model"] = model - kwargs["messages"] = messages - kwargs["original_function"] = self.text_completion - return self.function_with_retries(**kwargs) - else: - raise e + raise e async def atext_completion( self, @@ -3005,7 +3023,7 @@ async def async_function_with_retries(self, *args, **kwargs): async def make_call(self, original_function: Any, *args, **kwargs): """ - Handler for making a call to the .completion()/.embeddings() functions. + Handler for making a call to the .completion()/.embeddings()/etc. functions. """ model_group = kwargs.get("model") response = await original_function(*args, **kwargs) diff --git a/litellm/types/router.py b/litellm/types/router.py index 1f2c7224f946..f0737b3ef00a 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -608,3 +608,11 @@ class RouterModelGroupAliasItem(TypedDict): "staging", "production", ] + + +class RoutingStrategy(enum.Enum): + LEAST_BUSY = "least-busy" + LATENCY_BASED = "latency-based-routing" + COST_BASED = "cost-based-routing" + USAGE_BASED_ROUTING_V2 = "usage-based-routing-v2" + USAGE_BASED_ROUTING = "usage-based-routing" diff --git a/litellm/types/utils.py b/litellm/types/utils.py index c3118b453b96..cbc0f0274a64 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -11,6 +11,12 @@ CompletionUsage, PromptTokensDetails, ) +from openai.types.moderation import ( + Categories, + CategoryAppliedInputTypes, + CategoryScores, +) +from openai.types.moderation_create_response import Moderation, ModerationCreateResponse from pydantic import BaseModel, ConfigDict, PrivateAttr from typing_extensions import Callable, Dict, Required, TypedDict, override @@ -20,6 +26,7 @@ ChatCompletionUsageBlock, OpenAIChatCompletionChunk, ) +from .rerank import RerankResponse def _generate_id(): # private helper function @@ -811,7 +818,7 @@ def __init__( model: Optional[str] = None, usage: Optional[Usage] = None, response_ms=None, - data: Optional[List] = None, + data: Optional[Union[List, List[Embedding]]] = None, hidden_params=None, _response_headers=None, **params, diff --git a/pyproject.toml b/pyproject.toml index 1e533b046330..a70a75d14377 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [tool.poetry] name = "litellm" -version = "1.49.3" +version = "1.49.4" description = "Library to easily interface with LLM API providers" authors = ["BerriAI"] license = "MIT" @@ -91,7 +91,7 @@ requires = ["poetry-core", "wheel"] build-backend = "poetry.core.masonry.api" [tool.commitizen] -version = "1.49.3" +version = "1.49.4" version_files = [ "pyproject.toml:^version" ] diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py new file mode 100644 index 000000000000..fb88c3504198 --- /dev/null +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -0,0 +1,117 @@ +import ast +import os + + +def get_function_names_from_file(file_path): + """ + Extracts all function names from a given Python file. + """ + with open(file_path, "r") as file: + tree = ast.parse(file.read()) + + function_names = [] + + for node in ast.walk(tree): + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + function_names.append(node.name) + + return function_names + + +def get_all_functions_called_in_tests(base_dir): + """ + Returns a set of function names that are called in test functions + inside 'local_testing' and 'router_unit_test' directories, + specifically in files containing the word 'router'. + """ + called_functions = set() + test_dirs = ["local_testing", "router_unit_tests"] + + for test_dir in test_dirs: + dir_path = os.path.join(base_dir, test_dir) + if not os.path.exists(dir_path): + print(f"Warning: Directory {dir_path} does not exist.") + continue + + print("dir_path: ", dir_path) + for root, _, files in os.walk(dir_path): + for file in files: + if file.endswith(".py") and "router" in file.lower(): + print("file: ", file) + file_path = os.path.join(root, file) + with open(file_path, "r") as f: + try: + tree = ast.parse(f.read()) + except SyntaxError: + print(f"Warning: Syntax error in file {file_path}") + continue + if file == "test_router_validate_fallbacks.py": + print(f"tree: {tree}") + for node in ast.walk(tree): + if isinstance(node, ast.Call) and isinstance( + node.func, ast.Name + ): + called_functions.add(node.func.id) + elif isinstance(node, ast.Call) and isinstance( + node.func, ast.Attribute + ): + called_functions.add(node.func.attr) + + return called_functions + + +def get_functions_from_router(file_path): + """ + Extracts all functions defined in router.py. + """ + return get_function_names_from_file(file_path) + + +ignored_function_names = [ + "__init__", + "_acreate_file", + "_acreate_batch", + "acreate_assistants", + "adelete_assistant", + "aget_assistants", + "acreate_thread", + "aget_thread", + "a_add_message", + "aget_messages", + "arun_thread", +] + + +def main(): + router_file = "./litellm/router.py" # Update this path if it's located elsewhere + # router_file = "../../litellm/router.py" ## LOCAL TESTING + tests_dir = ( + "./tests/" # Update this path if your tests directory is located elsewhere + ) + # tests_dir = "../../tests/" # LOCAL TESTING + + router_functions = get_functions_from_router(router_file) + print("router_functions: ", router_functions) + called_functions_in_tests = get_all_functions_called_in_tests(tests_dir) + untested_functions = [ + fn for fn in router_functions if fn not in called_functions_in_tests + ] + + if untested_functions: + all_untested_functions = [] + for func in untested_functions: + if func not in ignored_function_names: + all_untested_functions.append(func) + untested_perc = (len(all_untested_functions)) / len(router_functions) + print("perc_covered: ", untested_perc) + if untested_perc < 0.3: + print("The following functions in router.py are not tested:") + raise Exception( + f"{untested_perc * 100:.2f}% of functions in router.py are not tested: {all_untested_functions}" + ) + else: + print("All functions in router.py are covered by tests.") + + +if __name__ == "__main__": + main() diff --git a/tests/llm_translation/test_anthropic_completion.py b/tests/llm_translation/test_anthropic_completion.py index 2d5dd570afd9..ca402903b476 100644 --- a/tests/llm_translation/test_anthropic_completion.py +++ b/tests/llm_translation/test_anthropic_completion.py @@ -366,41 +366,6 @@ def test_anthropic_tool_streaming(): assert tool_use["index"] == correct_tool_index -@pytest.mark.asyncio -async def test_anthropic_router_completion_e2e(): - litellm.set_verbose = True - - litellm.adapters = [{"id": "anthropic", "adapter": anthropic_adapter}] - - router = Router( - model_list=[ - { - "model_name": "claude-3-5-sonnet-20240620", - "litellm_params": { - "model": "gpt-3.5-turbo", - "mock_response": "hi this is macintosh.", - }, - } - ] - ) - messages = [{"role": "user", "content": "Hey, how's it going?"}] - - response = await router.aadapter_completion( - model="claude-3-5-sonnet-20240620", - messages=messages, - adapter_id="anthropic", - mock_response="This is a fake call", - ) - - print("Response: {}".format(response)) - - assert response is not None - - assert isinstance(response, AnthropicResponse) - - assert response.model == "gpt-3.5-turbo" - - def test_anthropic_tool_calling_translation(): kwargs = { "model": "claude-3-5-sonnet-20240620", diff --git a/tests/local_testing/test_audio_speech.py b/tests/local_testing/test_audio_speech.py index 4e45b995383b..ac37b1b0ce33 100644 --- a/tests/local_testing/test_audio_speech.py +++ b/tests/local_testing/test_audio_speech.py @@ -83,43 +83,6 @@ async def test_audio_speech_litellm(sync_mode, model, api_base, api_key): assert isinstance(response, HttpxBinaryResponseContent) -@pytest.mark.parametrize("mode", ["iterator"]) # "file", -@pytest.mark.asyncio -async def test_audio_speech_router(mode): - speech_file_path = Path(__file__).parent / "speech.mp3" - - from litellm import Router - - client = Router( - model_list=[ - { - "model_name": "tts", - "litellm_params": { - "model": "openai/tts-1", - }, - }, - ] - ) - - response = await client.aspeech( - model="tts", - voice="alloy", - input="the quick brown fox jumped over the lazy dogs", - api_base=None, - api_key=None, - organization=None, - project=None, - max_retries=1, - timeout=600, - client=None, - optional_params={}, - ) - - from litellm.llms.OpenAI.openai import HttpxBinaryResponseContent - - assert isinstance(response, HttpxBinaryResponseContent) - - @pytest.mark.parametrize( "sync_mode", [False, True], diff --git a/tests/local_testing/test_custom_callback_input.py b/tests/local_testing/test_custom_callback_input.py index c079123e7976..33ae09bc18d9 100644 --- a/tests/local_testing/test_custom_callback_input.py +++ b/tests/local_testing/test_custom_callback_input.py @@ -280,6 +280,9 @@ async def async_log_pre_api_call(self, model, messages, kwargs): async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): try: + print( + "in async_log_success_event", kwargs, response_obj, start_time, end_time + ) self.states.append("async_success") ## START TIME assert isinstance(start_time, datetime) @@ -522,6 +525,7 @@ async def test_async_chat_azure_stream(): @pytest.mark.asyncio async def test_async_chat_openai_stream_options(): try: + litellm.set_verbose = True customHandler = CompletionCustomHandler() litellm.callbacks = [customHandler] with patch.object( @@ -536,7 +540,7 @@ async def test_async_chat_openai_stream_options(): async for chunk in response: continue - + print("mock client args list=", mock_client.await_args_list) mock_client.assert_awaited_once() except Exception as e: pytest.fail(f"An exception occurred: {str(e)}") diff --git a/tests/local_testing/test_router.py b/tests/local_testing/test_router.py index 42148d9ab449..36ebdb382753 100644 --- a/tests/local_testing/test_router.py +++ b/tests/local_testing/test_router.py @@ -1310,19 +1310,38 @@ def test_aembedding_on_router(): router = Router(model_list=model_list) async def embedding_call(): + ## Test 1: user facing function response = await router.aembedding( model="text-embedding-ada-002", input=["good morning from litellm", "this is another item"], ) print(response) + ## Test 2: underlying function + response = await router._aembedding( + model="text-embedding-ada-002", + input=["good morning from litellm 2"], + ) + print(response) + router.reset() + asyncio.run(embedding_call()) print("\n Making sync Embedding call\n") + ## Test 1: user facing function response = router.embedding( model="text-embedding-ada-002", input=["good morning from litellm 2"], ) + print(response) + router.reset() + + ## Test 2: underlying function + response = router._embedding( + model="text-embedding-ada-002", + input=["good morning from litellm 2"], + ) + print(response) router.reset() except Exception as e: if "Your task failed as a result of our safety system." in str(e): @@ -1843,10 +1862,16 @@ async def test_router_amoderation(): ] router = Router(model_list=model_list) + ## Test 1: user facing function result = await router.amoderation( model="openai-moderations", input="this is valid good text" ) + ## Test 2: underlying function + result = await router._amoderation( + model="openai-moderations", input="this is valid good text" + ) + print("moderation result", result) diff --git a/tests/local_testing/test_whisper.py b/tests/local_testing/test_whisper.py index 087028928d4a..f66ad8b133fc 100644 --- a/tests/local_testing/test_whisper.py +++ b/tests/local_testing/test_whisper.py @@ -79,81 +79,6 @@ async def test_transcription(model, api_key, api_base, response_format, sync_mod assert transcript.text is not None -# This file includes the custom callbacks for LiteLLM Proxy -# Once defined, these can be passed in proxy_config.yaml -class MyCustomHandler(CustomLogger): - def __init__(self): - self.openai_client = None - - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - try: - # init logging config - print("logging a transcript kwargs: ", kwargs) - print("openai client=", kwargs.get("client")) - self.openai_client = kwargs.get("client") - - except Exception: - pass - - -proxy_handler_instance = MyCustomHandler() - - -# Set litellm.callbacks = [proxy_handler_instance] on the proxy -# need to set litellm.callbacks = [proxy_handler_instance] # on the proxy -@pytest.mark.asyncio -async def test_transcription_on_router(): - litellm.set_verbose = True - litellm.callbacks = [proxy_handler_instance] - print("\n Testing async transcription on router\n") - try: - model_list = [ - { - "model_name": "whisper", - "litellm_params": { - "model": "whisper-1", - }, - }, - { - "model_name": "whisper", - "litellm_params": { - "model": "azure/azure-whisper", - "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com/", - "api_key": os.getenv("AZURE_EUROPE_API_KEY"), - "api_version": "2024-02-15-preview", - }, - }, - ] - - router = Router(model_list=model_list) - - router_level_clients = [] - for deployment in router.model_list: - _deployment_openai_client = router._get_client( - deployment=deployment, - kwargs={"model": "whisper-1"}, - client_type="async", - ) - - router_level_clients.append(str(_deployment_openai_client)) - - response = await router.atranscription( - model="whisper", - file=audio_file, - ) - print(response) - - # PROD Test - # Ensure we ONLY use OpenAI/Azure client initialized on the router level - await asyncio.sleep(5) - print("OpenAI Client used= ", proxy_handler_instance.openai_client) - print("all router level clients= ", router_level_clients) - assert proxy_handler_instance.openai_client in router_level_clients - except Exception as e: - traceback.print_exc() - pytest.fail(f"Error occurred: {e}") - - @pytest.mark.asyncio() async def test_transcription_caching(): import litellm diff --git a/tests/logging_callback_tests/conftest.py b/tests/logging_callback_tests/conftest.py new file mode 100644 index 000000000000..eca0bc431a5a --- /dev/null +++ b/tests/logging_callback_tests/conftest.py @@ -0,0 +1,54 @@ +# conftest.py + +import importlib +import os +import sys + +import pytest + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +import litellm + + +@pytest.fixture(scope="function", autouse=True) +def setup_and_teardown(): + """ + This fixture reloads litellm before every function. To speed up testing by removing callbacks being chained. + """ + curr_dir = os.getcwd() # Get the current working directory + sys.path.insert( + 0, os.path.abspath("../..") + ) # Adds the project directory to the system path + + import litellm + from litellm import Router + + importlib.reload(litellm) + import asyncio + + loop = asyncio.get_event_loop_policy().new_event_loop() + asyncio.set_event_loop(loop) + print(litellm) + # from litellm import Router, completion, aembedding, acompletion, embedding + yield + + # Teardown code (executes after the yield point) + loop.close() # Close the loop created earlier + asyncio.set_event_loop(None) # Remove the reference to the loop + + +def pytest_collection_modifyitems(config, items): + # Separate tests in 'test_amazing_proxy_custom_logger.py' and other tests + custom_logger_tests = [ + item for item in items if "custom_logger" in item.parent.name + ] + other_tests = [item for item in items if "custom_logger" not in item.parent.name] + + # Sort tests based on their names + custom_logger_tests.sort(key=lambda x: x.name) + other_tests.sort(key=lambda x: x.name) + + # Reorder the items list + items[:] = custom_logger_tests + other_tests diff --git a/tests/logging_callback_tests/test_assemble_streaming_responses.py b/tests/logging_callback_tests/test_assemble_streaming_responses.py new file mode 100644 index 000000000000..5b52194bfa34 --- /dev/null +++ b/tests/logging_callback_tests/test_assemble_streaming_responses.py @@ -0,0 +1,362 @@ +""" +Testing for _assemble_complete_response_from_streaming_chunks + +- Test 1 - ModelResponse with 1 list of streaming chunks. Assert chunks are added to the streaming_chunks, after final chunk sent assert complete_streaming_response is not None +- Test 2 - TextCompletionResponse with 1 list of streaming chunks. Assert chunks are added to the streaming_chunks, after final chunk sent assert complete_streaming_response is not None +- Test 3 - Have multiple lists of streaming chunks, Assert that chunks are added to the correct list and that complete_streaming_response is None. After final chunk sent assert complete_streaming_response is not None +- Test 4 - build a complete response when 1 chunk is poorly formatted + +""" + +import json +import os +import sys +from datetime import datetime +from unittest.mock import AsyncMock + +from pydantic.main import Model + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path + + +import httpx +import pytest +from respx import MockRouter + +import litellm +from litellm import Choices, Message, ModelResponse, TextCompletionResponse, TextChoices + +from litellm.litellm_core_utils.litellm_logging import ( + _assemble_complete_response_from_streaming_chunks, +) + + +@pytest.mark.parametrize("is_async", [True, False]) +def test_assemble_complete_response_from_streaming_chunks_1(is_async): + """ + Test 1 - ModelResponse with 1 list of streaming chunks. Assert chunks are added to the streaming_chunks, after final chunk sent assert complete_streaming_response is not None + """ + + request_kwargs = { + "model": "test_model", + "messages": [{"role": "user", "content": "Hello, world!"}], + } + + list_streaming_chunks = [] + chunk = { + "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", + "choices": [ + litellm.utils.StreamingChoices( + delta=litellm.utils.Delta( + content="hello in response", + function_call=None, + role=None, + tool_calls=None, + ), + index=0, + logprobs=None, + ) + ], + "created": 1721353246, + "model": "gpt-3.5-turbo", + "object": "chat.completion.chunk", + "system_fingerprint": None, + "usage": None, + } + chunk = litellm.ModelResponse(**chunk, stream=True) + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks, + is_async=is_async, + ) + + # this is the 1st chunk - complete_streaming_response should be None + + print("list_streaming_chunks", list_streaming_chunks) + print("complete_streaming_response", complete_streaming_response) + assert complete_streaming_response is None + assert len(list_streaming_chunks) == 1 + assert list_streaming_chunks[0] == chunk + + # Add final chunk + chunk = { + "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", + "choices": [ + litellm.utils.StreamingChoices( + finish_reason="stop", + delta=litellm.utils.Delta( + content="end of response", + function_call=None, + role=None, + tool_calls=None, + ), + index=0, + logprobs=None, + ) + ], + "created": 1721353246, + "model": "gpt-3.5-turbo", + "object": "chat.completion.chunk", + "system_fingerprint": None, + "usage": None, + } + chunk = litellm.ModelResponse(**chunk, stream=True) + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks, + is_async=is_async, + ) + + print("list_streaming_chunks", list_streaming_chunks) + print("complete_streaming_response", complete_streaming_response) + + # this is the 2nd chunk - complete_streaming_response should not be None + assert complete_streaming_response is not None + assert len(list_streaming_chunks) == 2 + + assert isinstance(complete_streaming_response, ModelResponse) + assert isinstance(complete_streaming_response.choices[0], Choices) + + pass + + +@pytest.mark.parametrize("is_async", [True, False]) +def test_assemble_complete_response_from_streaming_chunks_2(is_async): + """ + Test 2 - TextCompletionResponse with 1 list of streaming chunks. Assert chunks are added to the streaming_chunks, after final chunk sent assert complete_streaming_response is not None + """ + + from litellm.utils import TextCompletionStreamWrapper + + _text_completion_stream_wrapper = TextCompletionStreamWrapper( + completion_stream=None, model="test_model" + ) + + request_kwargs = { + "model": "test_model", + "messages": [{"role": "user", "content": "Hello, world!"}], + } + + list_streaming_chunks = [] + chunk = { + "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", + "choices": [ + litellm.utils.StreamingChoices( + delta=litellm.utils.Delta( + content="hello in response", + function_call=None, + role=None, + tool_calls=None, + ), + index=0, + logprobs=None, + ) + ], + "created": 1721353246, + "model": "gpt-3.5-turbo", + "object": "chat.completion.chunk", + "system_fingerprint": None, + "usage": None, + } + chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = _text_completion_stream_wrapper.convert_to_text_completion_object(chunk) + + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks, + is_async=is_async, + ) + + # this is the 1st chunk - complete_streaming_response should be None + + print("list_streaming_chunks", list_streaming_chunks) + print("complete_streaming_response", complete_streaming_response) + assert complete_streaming_response is None + assert len(list_streaming_chunks) == 1 + assert list_streaming_chunks[0] == chunk + + # Add final chunk + chunk = { + "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", + "choices": [ + litellm.utils.StreamingChoices( + finish_reason="stop", + delta=litellm.utils.Delta( + content="end of response", + function_call=None, + role=None, + tool_calls=None, + ), + index=0, + logprobs=None, + ) + ], + "created": 1721353246, + "model": "gpt-3.5-turbo", + "object": "chat.completion.chunk", + "system_fingerprint": None, + "usage": None, + } + chunk = litellm.ModelResponse(**chunk, stream=True) + chunk = _text_completion_stream_wrapper.convert_to_text_completion_object(chunk) + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks, + is_async=is_async, + ) + + print("list_streaming_chunks", list_streaming_chunks) + print("complete_streaming_response", complete_streaming_response) + + # this is the 2nd chunk - complete_streaming_response should not be None + assert complete_streaming_response is not None + assert len(list_streaming_chunks) == 2 + + assert isinstance(complete_streaming_response, TextCompletionResponse) + assert isinstance(complete_streaming_response.choices[0], TextChoices) + + pass + + +@pytest.mark.parametrize("is_async", [True, False]) +def test_assemble_complete_response_from_streaming_chunks_3(is_async): + + request_kwargs = { + "model": "test_model", + "messages": [{"role": "user", "content": "Hello, world!"}], + } + + list_streaming_chunks_1 = [] + list_streaming_chunks_2 = [] + + chunk = { + "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", + "choices": [ + litellm.utils.StreamingChoices( + delta=litellm.utils.Delta( + content="hello in response", + function_call=None, + role=None, + tool_calls=None, + ), + index=0, + logprobs=None, + ) + ], + "created": 1721353246, + "model": "gpt-3.5-turbo", + "object": "chat.completion.chunk", + "system_fingerprint": None, + "usage": None, + } + chunk = litellm.ModelResponse(**chunk, stream=True) + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks_1, + is_async=is_async, + ) + + # this is the 1st chunk - complete_streaming_response should be None + + print("list_streaming_chunks_1", list_streaming_chunks_1) + print("complete_streaming_response", complete_streaming_response) + assert complete_streaming_response is None + assert len(list_streaming_chunks_1) == 1 + assert list_streaming_chunks_1[0] == chunk + assert len(list_streaming_chunks_2) == 0 + + # now add a chunk to the 2nd list + + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks_2, + is_async=is_async, + ) + + print("list_streaming_chunks_2", list_streaming_chunks_2) + print("complete_streaming_response", complete_streaming_response) + assert complete_streaming_response is None + assert len(list_streaming_chunks_2) == 1 + assert list_streaming_chunks_2[0] == chunk + assert len(list_streaming_chunks_1) == 1 + + # now add a chunk to the 1st list + + +@pytest.mark.parametrize("is_async", [True, False]) +def test_assemble_complete_response_from_streaming_chunks_4(is_async): + """ + Test 4 - build a complete response when 1 chunk is poorly formatted + + - Assert complete_streaming_response is None + - Assert list_streaming_chunks is not empty + """ + + request_kwargs = { + "model": "test_model", + "messages": [{"role": "user", "content": "Hello, world!"}], + } + + list_streaming_chunks = [] + + chunk = { + "id": "chatcmpl-9mWtyDnikZZoB75DyfUzWUxiiE2Pi", + "choices": [ + litellm.utils.StreamingChoices( + finish_reason="stop", + delta=litellm.utils.Delta( + content="end of response", + function_call=None, + role=None, + tool_calls=None, + ), + index=0, + logprobs=None, + ) + ], + "created": 1721353246, + "model": "gpt-3.5-turbo", + "object": "chat.completion.chunk", + "system_fingerprint": None, + "usage": None, + } + chunk = litellm.ModelResponse(**chunk, stream=True) + + # remove attribute id from chunk + del chunk.id + + complete_streaming_response = _assemble_complete_response_from_streaming_chunks( + result=chunk, + start_time=datetime.now(), + end_time=datetime.now(), + request_kwargs=request_kwargs, + streaming_chunks=list_streaming_chunks, + is_async=is_async, + ) + + print("complete_streaming_response", complete_streaming_response) + assert complete_streaming_response is None + + print("list_streaming_chunks", list_streaming_chunks) + + assert len(list_streaming_chunks) == 1 diff --git a/tests/router_unit_tests/README.md b/tests/router_unit_tests/README.md new file mode 100644 index 000000000000..0206bcbb8f36 --- /dev/null +++ b/tests/router_unit_tests/README.md @@ -0,0 +1,5 @@ +## Router component unit tests. + +Please name all files with the word 'router' in them. + +This is used to ensure all functions in the router are tested. \ No newline at end of file diff --git a/tests/router_unit_tests/gettysburg.wav b/tests/router_unit_tests/gettysburg.wav new file mode 100644 index 000000000000..9690f521e844 Binary files /dev/null and b/tests/router_unit_tests/gettysburg.wav differ diff --git a/tests/router_unit_tests/test_router_endpoints.py b/tests/router_unit_tests/test_router_endpoints.py new file mode 100644 index 000000000000..accd5ea40569 --- /dev/null +++ b/tests/router_unit_tests/test_router_endpoints.py @@ -0,0 +1,279 @@ +import sys +import os +import traceback +from dotenv import load_dotenv +from fastapi import Request +from datetime import datetime + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from litellm import Router, CustomLogger + +# Get the current directory of the file being run +pwd = os.path.dirname(os.path.realpath(__file__)) +print(pwd) + +file_path = os.path.join(pwd, "gettysburg.wav") + +audio_file = open(file_path, "rb") +from pathlib import Path +import litellm +import pytest +import asyncio + + +@pytest.fixture +def model_list(): + return [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "gpt-4o", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + { + "model_name": "dall-e-3", + "litellm_params": { + "model": "dall-e-3", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + { + "model_name": "cohere-rerank", + "litellm_params": { + "model": "cohere/rerank-english-v3.0", + "api_key": os.getenv("COHERE_API_KEY"), + }, + }, + { + "model_name": "claude-3-5-sonnet-20240620", + "litellm_params": { + "model": "gpt-3.5-turbo", + "mock_response": "hi this is macintosh.", + }, + }, + ] + + +# This file includes the custom callbacks for LiteLLM Proxy +# Once defined, these can be passed in proxy_config.yaml +class MyCustomHandler(CustomLogger): + def __init__(self): + self.openai_client = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + try: + # init logging config + print("logging a transcript kwargs: ", kwargs) + print("openai client=", kwargs.get("client")) + self.openai_client = kwargs.get("client") + + except Exception: + pass + + +proxy_handler_instance = MyCustomHandler() + + +# Set litellm.callbacks = [proxy_handler_instance] on the proxy +# need to set litellm.callbacks = [proxy_handler_instance] # on the proxy +@pytest.mark.asyncio +async def test_transcription_on_router(): + litellm.set_verbose = True + litellm.callbacks = [proxy_handler_instance] + print("\n Testing async transcription on router\n") + try: + model_list = [ + { + "model_name": "whisper", + "litellm_params": { + "model": "whisper-1", + }, + }, + { + "model_name": "whisper", + "litellm_params": { + "model": "azure/azure-whisper", + "api_base": "https://my-endpoint-europe-berri-992.openai.azure.com/", + "api_key": os.getenv("AZURE_EUROPE_API_KEY"), + "api_version": "2024-02-15-preview", + }, + }, + ] + + router = Router(model_list=model_list) + + router_level_clients = [] + for deployment in router.model_list: + _deployment_openai_client = router._get_client( + deployment=deployment, + kwargs={"model": "whisper-1"}, + client_type="async", + ) + + router_level_clients.append(str(_deployment_openai_client)) + + ## test 1: user facing function + response = await router.atranscription( + model="whisper", + file=audio_file, + ) + + ## test 2: underlying function + response = await router._atranscription( + model="whisper", + file=audio_file, + ) + print(response) + + # PROD Test + # Ensure we ONLY use OpenAI/Azure client initialized on the router level + await asyncio.sleep(5) + print("OpenAI Client used= ", proxy_handler_instance.openai_client) + print("all router level clients= ", router_level_clients) + assert proxy_handler_instance.openai_client in router_level_clients + except Exception as e: + traceback.print_exc() + pytest.fail(f"Error occurred: {e}") + + +@pytest.mark.parametrize("mode", ["iterator"]) # "file", +@pytest.mark.asyncio +async def test_audio_speech_router(mode): + + from litellm import Router + + client = Router( + model_list=[ + { + "model_name": "tts", + "litellm_params": { + "model": "openai/tts-1", + }, + }, + ] + ) + + response = await client.aspeech( + model="tts", + voice="alloy", + input="the quick brown fox jumped over the lazy dogs", + api_base=None, + api_key=None, + organization=None, + project=None, + max_retries=1, + timeout=600, + client=None, + optional_params={}, + ) + + from litellm.llms.OpenAI.openai import HttpxBinaryResponseContent + + assert isinstance(response, HttpxBinaryResponseContent) + + +@pytest.mark.asyncio() +async def test_rerank_endpoint(model_list): + from litellm.types.utils import RerankResponse + + router = Router(model_list=model_list) + + ## Test 1: user facing function + response = await router.arerank( + model="cohere-rerank", + query="hello", + documents=["hello", "world"], + top_n=3, + ) + + ## Test 2: underlying function + response = await router._arerank( + model="cohere-rerank", + query="hello", + documents=["hello", "world"], + top_n=3, + ) + + print("async re rank response: ", response) + + assert response.id is not None + assert response.results is not None + + RerankResponse.model_validate(response) + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_text_completion_endpoint(model_list, sync_mode): + router = Router(model_list=model_list) + + if sync_mode: + response = router.text_completion( + model="gpt-3.5-turbo", + prompt="Hello, how are you?", + mock_response="I'm fine, thank you!", + ) + else: + ## Test 1: user facing function + response = await router.atext_completion( + model="gpt-3.5-turbo", + prompt="Hello, how are you?", + mock_response="I'm fine, thank you!", + ) + + ## Test 2: underlying function + response_2 = await router._atext_completion( + model="gpt-3.5-turbo", + prompt="Hello, how are you?", + mock_response="I'm fine, thank you!", + ) + assert response_2.choices[0].text == "I'm fine, thank you!" + + assert response.choices[0].text == "I'm fine, thank you!" + + +@pytest.mark.asyncio +async def test_anthropic_router_completion_e2e(model_list): + from litellm.adapters.anthropic_adapter import anthropic_adapter + from litellm.types.llms.anthropic import AnthropicResponse + + litellm.set_verbose = True + + litellm.adapters = [{"id": "anthropic", "adapter": anthropic_adapter}] + + router = Router(model_list=model_list) + messages = [{"role": "user", "content": "Hey, how's it going?"}] + + ## Test 1: user facing function + response = await router.aadapter_completion( + model="claude-3-5-sonnet-20240620", + messages=messages, + adapter_id="anthropic", + mock_response="This is a fake call", + ) + + ## Test 2: underlying function + await router._aadapter_completion( + model="claude-3-5-sonnet-20240620", + messages=messages, + adapter_id="anthropic", + mock_response="This is a fake call", + ) + + print("Response: {}".format(response)) + + assert response is not None + + AnthropicResponse.model_validate(response) + + assert response.model == "gpt-3.5-turbo" diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py new file mode 100644 index 000000000000..97660471b9b9 --- /dev/null +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -0,0 +1,252 @@ +import sys +import os +import traceback +from dotenv import load_dotenv +from fastapi import Request +from datetime import datetime + +sys.path.insert( + 0, os.path.abspath("../..") +) # Adds the parent directory to the system path +from litellm import Router +import pytest +from unittest.mock import patch, MagicMock, AsyncMock + + +@pytest.fixture +def model_list(): + return [ + { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + { + "model_name": "gpt-4o", + "litellm_params": { + "model": "gpt-4o", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + { + "model_name": "dall-e-3", + "litellm_params": { + "model": "dall-e-3", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + }, + ] + + +def test_validate_fallbacks(model_list): + router = Router(model_list=model_list, fallbacks=[{"gpt-4o": "gpt-3.5-turbo"}]) + router.validate_fallbacks(fallback_param=[{"gpt-4o": "gpt-3.5-turbo"}]) + + +def test_routing_strategy_init(model_list): + """Test if all routing strategies are initialized correctly""" + from litellm.types.router import RoutingStrategy + + router = Router(model_list=model_list) + for strategy in RoutingStrategy._member_names_: + router.routing_strategy_init( + routing_strategy=strategy, routing_strategy_args={} + ) + + +def test_print_deployment(model_list): + """Test if the api key is masked correctly""" + + router = Router(model_list=model_list) + deployment = { + "model_name": "gpt-3.5-turbo", + "litellm_params": { + "model": "gpt-3.5-turbo", + "api_key": os.getenv("OPENAI_API_KEY"), + }, + } + printed_deployment = router.print_deployment(deployment) + assert 10 * "*" in printed_deployment["litellm_params"]["api_key"] + + +def test_completion(model_list): + """Test if the completion function is working correctly""" + router = Router(model_list=model_list) + response = router._completion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, how are you?"}], + mock_response="I'm fine, thank you!", + ) + assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.flaky(retries=6, delay=1) +@pytest.mark.asyncio +async def test_image_generation(model_list, sync_mode): + """Test if the underlying '_image_generation' function is working correctly""" + from litellm.types.utils import ImageResponse + + router = Router(model_list=model_list) + if sync_mode: + response = router._image_generation( + model="dall-e-3", + prompt="A cute baby sea otter", + ) + else: + response = await router._aimage_generation( + model="dall-e-3", + prompt="A cute baby sea otter", + ) + + ImageResponse.model_validate(response) + + +@pytest.mark.asyncio +async def test_router_acompletion_util(model_list): + """Test if the underlying '_acompletion' function is working correctly""" + router = Router(model_list=model_list) + response = await router._acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, how are you?"}], + mock_response="I'm fine, thank you!", + ) + assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" + + +@pytest.mark.asyncio +async def test_router_abatch_completion_one_model_multiple_requests_util(model_list): + """Test if the 'abatch_completion_one_model_multiple_requests' function is working correctly""" + router = Router(model_list=model_list) + response = await router.abatch_completion_one_model_multiple_requests( + model="gpt-3.5-turbo", + messages=[ + [{"role": "user", "content": "Hello, how are you?"}], + [{"role": "user", "content": "Hello, how are you?"}], + ], + mock_response="I'm fine, thank you!", + ) + print(response) + assert response[0]["choices"][0]["message"]["content"] == "I'm fine, thank you!" + assert response[1]["choices"][0]["message"]["content"] == "I'm fine, thank you!" + + +@pytest.mark.asyncio +async def test_router_schedule_acompletion(model_list): + """Test if the 'schedule_acompletion' function is working correctly""" + router = Router(model_list=model_list) + response = await router.schedule_acompletion( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, how are you?"}], + mock_response="I'm fine, thank you!", + priority=1, + ) + assert response["choices"][0]["message"]["content"] == "I'm fine, thank you!" + + +@pytest.mark.asyncio +async def test_router_arealtime(model_list): + """Test if the '_arealtime' function is working correctly""" + import litellm + + router = Router(model_list=model_list) + with patch.object(litellm, "_arealtime", AsyncMock()) as mock_arealtime: + mock_arealtime.return_value = "I'm fine, thank you!" + await router._arealtime( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, how are you?"}], + ) + + mock_arealtime.assert_awaited_once() + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_router_function_with_fallbacks(model_list, sync_mode): + """Test if the router 'async_function_with_fallbacks' + 'function_with_fallbacks' are working correctly""" + router = Router(model_list=model_list) + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello, how are you?"}], + "mock_response": "I'm fine, thank you!", + "num_retries": 0, + } + if sync_mode: + response = router.function_with_fallbacks( + original_function=router._completion, + **data, + ) + else: + response = await router.async_function_with_fallbacks( + original_function=router._acompletion, + **data, + ) + assert response.choices[0].message.content == "I'm fine, thank you!" + + +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.asyncio +async def test_router_function_with_retries(model_list, sync_mode): + """Test if the router 'async_function_with_retries' + 'function_with_retries' are working correctly""" + router = Router(model_list=model_list) + data = { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello, how are you?"}], + "mock_response": "I'm fine, thank you!", + "num_retries": 0, + } + if sync_mode: + response = router.function_with_retries( + original_function=router._completion, + **data, + ) + else: + response = await router.async_function_with_retries( + original_function=router._acompletion, + **data, + ) + assert response.choices[0].message.content == "I'm fine, thank you!" + + +@pytest.mark.asyncio +async def test_router_make_call(model_list): + """Test if the router 'make_call' function is working correctly""" + + ## ACOMPLETION + router = Router(model_list=model_list) + response = await router.make_call( + original_function=router._acompletion, + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "Hello, how are you?"}], + mock_response="I'm fine, thank you!", + ) + assert response.choices[0].message.content == "I'm fine, thank you!" + + ## ATEXT_COMPLETION + response = await router.make_call( + original_function=router._atext_completion, + model="gpt-3.5-turbo", + prompt="Hello, how are you?", + mock_response="I'm fine, thank you!", + ) + assert response.choices[0].text == "I'm fine, thank you!" + + ## AEMBEDDING + response = await router.make_call( + original_function=router._aembedding, + model="gpt-3.5-turbo", + input="Hello, how are you?", + mock_response=[0.1, 0.2, 0.3], + ) + assert response.data[0].embedding == [0.1, 0.2, 0.3] + + ## AIMAGE_GENERATION + response = await router.make_call( + original_function=router._aimage_generation, + model="dall-e-3", + prompt="A cute baby sea otter", + mock_response="https://example.com/image.png", + ) + assert response.data[0].url == "https://example.com/image.png"