a bgfQ@sddlZddlZddlZddlmZmZddlmZmZm Z m Z m Z m Z ddl Z ddlmZddlmZddlmZddlmZmZmZmZmZdgZGd d d eeZe eefed d d Ze eefed ddZe eefed ddZGdddeZ GdddeZ!edddZ"edddZ#edddZ$ee%dddZ&ee e%ed d!d"Z'eed#d$d%Z(ed&d'd(d)Gd*ddeZ)dS)+N)ABCabstractmethod)AnyCallableDictListMappingOptional) deprecated)CallbackManagerForLLMRun)LLM) BaseModel ConfigDictField PrivateAttrmodel_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$q/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/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%llm1sz_DatabricksClientBase.llm)N)__name__ __module__ __qualname____doc__str__annotations__rrr(r*rr rr.propertyboolr1r$r$r$r%rs  r)r#rcCs|dddS)Nchoicesrr!r$r#r$r$r%_transform_completions6sr<cCs|dddS)N candidatesrr!r$r;r$r$r%_transform_llama2_chat:sr>cCs|ddddS)Nr:rmessagecontentr$r;r$r$r%_transform_chat>srAcseZdZUdZeed<eed<eed<dZeed<dZe ed<dZ e eed <ed fd d Z e e d ddZeddeeeefedddZdee edefedddZZS) _DatabricksServingEndpointClientz:An API client that talks to a Databricks serving endpoint.host endpoint_namedatabricks_uriNclientFexternal_or_foundationtaskdatac stjfi|zddlm}||j|_Wn.ty\}ztd|WYd}~n d}~00|j|j}| dd dv|_ |j dur| d|_ dS)Nr)get_deploy_clientzMFailed to create the client. Please install mlflow with `pip install mlflow`.Z endpoint_type)Zexternal_modelZfoundation_model_apirH) super__init__Zmlflow.deploymentsrKrErF ImportErrorZ get_endpointrDgetlowerrGrH)r"rJrKeendpoint __class__r$r%rNLs  z)_DatabricksServingEndpointClient.__init__r/cCs |jdvS)N) llm/v1/chatllm/v1/completions llama2/chat)rHr0r$r$r%r1asz$_DatabricksServingEndpointClient.llmbeforemodevaluesrcCs6d|vr2|d}|d}d|d|d}||d<|S)NrrCrDhttps://z/serving-endpoints/z /invocationsr$)clsr]rCrDrr$r$r% set_api_urles 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)rSinputsrVrWZdataframe_recordsZ predictionsrrX) rGrFZpredictrDrHrAr< isinstancelistr>)r"rr,respZwrapped_requestr#predspredr$r$r%r.os$    z%_DatabricksServingEndpointClient.post)N)r2r3r4r5r6r7rFrrGr9rHr rNr8r1r classmethodrr`rr. __classcell__r$r$rTr%rBBs$    rBc@speZdZUdZeed<eed<eed<eddeeee fe ddd Z de e e d efe d d dZ d S)#_DatabricksClusterDriverProxyClientzBAn API client that talks to a Databricks cluster driver proxy app.rC cluster_idcluster_driver_portrYrZr\cCsBd|vr>|d}|d}|d}d|d|d|}||d<|S)NrrCrjrkr^z/driver-proxy-api/o/0//r$)r_r]rCrjportrr$r$r%r`sz/_DatabricksClusterDriverProxyClient.set_api_urlN.r+cCs||j|}|r||S|Sr-)r*r)r"rr,rdr$r$r%r.sz(_DatabricksClusterDriverProxyClient.post)N)r2r3r4r5r6r7rrgrrr`r rr.r$r$r$r%ris  rir/cCs4zddlm}|WSty.tdYn0dS)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_contextrorOrnr$r$r%get_repl_contexts  rpc Csttd}|sZztj}|s$tdWn2tyX}ztd|WYd}~n d}~00|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://rl)osgetenvrpZbrowserHostNamer Exceptionlstriprstrip)rCrRr$r$r%get_default_hosts  rvc Csbtd}r|Sztj}|s(tdWn2ty\}ztd|WYd}~n d}~00|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)rqrrrpZapiTokenrrs)rrRr$r$r%get_default_api_tokens rw)rJrcCs"t|tsdSd}tt||S)zJChecks if a data is a valid hexadecimal string using a regular expression.Fz^[0-9a-fA-F]+$)rbr6r9rematch)rJpatternr$r$r%_is_hex_strings r{)rJallow_dangerous_deserializationrc Cs|s tdz ddl}Wn2tyJ}ztd|WYd}~n d}~00z|t|WSty}ztd|WYd}~n d}~00dS)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 cloudpicklersloadsbytesfromhex)rJr|r~rRr$r$r% _load_pickled_fn_from_hex_strings  $r)fnrc Csz ddl}Wn2ty>}ztd|WYd}~n d}~00z||WSty}ztd|WYd}~n d}~00dS)z6Pickles a function and returns the hexadecimal string.rNr}zFailed to pickle the function: )r~rsrdumpshex)rr~rRr$r$r%_pickle_fn_to_hex_strings $rz0.3.3z1.0z#databricks_langchain.ChatDatabricks)ZsinceZremovalZalternative_importcseZdZUdZeedZeed<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<e%ddZ&e'eeefdddZ(e)dd e*eeefed!d"d#Z+ed$fd%d& Z,e'eeefdd'd(Z-e'e.eefdd)d*Z/e'edd+d,Z0d0ee eee e1eed-d.d/Z2Z3S)1ra 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_factoryrCrNrDrjrk model_kwargstransform_input_fn.r, databricksrEg temperaturenstop max_tokens extra_paramsrHFr|_clientZforbid)extrar/cCs6|j|jd}|jr|j|d<|jdur2|j|d<|S)N)rrrr)rrrr)r"paramsr$r$r% _llm_paramss   zDatabricks._llm_paramsrYrZr\c Cs|d}|d}|r&|r&tdnh|r4d|d<nZ|r:nTz tj}rP||d<tdWn2ty}ztd|WYd}~n d}~00|d}|r|rtdn<|rd|d<n.|durtdnt|d krtd |n|d }rd |vsJd d|vsJd|S)NrjrDz-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: rkz6Cannot set both endpoint_name and cluster_driver_port.zdnt |j |jdurTdnt |jdS)zReturn default params.N)rCrDrjrkrErrrrrrrHrr,)rCrDrjrkrErrrrrrrHrrr,r0r$r$r%_default_paramss&zDatabricks._default_paramscCs|jSr-)rr0r$r$r%_identifying_paramsszDatabricks._identifying_paramscCsdS)zReturn type of llm.rr$r0r$r$r% _llm_typeszDatabricks._llm_type)rr run_managerkwargsrcKsld|i}|jjr||j||jp*|j|||rD||d<|jrZ|jfi|}|jj||jdS)zAQueries the LLM endpoint with the given prompt and stop sequence.rr)r,) rr1updaterrrrr.r,)r"rrrrrr$r$r%_call$s   zDatabricks._call)NN)4r2r3r4r5rrvrCr6r7rwrrDr rjrkrrrrrr,rErfloatrrrrrdictrrHr|r9rrrrZ model_configr8rrrgrrNrrrrr rrhr$r$rTr%r sT 1       ,. )*rqrxrabcrrtypingrrrrrr rZlangchain_core._apir Zlangchain_core.callbacksr Zlangchain_core.language_modelsr Zpydanticr rrrr__all__rr6r<r>rArBrirprvrwr9r{rrrr$r$r$r%s<     I