a
    dg                     @  s   d Z ddlmZ ddlZddlmZmZmZmZm	Z	 ddl
mZ ddlmZ ddlmZmZ erhddlZdd	d
dZG dd deZdS )z6Experimental implementation of jsonformer wrapped LLM.    )annotationsN)TYPE_CHECKINGAnyListOptionalcast)HuggingFacePipeline)CallbackManagerForLLMRun)Fieldmodel_validator
jsonformer)returnc                  C  s,   zddl } W n ty&   tdY n0 | S )z(Lazily import of the jsonformer package.r   Nz\Could not import jsonformer python package. Please install it with `pip install jsonformer`.)r   ImportError)r    r   |/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_experimental/llms/jsonformer_decoder.pyimport_jsonformer   s    
r   c                   @  s   e Zd ZU dZedddZded< eddd	Zd
ed< eddd	Zded< e	dde
dddddZdddddddddZdS )
JsonFormerzrJsonformer wrapped LLM using HuggingFace Pipeline API.

    This pipeline is experimental and not yet stable.
    .zThe JSON Schema to complete.)descriptiondictjson_schema   z)Maximum number of new tokens to generate.)defaultr   intmax_new_tokensFzDebug mode.booldebugbefore)moder   )valuesr   c                 C  s
   t   |S )N)r   )clsr   r   r   r   check_jsonformer_installation(   s    z(JsonFormer.check_jsonformer_installationNstrzOptional[List[str]]z"Optional[CallbackManagerForLLMRun])promptstoprun_managerkwargsr   c           
      K  sN   t  }ddlm} t|| j}|j|j|j| j|| j	| j
d}| }	t|	S )Nr   )Text2TextGenerationPipeline)model	tokenizerr   r"   Zmax_number_tokensr   )r   Ztransformersr&   r   pipelineZ
Jsonformerr'   r(   r   r   r   jsondumps)
selfr"   r#   r$   r%   r   r&   r)   r'   textr   r   r   _call.   s    zJsonFormer._call)NN)__name__
__module____qualname____doc__r
   r   __annotations__r   r   r   classmethodr    r.   r   r   r   r   r      s   
  r   )r2   
__future__r   r*   typingr   r   r   r   r   Z-langchain_community.llms.huggingface_pipeliner   Z langchain_core.callbacks.managerr	   Zpydanticr
   r   r   r   r   r   r   r   r   <module>   s   