U h5 @sddlZddlZddlmZmZmZmZmZmZm Z ddl Z ddl m Z ddl mZddlmZddlmZddlmZdZejGd d d Zeed d d ZdeeeeeeeeeeffdddZGdddeZdS)N)AnyDictListMappingOptionalUnioncast)CallbackManagerForLLMRun)LLM)root_validator)get_from_dict_or_env)enforce_stop_tokens<c@sBeZdZUdZeed<eed<ddddZedddd ZdS) AviaryBackendzAviary backend. Attributes: backend_url: The URL for the Aviary backend. bearer: The bearer token for the Aviary backend. backend_urlbearerNreturncCsd|ji|_dS)N Authorization)rheaderselfrC/tmp/pip-unpacked-wheel-9gdii04g/langchain_community/llms/aviary.py __post_init__szAviaryBackend.__post_init__cCsTtd}|stdtdd}|r0d|nd}||dsDdnd7}|||S)N AVIARY_URLzAVIARY_URL must be set AVIARY_TOKENzBearer /)osgetenvAssertionErrorendswith)cls aviary_url aviary_tokenrrrrfrom_envs    zAviaryBackend.from_env) __name__ __module__ __qualname____doc__str__annotations__r classmethodr&rrrrrs rrc Cst}|jd}tj||jtd}z |}Wn<tjk rl}zt d|d|j |W5d}~XYnXt dd| D}|S)zList available modelsz-/routes)headerstimeoutError decoding JSON from . Text response: NcSs&g|]}d|kr|dddqS)--r)lstripreplace).0krrr 8szget_models..) rr&rrequestsgetrTIMEOUTjsonJSONDecodeError RuntimeErrortextsortedkeys)backend request_urlresponseresulterrr get_models,s  rFTr)modelpromptuse_prompt_formatversionrc Cst}|j|ddd|d}tj||j||dtd}z |WStj k r}zt d|d|j |W5d}~XYnXdS) z#Get completions from Aviary models.rr2query)rHrI)r.r;r/r0r1N) rr&rr4r8postrr:r;r<r=r>)rGrHrIrJrAurlrCrErrrget_completions=s  rNc@seZdZUdZdZeed<dZeeed<dZ eeed<dZ e ed<dZ eeed <Gd d d Z edd eed ddZeeeefdddZeedddZdeeeeeeeedddZdS)AviaryabAviary hosted models. Aviary is a backend for hosted models. You can find out more about aviary at http://github.com/ray-project/aviary To get a list of the models supported on an aviary, follow the instructions on the website to install the aviary CLI and then use: `aviary models` AVIARY_URL and AVIARY_TOKEN environment variables must be set. Attributes: model: The name of the model to use. Defaults to "amazon/LightGPT". aviary_url: The URL for the Aviary backend. Defaults to None. aviary_token: The bearer token for the Aviary backend. Defaults to None. use_prompt_format: If True, the prompt template for the model will be ignored. Defaults to True. version: API version to use for Aviary. Defaults to None. Example: .. code-block:: python from langchain_community.llms import Aviary os.environ["AVIARY_URL"] = "" os.environ["AVIARY_TOKEN"] = "" light = Aviary(model='amazon/LightGPT') output = light('How do you make fried rice?') zamazon/LightGPTrGNr$r%TrIrJc@seZdZdZdS)z Aviary.ConfigZforbidN)r'r(r)extrarrrrConfig}srQ)pre)valuesrc Cst|dd}t|dd}|tjd<|tjd<z t}Wn.tjjk rd}z t|W5d}~XYnX|d}|r||krt|d|dd|S) z?Validate that api key and python package exists in environment.r$rr%rNrGz does not support model .) r renvironrFr8 exceptionsRequestException ValueErrorr9)r#rSr$r%Z aviary_modelsrErGrrrvalidate_environments       zAviary.validate_environmentrcCs|j|jdS)zGet the identifying parameters.)Z model_namer$)rGr$rrrr_identifying_paramsszAviary._identifying_paramscCsd|jddS)zReturn type of llm.zaviary-r-)rGr4rrrr _llm_typeszAviary._llm_type)rHstop run_managerkwargsrcKsPd|ji}|jr|j|d<tf|j|d|}tt|d}|rLt||}|S)a Call out to Aviary Args: prompt: The prompt to pass into the model. Returns: The string generated by the model. Example: .. code-block:: python response = aviary("Tell me a joke.") rIrJ)rGrHZgenerated_text)rIrJrNrGrr+r )rrHr]r^r_outputr>rrr_calls   z Aviary._call)NN)r'r(r)r*rGr+r,r$rr%rIboolrJrQr rrYpropertyrrrZr\rr rarrrrrOUs,    rO)Tr)Z dataclassesrtypingrrrrrrrr8Zlangchain_core.callbacksr Z#langchain_core.language_models.llmsr Zlangchain_core.pydantic_v1r Zlangchain_core.utilsr Zlangchain_community.llms.utilsr r:Z dataclassrr+rFrbfloatintrNrOrrrrs,$