U hCP@sxddlZddlZddlZddlmZmZddlmZmZm Z m Z m Z m Z ddl Z ddlmZddlmZddlmZmZmZmZmZdgZGdd d eeZe eefed d d Ze eefed d dZe eefed ddZGdddeZGdddeZedddZ edddZ!edddZ"ee#dddZ$ee e#edd d!Z%eed"d#d$Z&Gd%ddeZ'dS)&N)ABCabstractmethod)AnyCallableDictListMappingOptional)CallbackManagerForLLMRun)LLM) BaseModelField PrivateAttrroot_validator validator Databricksc@seZdZUdZeed<eed<eeeedddZeeddd Zeeed d d Z e dee e defedddZ eedddZd S)_DatabricksClientBasez0A base JSON API client that talks to Databricks.api_url api_token)methodurlrequestreturncCsHdd|ji}tj||||d}|js@td|jd|j|S)N AuthorizationzBearer )rrheadersjsonzHTTP z error: )rrequestsrok ValueError status_codetextr)selfrrrrresponser#G/tmp/pip-unpacked-wheel-9gdii04g/langchain_community/llms/databricks.pyrsz_DatabricksClientBase.request)rrcCs|d|dS)NGETr)r!rr#r#r$_get%sz_DatabricksClientBase._get)rrrcCs|d||S)NPOSTr&)r!rrr#r#r$_post(sz_DatabricksClientBase._postN.rtransform_output_fnrcCsdSNr#)r!rr+r#r#r$post+sz_DatabricksClientBase.postrcCsdS)NFr#r!r#r#r$llm0sz_DatabricksClientBase.llm)N)__name__ __module__ __qualname____doc__str__annotations__rrr'r)rr rr-propertyboolr0r#r#r#r$rs  r)r"rcCs|dddS)Nchoicesrr r#r"r#r#r$_transform_completions5sr;cCs|dddS)N candidatesrr r#r:r#r#r$_transform_llama2_chat9sr=cCs|ddddS)Nr9rmessagecontentr#r:r#r#r$_transform_chat=sr@cseZdZUdZeed<eed<eed<dZeed<dZe ed<dZ e eed <ed fd d Z e e d ddZeddeeefeeefdddZdee edefedddZZS) _DatabricksServingEndpointClientz:An API client that talks to a Databricks serving endpoint.host endpoint_namedatabricks_uriNclientFexternal_or_foundationtaskdatac stjf|zddlm}||j|_Wn,tk rV}ztd|W5d}~XYnX|j|j}| dd dk|_ |j dkr| d|_ dS)Nr)get_deploy_clientzMFailed to create the client. Please install mlflow with `pip install mlflow`.Z endpoint_type)Zexternal_modelZfoundation_model_apirG) super__init__Zmlflow.deploymentsrJrDrE ImportErrorZ get_endpointrCgetlowerrFrG)r!rIrJeendpoint __class__r#r$rMKs  z)_DatabricksServingEndpointClient.__init__r.cCs |jdkS)N) llm/v1/chatllm/v1/completions llama2/chat)rGr/r#r#r$r0`sz$_DatabricksServingEndpointClient.llmTprevaluesrcCs6d|kr2|d}|d}d|d|d}||d<|S)NrrBrChttps://z/serving-endpoints/z /invocationsr#)clsr[rBrCrr#r#r$ set_api_urlds z,_DatabricksServingEndpointClient.set_api_url.r*cCs|jrL|jj|j|d}|r$||S|jdkr6t|S|jdkrHt|S|Sd|gi}|jj|j|d}|d}t|tr|dn|}|jdkrt |S|r||S|SdS)N)rRinputsrUrVZdataframe_recordsZ predictionsrrW) rFrEZpredictrCrGr@r; isinstancelistr=)r!rr+respZwrapped_requestr"predspredr#r#r$r-ms&    z%_DatabricksServingEndpointClient.post)N)r1r2r3r4r5r6rErrFr8rGr rMr7r0rrr^rr- __classcell__r#r#rSr$rAAs$   " rAc@steZdZUdZeed<eed<eed<eddeeefeeefddd Z dee e d efed d dZ d S)#_DatabricksClusterDriverProxyClientzBAn API client that talks to a Databricks cluster driver proxy app.rB cluster_idcluster_driver_portTrXrZcCsBd|kr>|d}|d}|d}d|d|d|}||d<|S)NrrBrgrhr\z/driver-proxy-api/o/0//r#)r]r[rBrgportrr#r#r$r^sz/_DatabricksClusterDriverProxyClient.set_api_urlN.r*cCs||j|}|r||S|Sr,)r)r)r!rr+rbr#r#r$r-sz(_DatabricksClusterDriverProxyClient.post)N) r1r2r3r4r5r6rrrr^r rr-r#r#r#r$rfs " rfr.cCs6zddlm}|WStk r0tdYnXdS)zgGet the notebook REPL context if running inside a Databricks notebook. Returns None otherwise. r get_contextzBCannot access dbruntime, not running inside a Databricks notebook.N)Z!dbruntime.databricks_repl_contextrlrNrkr#r#r$get_repl_contexts rmc Csrtd}|sXztj}|s$tdWn0tk rV}ztd|W5d}~XYnX|ddd}|S)z{Get the default Databricks workspace hostname. Raises an error if the hostname cannot be automatically determined. ZDATABRICKS_HOSTz(context doesn't contain browserHostName.zshost was not set and cannot be automatically inferred. Set environment variable 'DATABRICKS_HOST'. Received error: Nr\zhttp://ri)osgetenvrmZbrowserHostNamer Exceptionlstriprstrip)rBrQr#r#r$get_default_hosts  rsc Cs`td}r|Sztj}|s(tdWn0tk rZ}ztd|W5d}~XYnX|S)z{Get the default Databricks personal access token. Raises an error if the token cannot be automatically determined. ZDATABRICKS_TOKENz!context doesn't contain apiToken.zyapi_token was not set and cannot be automatically inferred. Set environment variable 'DATABRICKS_TOKEN'. Received error: N)rnrormZapiTokenrrp)rrQr#r#r$get_default_api_tokens rt)rIrcCs"t|tsdSd}tt||S)zJChecks if a data is a valid hexadecimal string using a regular expression.Fz^[0-9a-fA-F]+$)r`r5r8rematch)rIpatternr#r#r$_is_hex_strings rx)rIallow_dangerous_deserializationrc Cs|s tdz ddl}Wn0tk rH}ztd|W5d}~XYnXz|t|WStk r}ztd|W5d}~XYnXdS)z3Loads a pickled function from a hexadecimal string.aWThis code relies on the pickle module. You will need to set allow_dangerous_deserialization=True if you want to opt-in to allow deserialization of data using pickle.Data can be compromised by a malicious actor if not handled properly to include a malicious payload that when deserialized with pickle can execute arbitrary code on your machine.rN*Please install cloudpickle>=2.0.0. Error: zFFailed to load the pickled function from a hexadecimal string. Error: )r cloudpicklerploadsbytesfromhex)rIryr{rQr#r#r$ _load_pickled_fn_from_hex_strings   r)fnrc Csz ddl}Wn0tk r<}ztd|W5d}~XYnXz||WStk r~}ztd|W5d}~XYnXdS)z6Pickles a function and returns the hexadecimal string.rNrzzFailed to pickle the function: )r{rprdumpshex)rr{rQr#r#r$_pickle_fn_to_hex_strings  rcsbeZdZUdZeedZeed<ee dZ eed<dZ e eed<dZ e eed<dZe eed<dZe eeefed <dZe eed <dZe ed efed <d Zeed<dZeed<dZeed<dZe eeed<dZe eed<eedZeeefed<dZe eed<dZ e!ed<e"Z#e$ed<GdddZ%e&eeefdddZ'e(ddd eeeefe ed!d"d#Z)e(ddd eeeefe ed!d$d%Z*e(d dd e eeefe eeefd&d'd(Z+ed)fd*d+ Z,e&eeefdd,d-Z-e&e.eefdd.d/Z/e&edd0d1Z0d5ee eee e1eed2d3d4Z2Z3S)6ra Databricks serving endpoint or a cluster driver proxy app for LLM. It supports two endpoint types: * **Serving endpoint** (recommended for both production and development). We assume that an LLM was deployed to a serving endpoint. To wrap it as an LLM you must have "Can Query" permission to the endpoint. Set ``endpoint_name`` accordingly and do not set ``cluster_id`` and ``cluster_driver_port``. If the underlying model is a model registered by MLflow, the expected model signature is: * inputs:: [{"name": "prompt", "type": "string"}, {"name": "stop", "type": "list[string]"}] * outputs: ``[{"type": "string"}]`` If the underlying model is an external or foundation model, the response from the endpoint is automatically transformed to the expected format unless ``transform_output_fn`` is provided. * **Cluster driver proxy app** (recommended for interactive development). One can load an LLM on a Databricks interactive cluster and start a local HTTP server on the driver node to serve the model at ``/`` using HTTP POST method with JSON input/output. Please use a port number between ``[3000, 8000]`` and let the server listen to the driver IP address or simply ``0.0.0.0`` instead of localhost only. To wrap it as an LLM you must have "Can Attach To" permission to the cluster. Set ``cluster_id`` and ``cluster_driver_port`` and do not set ``endpoint_name``. The expected server schema (using JSON schema) is: * inputs:: {"type": "object", "properties": { "prompt": {"type": "string"}, "stop": {"type": "array", "items": {"type": "string"}}}, "required": ["prompt"]}` * outputs: ``{"type": "string"}`` If the endpoint model signature is different or you want to set extra params, you can use `transform_input_fn` and `transform_output_fn` to apply necessary transformations before and after the query. )default_factoryrBrNrCrgrh model_kwargstransform_input_fn.r+ databricksrDg temperaturenstop max_tokens extra_paramsrGFry_clientc@seZdZdZdZdS)zDatabricks.ConfigZforbidTN)r1r2r3extraZunderscore_attrs_are_privater#r#r#r$Configsrr.cCs6|j|jd}|jr|j|d<|jdk r2|j|d<|S)N)rrrr)rrrr)r!paramsr#r#r$ _llm_paramss   zDatabricks._llm_paramsT)always)vr[rc Cs~|r|drtdnd|dr"dS|r*|Sztj}r>|WStdWn0tk rx}ztd|W5d}~XYnXdS)NrCz-Cannot set both endpoint_name and cluster_id.z"Context doesn't contain clusterId.zuNeither endpoint_name nor cluster_id was set. And the cluster_id cannot be automatically determined. Received error: )rrmZ clusterIdrp)r]rr[rQr#r#r$set_cluster_ids    zDatabricks.set_cluster_idcCsX|r|drtdn>|dr"dS|dkr4tdn t|dkrPtd|n|SdS)NrCz6Cannot set both endpoint_name and cluster_driver_port.zdnt |j |jdkrTdnt |jdS)zReturn default params.N)rBrCrgrhrDrrrrrrrGrr+)rBrCrgrhrDrrrrrrrGrrr+r/r#r#r$_default_paramss&zDatabricks._default_paramscCs|jSr,)rr/r#r#r$_identifying_paramsszDatabricks._identifying_paramscCsdS)zReturn type of llm.rr#r/r#r#r$ _llm_typeszDatabricks._llm_type)rr run_managerkwargsrcKshd|i}|jjr||j||jp*|j|||rD||d<|jrV|jf|}|jj||jdS)zAQueries the LLM endpoint with the given prompt and stop sequence.rr)r+) rr0updaterrrrr-r+)r!rrrrrr#r#r$_calls    zDatabricks._call)NN)4r1r2r3r4r rsrBr5r6rtrrCr rgrhrrrrrr+rDrfloatrrrrrdictrrGryr8rrrrr7rrrrrrMrrrrr rrer#r#rSr$rsV 1           *. )(rnrurabcrrtypingrrrrrr rZlangchain_core.callbacksr Zlangchain_core.language_modelsr Zlangchain_core.pydantic_v1r r rrr__all__rr5r;r=r@rArfrmrsrtr8rxrrrr#r#r#r$s2    H