a
    bgY                     @   s  d dl Z d dlZd dlZd dlZd dlZd dlZd dlZd dlmZ d dlmZ d dl	m
Z
 d dlmZmZmZmZmZmZmZmZmZ d dlmZmZmZ d dlmZ d dlmZ d d	lmZmZ d d
l m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+ d dl,m-Z-m.Z. d dl/m0Z0m1Z1m2Z2 d dl3m4Z4m5Z5 d dl6m7Z7 d dl8m9Z9m:Z:m;Z; e<e=Z>dZ?dZ@e#eAdddZBeeCef e#dddZDeeCef ee$ e$dddZEG dd deZFG dd dZGdS )    N)datetime)Queue)mktime)	AnyDict	GeneratorIteratorListMappingOptionalTypecast)	urlencodeurlparse
urlunparse)format_date_time)CallbackManagerForLLMRun)BaseChatModelgenerate_from_stream)	AIMessageAIMessageChunkBaseMessageBaseMessageChunkChatMessageChatMessageChunkFunctionMessageChunkHumanMessageHumanMessageChunkSystemMessageToolMessageChunk)make_invalid_tool_callparse_tool_call)ChatGenerationChatGenerationChunk
ChatResult)get_from_dict_or_envget_pydantic_field_names)
get_fields)
ConfigDictFieldmodel_validatorz$wss://spark-api.xf-yun.com/v3.5/chatzgeneralv3.5)messagereturnc                 C   s   t | trd| jd}nt | tr0d| jd}nt | trd| jd}d| jv rr| jd |d< |d dkrrd |d< d| jv r| jd |d< |d dkrd |d< n&t | trd| jd}ntd	|  |S )
Nuser)rolecontent	assistantfunction_callr/    
tool_callssystemzGot unknown type )
isinstancer   r/   r   r   additional_kwargsr   
ValueError)r+   Zmessage_dict r8   v/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/chat_models/sparkllm.pyconvert_message_to_dict9   s$    






r:   )_dictr,   c                 C   s  | d }| d }|dkr"t |dS |dkrg }i }| d }rLt||d< g }| d }r||d< | d D ]T}z|t|dd	 W qn ty }	 z |t|t|	 W Y d }	~	qnd }	~	0 0 qnni }|pd
}
t|
|||dS |dkrt	|dS t
||dS d S )Nr.   r/   r-   r/   r0   r1   r3   T)Z	return_idr2   )r/   r6   r3   invalid_tool_callsr4   r/   r.   )r   getdictappendr!   	Exceptionr    strr   r   r   )r;   msg_rolemsg_contentr=   r6   r1   r3   Zraw_tool_callsZraw_tool_caller/   r8   r8   r9   convert_dict_to_messageS   s<    

rG   )r;   default_classr,   c                 C   s  t t| d}t t| dp d}i }| drbt| d }d|v rZ|d d u rZd|d< ||d< | drx| d |d< |dks|tkrt|dS |d	ks|tkrt||d
S |dks|tkrt|| d dS |dks|tkrt|| d dS | s|tkr
t||dS ||dS d S )Nr.   r/   r2   r1   namer3   r-   r<   r0   )r/   r6   function)r/   rI   Ztooltool_call_id)r/   rK   r>   )	r   rC   r?   r@   r   r   r   r   r   )r;   rH   rD   rE   r6   r1   r8   r8   r9   _convert_delta_to_message_chunkv   s*    


rL   c                   @   s  e Zd ZU dZeedddZeee	e	f dddZ
dZeed< edd	d
Zee	 ed< eddd
Zee	 ed< eddd
Zee	 ed< eddd
Zee	 ed< eddd
Zee	 ed< dZe	ed< dZeed< edddZeed< eddZeed< dZeed < eed!Zee	ef ed"< ed#d$Ze d%d&eeed'd(d)Z!e d%d&eee	ef ed'd*d+Z"d4e#e$ ee#e	  ee% ee&e' d,d-d.Z(d5e#e$ ee#e	  ee% ee ee)d/d0d1Z*ee	dd2d3Z+dS )6ChatSparkLLMu-  IFlyTek Spark chat model integration.

    Setup:
        To use, you should have the environment variable``IFLYTEK_SPARK_API_KEY``,
        ``IFLYTEK_SPARK_API_SECRET`` and ``IFLYTEK_SPARK_APP_ID``.

    Key init args — completion params:
        model: Optional[str]
            Name of IFLYTEK SPARK model to use.
        temperature: Optional[float]
            Sampling temperature.
        top_k: Optional[float]
            What search sampling control to use.
        streaming: Optional[bool]
             Whether to stream the results or not.

    Key init args — client params:
        api_key: Optional[str]
            IFLYTEK SPARK API KEY. If not passed in will be read from env var IFLYTEK_SPARK_API_KEY.
        api_secret: Optional[str]
            IFLYTEK SPARK API SECRET. If not passed in will be read from env var IFLYTEK_SPARK_API_SECRET.
        api_url: Optional[str]
            Base URL for API requests.
        timeout: Optional[int]
            Timeout for requests.

    See full list of supported init args and their descriptions in the params section.

    Instantiate:
        .. code-block:: python

            from langchain_community.chat_models import ChatSparkLLM

            chat = ChatSparkLLM(
                api_key="your-api-key",
                api_secret="your-api-secret",
                model='Spark4.0 Ultra',
                # temperature=...,
                # other params...
            )

    Invoke:
        .. code-block:: python

            messages = [
                ("system", "你是一名专业的翻译家，可以将用户的中文翻译为英文。"),
                ("human", "我喜欢编程。"),
            ]
            chat.invoke(messages)

        .. code-block:: python

            AIMessage(
                content='I like programming.',
                response_metadata={
                    'token_usage': {
                        'question_tokens': 3,
                        'prompt_tokens': 16,
                        'completion_tokens': 4,
                        'total_tokens': 20
                    }
                },
                id='run-af8b3531-7bf7-47f0-bfe8-9262cb2a9d47-0'
            )

    Stream:
        .. code-block:: python

            for chunk in chat.stream(messages):
                print(chunk)

        .. code-block:: python

            content='I' id='run-fdbb57c2-2d32-4516-b894-6c5a67605d83'
            content=' like programming' id='run-fdbb57c2-2d32-4516-b894-6c5a67605d83'
            content='.' id='run-fdbb57c2-2d32-4516-b894-6c5a67605d83'

        .. code-block:: python

            stream = chat.stream(messages)
            full = next(stream)
            for chunk in stream:
                full += chunk
            full

        .. code-block:: python

            AIMessageChunk(
                content='I like programming.',
                id='run-aca2fa82-c2e4-4835-b7e2-865ddd3c46cb'
            )

    Response metadata
        .. code-block:: python

            ai_msg = chat.invoke(messages)
            ai_msg.response_metadata

        .. code-block:: python

            {
                'token_usage': {
                    'question_tokens': 3,
                    'prompt_tokens': 16,
                    'completion_tokens': 4,
                    'total_tokens': 20
                }
            }

    )r,   c                 C   s   dS )z9Return whether this model can be serialized by Langchain.Fr8   )clsr8   r8   r9   is_lc_serializable  s    zChatSparkLLM.is_lc_serializablec                 C   s   ddddddS )NIFLYTEK_SPARK_APP_IDIFLYTEK_SPARK_API_KEYIFLYTEK_SPARK_API_SECRETIFLYTEK_SPARK_API_URLIFLYTEK_SPARK_LLM_DOMAIN)spark_app_idspark_api_keyspark_api_secretspark_api_urlspark_llm_domainr8   selfr8   r8   r9   
lc_secrets  s    zChatSparkLLM.lc_secretsNclientapp_id)defaultaliasrU   api_keyrV   
api_secretrW   api_urlrX   modelrY   Zlc_userspark_user_idF	streaming   timeout)r`   request_timeoutg      ?r_   temperature   top_k)default_factorymodel_kwargsT)Zpopulate_by_namebefore)mode)valuesr,   c                 C   s   t |ddgd|d< t |ddgd|d< t |ddgd	|d< t |d
dt|d
< t |ddt|d< dd t|  D }|d|d d< |d|d d< t|d |d |d |d
 |d |d d|d< |S )NrU   r^   rP   rV   ra   rQ   rW   rb   rR   rX   rS   rY   rT   c                 S   s"   i | ]\}}|j d ur||j qS )Nrj   ).0rI   fieldr8   r8   r9   
<dictcomp>O  s   
z5ChatSparkLLM.validate_environment.<locals>.<dictcomp>rk   ro   rm   r^   ra   rb   rc   spark_domainro   r]   )r%   SPARK_API_URLSPARK_LLM_DOMAINr'   itemsr?   _SparkLLMClient)rN   rr   default_valuesr8   r8   r9   validate_environment/  sR    

z!ChatSparkLLM.validate_environmentc              
   C   s   t | }|di }t|D ]P}||v r8td| d||vrtd| d| d| d ||||< q|| }|rtd| d	||d< |S )
z>Build extra kwargs from additional params that were passed in.ro   zFound z supplied twice.z	WARNING! z/ is not default parameter.
                    zJ was transferred to model_kwargs.
                    Please confirm that z is what you intended.zParameters za should be specified explicitly. Instead they were passed in as part of `model_kwargs` parameter.)	r&   r?   listr7   loggerwarningpopintersectionkeys)rN   rr   Zall_required_field_namesextra
field_nameZinvalid_model_kwargsr8   r8   r9   build_extrae  s*    
zChatSparkLLM.build_extra)messagesstoprun_managerkwargsr,   c           
      k   s   t }| jjdd |D | j| jdd | jj| jdD ]H}d|vrDq6|d }t||}t|d}	|rx|j	t
|j|	d |	V  q6d S )	Nc                 S   s   g | ]}t |qS r8   r:   rs   mr8   r8   r9   
<listcomp>      z(ChatSparkLLM._stream.<locals>.<listcomp>T)rf   rh   datar+   )chunk)r   r]   arunre   ro   	subscriberi   rL   r#   Zon_llm_new_tokenrC   r/   )
r[   r   r   r   r   Zdefault_chunk_classr/   deltar   Zcg_chunkr8   r8   r9   _stream  s     

zChatSparkLLM._stream)r   r   r   streamr   r,   c                 K   s   |s
| j r*| jf |||d|}t|S | jdd |D | j| jd i }i }| jj| jdD ]*}	d|	v rz|	d |d< d|	vrqb|	d }qbt	|}
t
|
d	g}t||d
S )N)r   r   r   c                 S   s   g | ]}t |qS r8   r   r   r8   r8   r9   r     r   z*ChatSparkLLM._generate.<locals>.<listcomp>Fr   usageZtoken_usager   r   )generations
llm_output)rf   r   r   r]   r   re   ro   r   ri   rG   r"   r$   )r[   r   r   r   r   r   Zstream_iter
completionr   r/   r+   r   r8   r8   r9   	_generate  s0    

zChatSparkLLM._generatec                 C   s   dS )Nzspark-llm-chatr8   rZ   r8   r8   r9   	_llm_type  s    zChatSparkLLM._llm_type)NN)NNN),__name__
__module____qualname____doc__classmethodboolrO   propertyr   rC   r\   r]   r   __annotations__r)   rU   r   rV   rW   rX   rY   re   rf   ri   intrk   floatrm   r@   ro   r(   Zmodel_configr*   r}   r   r	   r   r   r   r#   r   r$   r   r   r8   r8   r8   r9   rM      s`   
o	4  
   
 rM   c                   @   s  e Zd ZdZd"eeeee ee ee dddZeeeeedddZ	d#e
e eee edd
ddZd$e
e eee eejd
ddZeee ddddZeeeddddZeddddZeeddddZd%eeee edddZd&ee eeddf dd d!ZdS )'r{   z
    Use websocket-client to call the SparkLLM interface provided by Xfyun,
    which is the iFlyTek's open platform for AI capabilities
    Nrv   c                 C   sv   zdd l }|| _W n ty,   tdY n0 |s6tn|| _|| _|| _|pNt| _t	 | _
ddd| _|| _|| _d S )Nr   zhCould not import websocket client python package. Please install it with `pip install websocket-client`.r2   r0   r>   )	websocketwebsocket_clientImportErrorrx   rc   r^   ro   ry   rw   r   queueblocking_messagera   rb   )r[   r^   ra   rb   rc   rw   ro   r   r8   r8   r9   __init__  s    	


z_SparkLLMClient.__init__)rc   ra   rb   r,   c                 C   s   t tt  }t| }|j}|j}d| d| d| d}tj	|
d|
dtjd }t|jdd}	d| d	|	 d
}
t|

djdd}|||d}t|}t|j|j|j|j||jf}|S )zK
        Generate a request url with an api key and an api secret.
        zhost: z
date: z
GET z	 HTTP/1.1zutf-8)	digestmod)encodingz	api_key="zQ", algorithm="hmac-sha256",         headers="host date request-line", signature="")authorizationdatehost)r   r   r   now	timetupler   netlocpathhmacnewencodehashlibsha256digestbase64	b64encodedecoder   r   schemeparamsfragment)rc   ra   rb   r   
parsed_urlr   r   Zsignature_originZsignature_shaZsignature_sha_base64Zauthorization_originr   params_dictZencoded_paramsurlr8   r8   r9   _create_url  s:    

z_SparkLLMClient._create_urlF)r   user_idro   rf   r,   c                 C   sl   | j d | j jt| j| j| j| j| j	| j
| jd}||_||_|d u rT| jn||_||_|  d S )NF)
on_messageon_erroron_closeon_open)r   ZenableTraceZWebSocketAppr{   r   rc   ra   rb   r   r   r   r   r   r   ro   rf   run_forever)r[   r   r   ro   rf   wsr8   r8   r9   run  s"    z_SparkLLMClient.runc                 C   s$   t j| j||||fd}|  |S )N)targetargs)	threadingThreadr   start)r[   r   r   ro   rf   Z	ws_threadr8   r8   r9   r   '  s    	z_SparkLLMClient.arun)r   errorr,   c                 C   s   | j d|i |  d S )Nr   )r   putclose)r[   r   r   r8   r8   r9   r   :  s    z_SparkLLMClient.on_error)r   close_status_codeclose_reasonr,   c                 C   s(   t d||di | jddi d S )Nlog)r   r   doneT)r   debugr   r   )r[   r   r   r   r8   r8   r9   r   >  s    z_SparkLLMClient.on_close)r   r,   c                 C   s6   ddd| _ t| j|j|j|jd}|| d S )Nr2   r0   r>   )r   r   ro   )r   jsondumps
gen_paramsr   r   ro   send)r[   r   r   r8   r8   r9   r   I  s    z_SparkLLMClient.on_open)r   r+   r,   c           	      C   s  t |}|d d }|dkrL| jdd| d|d d  i |  n|d d	 }|d
 }|d d d }|jr| jd|d d i n| jd  |7  < |dkr|js| jd| ji |r|di di di ni }| jd|i |  d S )Nheadercoder   r   zCode: z	, Error: r+   payloadchoicesstatustextr/   r      r   )r   loadsr   r   r   rf   r   r?   )	r[   r   r+   r   r   r   r   r/   Z
usage_datar8   r8   r9   r   R  s,    

z_SparkLLMClient.on_message)r   r   ro   r,   c                 C   sP   | j |ddd| jiidd|iid}|r<|d d | td|  |S )	N)r^   uidZchatdomainr+   r   )r   	parameterr   r   zSpark Request Parameters: )r^   rw   updater   r   )r[   r   r   ro   r   r8   r8   r9   r   m  s    

z_SparkLLMClient.gen_paramsrg   )rh   r,   c              
   c   s   z| j j|d}W n6 t jyH } ztd| dW Y d }~n
d }~0 0 d|v r^t|d d|v rn|V  q d|v rxqd|vrq|V  q d S )Nr   z-SparkLLMClient wait LLM api response timeout z secondsr   r   r   r   )r   r?   EmptyTimeoutErrorConnectionError)r[   rh   r/   _r8   r8   r9   r   {  s     
z_SparkLLMClient.subscribe)NNN)NF)NF)N)rg   )r   r   r   r   rC   r   r@   r   staticmethodr   r	   r   r   r   r   r   r   r   r   r   r   r   r   r~   r   r   r   r8   r8   r8   r9   r{     sT   
   0    	 
r{   )Hr   r   r   r   loggingr   r   r   r   timer   typingr   r   r   r   r	   r
   r   r   r   urllib.parser   r   r   Zwsgiref.handlersr   Zlangchain_core.callbacksr   Z*langchain_core.language_models.chat_modelsr   r   Zlangchain_core.messagesr   r   r   r   r   r   r   r   r   r   r   Z*langchain_core.output_parsers.openai_toolsr    r!   Zlangchain_core.outputsr"   r#   r$   Zlangchain_core.utilsr%   r&   Zlangchain_core.utils.pydanticr'   Zpydanticr(   r)   r*   	getLoggerr   r   rx   ry   r@   r:   rC   rG   rL   rM   r{   r8   r8   r8   r9   <module>   s@   ,4
$  0