diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index 165301efb5c..20e0d412edc 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -1,9 +1,27 @@ """ Translate from OpenAI's `/v1/chat/completions` to Groq's `/v1/chat/completions` """ -from typing import Any, Coroutine, List, Literal, Optional, Tuple, Union, cast, overload +from typing import ( + Any, + Coroutine, + List, + Literal, + Optional, + Tuple, + Union, + cast, + overload, + Iterator, + AsyncIterator, +) import httpx + +from litellm.llms.openai.chat.gpt_transformation import ( + OpenAIChatCompletionStreamingHandler, +) +from litellm.llms.openai.common_utils import OpenAIError + from pydantic import BaseModel import litellm @@ -16,7 +34,7 @@ from litellm.types.llms.openai import ( ChatCompletionToolParam, ChatCompletionToolParamFunctionChunk, ) -from litellm.types.utils import ModelResponse +from litellm.types.utils import ModelResponse, ModelResponseStream from ...openai_like.chat.transformation import OpenAILikeChatConfig @@ -65,6 +83,18 @@ class GroqChatConfig(OpenAILikeChatConfig): def get_config(cls): return super().get_config() + def get_model_response_iterator( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], + sync_stream: bool, + json_mode: Optional[bool] = False, + ) -> Any: + return GroqChatCompletionStreamingHandler( + streaming_response=streaming_response, + sync_stream=sync_stream, + json_mode=json_mode, + ) + def get_supported_openai_params(self, model: str) -> list: base_params = super().get_supported_openai_params(model) try: @@ -209,7 +239,6 @@ class GroqChatConfig(OpenAILikeChatConfig): ) return optional_params - def transform_response( self, @@ -239,12 +268,17 @@ class GroqChatConfig(OpenAILikeChatConfig): json_mode=json_mode, ) - mapped_service_tier: Literal["auto", "default", "flex"] = self._map_groq_service_tier(original_service_tier=getattr(model_response, "service_tier")) + mapped_service_tier: Literal[ + "auto", "default", "flex" + ] = self._map_groq_service_tier( + original_service_tier=getattr(model_response, "service_tier") + ) setattr(model_response, "service_tier", mapped_service_tier) return model_response - - def _map_groq_service_tier(self, original_service_tier: Optional[str]) -> Literal["auto", "default", "flex"]: + def _map_groq_service_tier( + self, original_service_tier: Optional[str] + ) -> Literal["auto", "default", "flex"]: """ Ensure groq service tier is OpenAI compatible. """ @@ -252,5 +286,16 @@ class GroqChatConfig(OpenAILikeChatConfig): return "auto" if original_service_tier not in ["auto", "default", "flex"]: return "auto" - - return cast(Literal["auto", "default", "flex"], original_service_tier) \ No newline at end of file + + return cast(Literal["auto", "default", "flex"], original_service_tier) + + +class GroqChatCompletionStreamingHandler(OpenAIChatCompletionStreamingHandler): + def chunk_parser(self, chunk: dict) -> ModelResponseStream: + error = chunk.get("error") + if error: + raise OpenAIError( + status_code=error.get("code"), message=error.get("message"), body=error + ) + + return super().chunk_parser(chunk)