a
    bg>%                     @   s   d dl Z d dlZd dl mZ d dlmZmZmZmZmZm	Z	 d dl
Z
d dlmZ d dlmZmZ d dlmZmZmZ d dl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m Z m!Z! e"e"dddZ#ee ee"ef dddZ$G dd deZ%dS )    N)JSONDecodeError)AnyAsyncIteratorDictIteratorListOptional)ClientSession)AsyncCallbackManagerForLLMRunCallbackManagerForLLMRun)BaseChatModelagenerate_from_streamgenerate_from_stream)	AIMessageAIMessageChunkBaseMessage)ChatGenerationChatGenerationChunk
ChatResult)convert_to_secret_str)
ConfigDictField	SecretStr)rolereturnc                 C   s.   dddd}| |v r||  S t d|  d S )NZ	assistanthuman)Zair   ZchatzUnknown role type: )
ValueError)r   map r   |/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/chat_models/symblai_nebula.py_convert_role   s    r    )messagesr   c                 C   s~   d}g }| d d D ]D}|j dkr@t|jtr6|j}qXtdq|t|j |jd q| d j}|d|d ||dS )N systemzSystem prompt must be a string)r   textr   )Zsystem_promptr!   )type
isinstancecontentstrr   appendr    )r!   r$   Zformatted_messagesmessager%   r   r   r   _format_nebula_messages   s    


r,   c                       sN  e Zd ZU dZdZeed< dZee	 ed< dZ
eed< dZeed	< ed
ddZee ed< edddZed
d fddZeedddZeedddZd!ee eee  ee eee dddZd"ee eee  ee eee dddZ d#ee eee  ee ee!dddZ"d$ee eee  ee ee!ddd Z#  Z$S )%
ChatNebulaam  `Nebula` chat large language model - https://docs.symbl.ai/docs/nebula-llm

    API Reference: https://docs.symbl.ai/reference/nebula-chat

    To use, set the environment variable ``NEBULA_API_KEY``,
    or pass it as a named parameter to the constructor.
    To request an API key, visit https://platform.symbl.ai/#/login
    Example:
        .. code-block:: python

            from langchain_community.chat_models import ChatNebula
            from langchain_core.messages import SystemMessage, HumanMessage

            chat = ChatNebula(max_new_tokens=1024, temperature=0.5)

            messages = [
            SystemMessage(
                content="You are a helpful assistant."
            ),
            HumanMessage(
                "Answer the following question. How can I help save the world."
            ),
            ]
            chat.invoke(messages)
    i   max_new_tokensr   temperatureF	streamingzhttps://api-nebula.symbl.ainebula_api_urlNzNebula API Token)descriptionnebula_api_keyT)Zpopulate_by_nameZarbitrary_types_allowed)kwargsr   c                    sP   d|v rt |d}ndtjv r2t tjd }nd }t jf d|i| d S )Nr3   ZNEBULA_API_KEY)r   poposenvironsuper__init__)selfr4   Zapi_key	__class__r   r   r9   a   s    
zChatNebula.__init__)r   c                 C   s   dS )zReturn type of chat model.znebula-chatr   r:   r   r   r   	_llm_typej   s    zChatNebula._llm_typec                 C   s   | j r| j  S dS )Nr"   )r3   Zget_secret_valuer=   r   r   r   _api_keyo   s    
zChatNebula._api_key)r!   stoprun_managerr4   r   c              	   k   s   | j  d}| jdd}t|d}| j| jd||}dd | D }t|}	tj	d|||	d	d
}
|

  |
 D ]j}| dd }zt|}W n ty   Y qxY n0 |d }tt|dd}|r|j||d |V  qxdS )z#Call out to Nebula's chat endpoint./v1/model/chat/streamingapplication/jsonZApiKeyzContent-Typer!   r.   r/   c                 S   s   i | ]\}}|d ur||qS Nr   .0kvr   r   r   
<dictcomp>       z&ChatNebula._stream.<locals>.<dictcomp>POSTT)headersdatastream   Ndeltar(   r+   chunk)r1   r?   r,   r.   r/   itemsjsondumpsrequestsrequestraise_for_status
iter_linesdecodeloadsr   r   r   on_llm_new_token)r:   r!   r@   rA   r4   urlrO   formatted_datapayloadjson_payloadresponsechunk_responsechunk_decodedrW   tokencg_chunkr   r   r   _streamu   s:    



zChatNebula._streamc                 K  sl  | j  d}| jdd}t|d}| j| jd||}dd | D }t|}	t 4 I d H }
|
j	||	|dd	4 I d H }|
  |j2 zz3 d H W }| d
d  }zt|}W n ty   Y qY n0 |d }tt|dd}|r|j||dI d H  |V  q6 W d   I d H  q:1 I d H s00    Y  W d   I d H  qh1 I d H s^0    Y  d S )NrB   rC   rD   rE   rF   c                 S   s   i | ]\}}|d ur||qS rG   r   rH   r   r   r   rL      rM   z'ChatNebula._astream.<locals>.<dictcomp>T)rP   rO   rQ   rR   rS   rT   rU   rV   )r1   r?   r,   r.   r/   rX   rY   rZ   r	   postr]   r(   r_   r`   r   r   r   ra   )r:   r!   r@   rA   r4   rb   rO   rc   rd   re   sessionrf   rg   rh   rW   ri   rj   r   r   r   _astream   s>    


zChatNebula._astreamc                 K   s   | j r&| j|f||d|}t|S | j d}| jdd}t|d}| j| jd||}	dd |	 D }	t	
|	}
tjd	|||
d
}|  |	 }ttt|d ddg|dS )Nr@   rA   /v1/model/chatrC   rD   rE   rF   c                 S   s   i | ]\}}|d ur||qS rG   r   rH   r   r   r   rL      rM   z(ChatNebula._generate.<locals>.<dictcomp>rN   )rO   rP   r!   rT   rU   ZgenerationsZ
llm_output)r0   rk   r   r1   r?   r,   r.   r/   rX   rY   rZ   r[   r\   r]   r   r   r   )r:   r!   r@   rA   r4   stream_iterrb   rO   rc   rd   re   rf   rP   r   r   r   	_generate   s8    

zChatNebula._generatec              
      sP  | j r,| j|f||d|}t|I d H S | j d}| jdd}t|d}| j| jd||}	dd |	 D }	t	
|	}
t 4 I d H }|j||
|d	4 I d H ^}|  |	 I d H }ttt|d
 ddg|dW  d   I d H  W  d   I d H  S 1 I d H s0    Y  W d   I d H  qL1 I d H sB0    Y  d S )Nro   rp   rC   rD   rE   rF   c                 S   s   i | ]\}}|d ur||qS rG   r   rH   r   r   r   rL      rM   z)ChatNebula._agenerate.<locals>.<dictcomp>)rP   rO   r!   rT   rU   rq   )r0   rn   r   r1   r?   r,   r.   r/   rX   rY   rZ   r	   rl   r]   r   r   r   )r:   r!   r@   rA   r4   rr   rb   rO   rc   rd   re   rm   rf   rP   r   r   r   
_agenerate   sB    

zChatNebula._agenerate)NN)NN)NN)NN)%__name__
__module____qualname____doc__r.   int__annotations__r/   r   floatr0   boolr1   r)   r   r3   r   r   Zmodel_configr   r9   propertyr>   r?   r   r   r   r   r   rk   r
   r   rn   r   rs   rt   __classcell__r   r   r;   r   r-   5   sf   
	  
,  
*  
&  
r-   )&rY   r6   r   typingr   r   r   r   r   r   r[   Zaiohttpr	   Zlangchain_core.callbacksr
   r   Z*langchain_core.language_models.chat_modelsr   r   r   Zlangchain_core.messagesr   r   r   Zlangchain_core.outputsr   r   r   Zlangchain_core.utilsr   Zpydanticr   r   r   r)   r    r,   r-   r   r   r   r   <module>   s    