from__future__importannotationsfromtypingimportAny,Dict,List,Mapping,Optionalfromlangchain_core.callbacksimport(AsyncCallbackManagerForLLMRun,CallbackManagerForLLMRun,)fromlangchain_core.language_models.llmsimportLLMfrompydanticimportBaseModel# Ignoring type because below is valid pydantic code# Unexpected keyword argument "extra" for "__init_subclass__" of "object"
[docs]classParams(BaseModel,extra="allow"):# type: ignore[call-arg]"""Parameters for the Javelin AI Gateway LLM."""temperature:float=0.0stop:Optional[List[str]]=Nonemax_tokens:Optional[int]=None
[docs]classJavelinAIGateway(LLM):"""Javelin AI Gateway LLMs. To use, you should have the ``javelin_sdk`` python package installed. For more information, see https://docs.getjavelin.io Example: .. code-block:: python from langchain_community.llms import JavelinAIGateway completions = JavelinAIGateway( gateway_uri="<your-javelin-ai-gateway-uri>", route="<your-javelin-ai-gateway-completions-route>", params={ "temperature": 0.1 } ) """route:str"""The route to use for the Javelin AI Gateway API."""client:Optional[Any]=None"""The Javelin AI Gateway client."""gateway_uri:Optional[str]=None"""The URI of the Javelin AI Gateway API."""params:Optional[Params]=None"""Parameters for the Javelin AI Gateway API."""javelin_api_key:Optional[str]=None"""The API key for the Javelin AI Gateway API."""def__init__(self,**kwargs:Any):try:fromjavelin_sdkimport(JavelinClient,UnauthorizedError,)exceptImportError:raiseImportError("Could not import javelin_sdk python package. ""Please install it with `pip install javelin_sdk`.")super().__init__(**kwargs)ifself.gateway_uri:try:self.client=JavelinClient(base_url=self.gateway_uri,api_key=self.javelin_api_key)exceptUnauthorizedErrorase:raiseValueError("Javelin: Incorrect API Key.")frome@propertydef_default_params(self)->Dict[str,Any]:"""Get the default parameters for calling Javelin AI Gateway API."""params:Dict[str,Any]={"gateway_uri":self.gateway_uri,"route":self.route,"javelin_api_key":self.javelin_api_key,**(self.params.dict()ifself.paramselse{}),}returnparams@propertydef_identifying_params(self)->Mapping[str,Any]:"""Get the identifying parameters."""returnself._default_paramsdef_call(self,prompt:str,stop:Optional[List[str]]=None,run_manager:Optional[CallbackManagerForLLMRun]=None,**kwargs:Any,)->str:"""Call the Javelin AI Gateway API."""data:Dict[str,Any]={"prompt":prompt,**(self.params.dict()ifself.paramselse{}),}ifs:=(stopor(self.params.stopifself.paramselseNone)):data["stop"]=sifself.clientisnotNone:resp=self.client.query_route(self.route,query_body=data)else:raiseValueError("Javelin client is not initialized.")resp_dict=resp.dict()try:returnresp_dict["llm_response"]["choices"][0]["text"]exceptKeyError:return""asyncdef_acall(self,prompt:str,stop:Optional[List[str]]=None,run_manager:Optional[AsyncCallbackManagerForLLMRun]=None,**kwargs:Any,)->str:"""Call async the Javelin AI Gateway API."""data:Dict[str,Any]={"prompt":prompt,**(self.params.dict()ifself.paramselse{}),}ifs:=(stopor(self.params.stopifself.paramselseNone)):data["stop"]=sifself.clientisnotNone:resp=awaitself.client.aquery_route(self.route,query_body=data)else:raiseValueError("Javelin client is not initialized.")resp_dict=resp.dict()try:returnresp_dict["llm_response"]["choices"][0]["text"]exceptKeyError:return""@propertydef_llm_type(self)->str:"""Return type of llm."""return"javelin-ai-gateway"