U h@srddlZddlmZmZmZmZddlmZddlm Z ddl m Z m Z ddl mZeeZGddde ZdS) N)AnyDictListOptional)CallbackManagerForLLMRun)LLM)get_from_dict_or_envpre_init)enforce_stop_tokensc@s eZdZUdZeed<dZeeed<dZ ee eefed<dZ e ed<d Z eed <dZeeed <dZeeeed <Gd ddZee e dddZee eefdddZee eefdddZeedddZdeeeeeeeedddZdS)PredictionGuardaPrediction Guard large language models. To use, you should have the ``predictionguard`` python package installed, and the environment variable ``PREDICTIONGUARD_TOKEN`` set with your access token, or pass it as a named parameter to the constructor. To use Prediction Guard's API along with OpenAI models, set the environment variable ``OPENAI_API_KEY`` with your OpenAI API key as well. Example: .. code-block:: python pgllm = PredictionGuard(model="MPT-7B-Instruct", token="my-access-token", output={ "type": "boolean" }) clientzMPT-7B-InstructmodelNoutput max_tokensg? temperaturetokenstopc@seZdZdZdS)zPredictionGuard.ConfigZforbidN)__name__ __module__ __qualname__extrarrL/tmp/pip-unpacked-wheel-9gdii04g/langchain_community/llms/predictionguard.pyConfig2sr)valuesreturncCsJt|dd}zddl}|j|d|d<Wntk rDtdYnX|S)zHValidate that the access token and python package exists in environment.rZPREDICTIONGUARD_TOKENrN)rr zfCould not import predictionguard python package. Please install it with `pip install predictionguard`.)rpredictionguardZClient ImportError)clsrrpgrrrvalidate_environment5s  z$PredictionGuard.validate_environment)rcCs|j|jdS)z@Get the default parameters for calling the Prediction Guard API.rrr"selfrrr_default_paramsDszPredictionGuard._default_paramscCsd|ji|jS)zGet the identifying parameters.r )r r%r#rrr_identifying_paramsLsz#PredictionGuard._identifying_paramscCsdS)zReturn type of llm.rrr#rrr _llm_typeQszPredictionGuard._llm_type)promptr run_managerkwargsrc Ksddl}|j}|jdk r*|dk r*tdn|jdk r@|j|d<n||d<|jjf|j||j|d|dd|}|ddd }|dk s|jdk rt||d}|S) a&Call out to Prediction Guard's model API. Args: prompt: The prompt to pass into the model. Returns: The string generated by the model. Example: .. code-block:: python response = pgllm.invoke("Tell me a joke.") rNz2`stop` found in both the input and default params.Zstop_sequencesrr)r r(rrrchoicestext) rr%r ValueErrorZ Completioncreater rr ) r$r(rr)r*r paramsresponser,rrr_callVs(   zPredictionGuard._call)NN)rrr__doc__r__annotations__r rstrrrrintrfloatrrrrr r!propertyr%r&r'rr1rrrrr s4    r )loggingtypingrrrrZlangchain_core.callbacksrZ#langchain_core.language_models.llmsrZlangchain_core.utilsrr Zlangchain_community.llms.utilsr getLoggerrloggerr rrrrs