Spaces:
Sleeping
Sleeping
| import json | |
| from typing import TYPE_CHECKING, Any, AsyncIterator, Iterator, List, Optional, Union | |
| import httpx | |
| from litellm.litellm_core_utils.prompt_templates.common_utils import ( | |
| convert_content_list_to_str, | |
| ) | |
| from litellm.llms.base_llm.base_model_iterator import FakeStreamResponseIterator | |
| from litellm.llms.base_llm.chat.transformation import BaseConfig, BaseLLMException | |
| from litellm.types.llms.openai import AllMessageValues | |
| from litellm.types.utils import ( | |
| ChatCompletionToolCallChunk, | |
| ChatCompletionUsageBlock, | |
| Choices, | |
| GenericStreamingChunk, | |
| Message, | |
| ModelResponse, | |
| Usage, | |
| ) | |
| from litellm.utils import token_counter | |
| from ..common_utils import ClarifaiError | |
| if TYPE_CHECKING: | |
| from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj | |
| LoggingClass = LiteLLMLoggingObj | |
| else: | |
| LoggingClass = Any | |
| class ClarifaiConfig(BaseConfig): | |
| """ | |
| Reference: https://clarifai.com/meta/Llama-2/models/llama2-70b-chat | |
| """ | |
| max_tokens: Optional[int] = None | |
| temperature: Optional[int] = None | |
| top_k: Optional[int] = None | |
| def __init__( | |
| self, | |
| max_tokens: Optional[int] = None, | |
| temperature: Optional[int] = None, | |
| top_k: Optional[int] = None, | |
| ) -> None: | |
| locals_ = locals().copy() | |
| for key, value in locals_.items(): | |
| if key != "self" and value is not None: | |
| setattr(self.__class__, key, value) | |
| def get_config(cls): | |
| return super().get_config() | |
| def get_supported_openai_params(self, model: str) -> list: | |
| return [ | |
| "temperature", | |
| "max_tokens", | |
| ] | |
| def map_openai_params( | |
| self, | |
| non_default_params: dict, | |
| optional_params: dict, | |
| model: str, | |
| drop_params: bool, | |
| ) -> dict: | |
| for param, value in non_default_params.items(): | |
| if param == "temperature": | |
| optional_params["temperature"] = value | |
| elif param == "max_tokens": | |
| optional_params["max_tokens"] = value | |
| return optional_params | |
| def _completions_to_model(self, prompt: str, optional_params: dict) -> dict: | |
| params = {} | |
| if temperature := optional_params.get("temperature"): | |
| params["temperature"] = temperature | |
| if max_tokens := optional_params.get("max_tokens"): | |
| params["max_tokens"] = max_tokens | |
| return { | |
| "inputs": [{"data": {"text": {"raw": prompt}}}], | |
| "model": {"output_info": {"params": params}}, | |
| } | |
| def _convert_model_to_url(self, model: str, api_base: str): | |
| user_id, app_id, model_id = model.split(".") | |
| return f"{api_base}/users/{user_id}/apps/{app_id}/models/{model_id}/outputs" | |
| def transform_request( | |
| self, | |
| model: str, | |
| messages: List[AllMessageValues], | |
| optional_params: dict, | |
| litellm_params: dict, | |
| headers: dict, | |
| ) -> dict: | |
| prompt = " ".join(convert_content_list_to_str(message) for message in messages) | |
| ## Load Config | |
| config = self.get_config() | |
| for k, v in config.items(): | |
| if k not in optional_params: | |
| optional_params[k] = v | |
| data = self._completions_to_model( | |
| prompt=prompt, optional_params=optional_params | |
| ) | |
| return data | |
| def validate_environment( | |
| self, | |
| headers: dict, | |
| model: str, | |
| messages: List[AllMessageValues], | |
| optional_params: dict, | |
| litellm_params: dict, | |
| api_key: Optional[str] = None, | |
| api_base: Optional[str] = None, | |
| ) -> dict: | |
| headers = { | |
| "accept": "application/json", | |
| "content-type": "application/json", | |
| } | |
| if api_key: | |
| headers["Authorization"] = f"Bearer {api_key}" | |
| return headers | |
| def get_error_class( | |
| self, error_message: str, status_code: int, headers: Union[dict, httpx.Headers] | |
| ) -> BaseLLMException: | |
| return ClarifaiError(message=error_message, status_code=status_code) | |
| def transform_response( | |
| self, | |
| model: str, | |
| raw_response: httpx.Response, | |
| model_response: ModelResponse, | |
| logging_obj: LoggingClass, | |
| request_data: dict, | |
| messages: List[AllMessageValues], | |
| optional_params: dict, | |
| litellm_params: dict, | |
| encoding: str, | |
| api_key: Optional[str] = None, | |
| json_mode: Optional[bool] = None, | |
| ) -> ModelResponse: | |
| logging_obj.post_call( | |
| input=messages, | |
| api_key=api_key, | |
| original_response=raw_response.text, | |
| additional_args={"complete_input_dict": request_data}, | |
| ) | |
| ## RESPONSE OBJECT | |
| try: | |
| completion_response = raw_response.json() | |
| except httpx.HTTPStatusError as e: | |
| raise ClarifaiError( | |
| message=str(e), | |
| status_code=raw_response.status_code, | |
| ) | |
| except Exception as e: | |
| raise ClarifaiError( | |
| message=str(e), | |
| status_code=422, | |
| ) | |
| # print(completion_response) | |
| try: | |
| choices_list = [] | |
| for idx, item in enumerate(completion_response["outputs"]): | |
| if len(item["data"]["text"]["raw"]) > 0: | |
| message_obj = Message(content=item["data"]["text"]["raw"]) | |
| else: | |
| message_obj = Message(content=None) | |
| choice_obj = Choices( | |
| finish_reason="stop", | |
| index=idx + 1, # check | |
| message=message_obj, | |
| ) | |
| choices_list.append(choice_obj) | |
| model_response.choices = choices_list # type: ignore | |
| except Exception as e: | |
| raise ClarifaiError( | |
| message=str(e), | |
| status_code=422, | |
| ) | |
| # Calculate Usage | |
| prompt_tokens = token_counter(model=model, messages=messages) | |
| completion_tokens = len( | |
| encoding.encode(model_response["choices"][0]["message"].get("content")) | |
| ) | |
| model_response.model = model | |
| setattr( | |
| model_response, | |
| "usage", | |
| Usage( | |
| prompt_tokens=prompt_tokens, | |
| completion_tokens=completion_tokens, | |
| total_tokens=prompt_tokens + completion_tokens, | |
| ), | |
| ) | |
| return model_response | |
| def get_model_response_iterator( | |
| self, | |
| streaming_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], | |
| sync_stream: bool, | |
| json_mode: Optional[bool] = False, | |
| ) -> Any: | |
| return ClarifaiModelResponseIterator( | |
| model_response=streaming_response, | |
| json_mode=json_mode, | |
| ) | |
| class ClarifaiModelResponseIterator(FakeStreamResponseIterator): | |
| def __init__( | |
| self, | |
| model_response: Union[Iterator[str], AsyncIterator[str], ModelResponse], | |
| json_mode: Optional[bool] = False, | |
| ): | |
| super().__init__( | |
| model_response=model_response, | |
| json_mode=json_mode, | |
| ) | |
| def chunk_parser(self, chunk: dict) -> GenericStreamingChunk: | |
| try: | |
| text = "" | |
| tool_use: Optional[ChatCompletionToolCallChunk] = None | |
| is_finished = False | |
| finish_reason = "" | |
| usage: Optional[ChatCompletionUsageBlock] = None | |
| provider_specific_fields = None | |
| text = ( | |
| chunk.get("outputs", "")[0] | |
| .get("data", "") | |
| .get("text", "") | |
| .get("raw", "") | |
| ) | |
| index: int = 0 | |
| return GenericStreamingChunk( | |
| text=text, | |
| tool_use=tool_use, | |
| is_finished=is_finished, | |
| finish_reason=finish_reason, | |
| usage=usage, | |
| index=index, | |
| provider_specific_fields=provider_specific_fields, | |
| ) | |
| except json.JSONDecodeError: | |
| raise ValueError(f"Failed to decode JSON from chunk: {chunk}") | |