Underground-Digital/Workflow-Engine

1from typing import IO, Optional2 3from requests import Request, Session4from yarl import URL5 6from core.model_runtime.entities.common_entities import I18nObject7from core.model_runtime.entities.model_entities import AIModelEntity, FetchFrom, ModelType8from core.model_runtime.errors.invoke import (9    InvokeAuthorizationError,10    InvokeBadRequestError,11    InvokeConnectionError,12    InvokeError,13    InvokeRateLimitError,14    InvokeServerUnavailableError,15)16from core.model_runtime.errors.validate import CredentialsValidateFailedError17from core.model_runtime.model_providers.__base.speech2text_model import Speech2TextModel18 19 20class LocalAISpeech2text(Speech2TextModel):21    """22    Model class for Local AI Text to speech model.23    """24 25    def _invoke(self, model: str, credentials: dict, file: IO[bytes], user: Optional[str] = None) -> str:26        """27        Invoke large language model28 29        :param model: model name30        :param credentials: model credentials31        :param file: audio file32        :param user: unique user id33        :return: text for given audio file34        """35 36        url = str(URL(credentials["server_url"]) / "v1/audio/transcriptions")37        data = {"model": model}38        files = {"file": file}39 40        session = Session()41        request = Request("POST", url, data=data, files=files)42        prepared_request = session.prepare_request(request)43        response = session.send(prepared_request)44 45        if "error" in response.json():46            raise InvokeServerUnavailableError("Empty response")47 48        return response.json()["text"]49 50    def validate_credentials(self, model: str, credentials: dict) -> None:51        """52        Validate model credentials53 54        :param model: model name55        :param credentials: model credentials56        :return:57        """58        try:59            audio_file_path = self._get_demo_file_path()60 61            with open(audio_file_path, "rb") as audio_file:62                self._invoke(model, credentials, audio_file)63        except Exception as ex:64            raise CredentialsValidateFailedError(str(ex))65 66    @property67    def _invoke_error_mapping(self) -> dict[type[InvokeError], list[type[Exception]]]:68        return {69            InvokeConnectionError: [InvokeConnectionError],70            InvokeServerUnavailableError: [InvokeServerUnavailableError],71            InvokeRateLimitError: [InvokeRateLimitError],72            InvokeAuthorizationError: [InvokeAuthorizationError],73            InvokeBadRequestError: [InvokeBadRequestError],74        }75 76    def get_customizable_model_schema(self, model: str, credentials: dict) -> Optional[AIModelEntity]:77        """78        used to define customizable model schema79        """80        entity = AIModelEntity(81            model=model,82            label=I18nObject(en_US=model),83            fetch_from=FetchFrom.CUSTOMIZABLE_MODEL,84            model_type=ModelType.SPEECH2TEXT,85            model_properties={},86            parameter_rules=[],87        )88 89        return entity90