
    7|h'                       d Z ddlmZ ddlZddlZddl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ZddlmZ ddlmZ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 dd
l m!Z! ddl"m#Z#m$Z$m%Z% ddl&m'Z' dZ(dZ) ejT                  e+      Z, G d de#      Z-y)zKonkoAI chat wrapper.    )annotationsN)	AnyDictIteratorListOptionalSetTupleUnioncast)CallbackManagerForLLMRun)AIMessageChunkBaseMessage)ChatGenerationChunk
ChatResult)convert_to_secret_strget_from_dict_or_envpre_init)Field	SecretStr)convert_message_to_dict)
ChatOpenAI_convert_delta_to_message_chunkgenerate_from_stream)is_openai_v1zhttps://api.konko.ai/v1zmeta-llama/Llama-2-13b-chat-hfc                      e Zd ZU dZed(d       Zed)d       ZdZde	d<    e
ed	      Zd
e	d<   	 dZde	d<   	  e
e      Zde	d<   	 dZde	d<   dZde	d<   dZde	d<   	 dZde	d<   	 dZde	d<   	 dZde	d<   	 ed*d       Zed+d       Zeddef	 	 	 	 	 	 	 d,d       Z	 d-	 	 	 	 	 d.d Z	 	 d/	 	 	 	 	 	 	 	 	 d0d!Z	 	 	 d1	 	 	 	 	 	 	 	 	 	 	 d2d"Z	 	 	 	 	 	 d3d#Zed+d$       Z ed+d%       Z!	 d-	 	 	 	 	 d4 fd&Z"ed5d'       Z# xZ$S )6	ChatKonkoa  `ChatKonko` Chat large language models API.

    To use, you should have the ``konko`` python package installed, and the
    environment variable ``KONKO_API_KEY`` and ``OPENAI_API_KEY`` set with your API key.

    Any parameters that are valid to be passed to the konko.create call can be passed
    in, even if not explicitly saved on this class.

    Example:
        .. code-block:: python

            from langchain_community.chat_models import ChatKonko
            llm = ChatKonko(model="meta-llama/Llama-2-13b-chat-hf")
    c                    dddS )NKONKO_API_KEYOPENAI_API_KEY)konko_api_keyopenai_api_key selfs    d/var/www/html/test/engine/venv/lib/python3.12/site-packages/langchain_community/chat_models/konko.py
lc_secretszChatKonko.lc_secrets=   s    !0DTUU    boolc                     y)z9Return whether this model can be serialized by Langchain.Fr#   )clss    r&   is_lc_serializablezChatKonko.is_lc_serializableA   s     r(   Nr   clientmodel)defaultaliasstrgffffff?floattemperature)default_factoryDict[str, Any]model_kwargszOptional[str]r"   r!      intmax_retriesF	streaming   n   
max_tokensc                   t        t        |dd            |d<   	 ddl}	 t	               r|j
                  j                  |d<   n|j                  |d<   t        |d      st        j                  d	       |d
   dk  rt        d      |d
   dkD  r|d   rt        d      |S # t        $ r t        d      w xY w# t        $ r t        d      w xY w)z?Validate that api key and python package exists in environment.r!   r   r   NzRCould not import konko python package. Please install it with `pip install konko`.r-   z`konko` has no `ChatCompletion` attribute, this is likely due to an old version of the konko package. Try upgrading it with `pip install --upgrade konko`._is_legacy_openaizhYou are using an older version of the 'konko' package. Please consider upgrading to access new features.r<   r;   zn must be at least 1.r:   zn must be 1 when streaming.)r   r   konkoImportErrorr   chatcompletionsChatCompletionAttributeError
ValueErrorhasattrwarningswarn)r+   valuesrA   s      r&   validate_environmentzChatKonko.validate_environmentX   s     #8 /J#
	
	~#(::#9#9x #(#7#7x  u12MMD
 #;?455#;?vk2:;;7  	> 	  	6 	s   B! 3B9 !B69Cc                    | j                   | j                  | j                  | j                  | j                  d| j
                  S )z1Get the default parameters for calling Konko API.)r.   r>   streamr<   r3   )r.   r>   r:   r<   r3   r6   r$   s    r&   _default_paramszChatKonko._default_params~   sC     ZZ//nn++
 
 	
r(   c                   |s	 t        t        j                  d         }nt	        |t
              rt        |      }| s	 t        t        j                  d         } nt	        | t
              rt        |       } | d}dd| j                          i}|r!t        t        |      j                         |d<   t        j                  ||      }|j                  d	k7  rt        d
| d|j                         |j                         d   D ch c]  }|d   	 c}S # t        $ r Y w xY w# t        $ r t        d      w xY wc c}w )z$Get available models from Konko API.r    r   z^Konko API key must be passed as keyword argument or set in environment variable KONKO_API_KEY.z/modelsAuthorizationzBearer zX-OpenAI-Api-Key)headers   zError getting models from z: dataid)r   osenvironKeyError
isinstancer1   rG   get_secret_valuer   r   requestsgetstatus_codejson)r!   r"   konko_api_base
models_urlrR   models_responser.   s          r&   get_available_modelszChatKonko.get_available_models   sa    !6rzzBR7S!T ,2>BN  5bjj6Q R s+1-@M&'w/
 w}'E'E'G&HI
 *.>+  &' #,,z7C&&#-,ZL?;V;V:WX  *9)=)=)?)GHdHHG     A 4 Is"   D D( 
E 	D%$D%(D=c                      d fd} |di |S )Nc                 <     j                   j                  di | S )Nr#   )r-   create)kwargsr%   s    r&   _completion_with_retryz?ChatKonko.completion_with_retry.<locals>._completion_with_retry   s    %4;;%%///r(   )rf   r   returnr   r#   r#   )r%   run_managerrf   rg   s   `   r&   completion_with_retryzChatKonko.completion_with_retry   s    	0 &///r(   c              +    K   | j                  ||      \  }}i ||ddi}t        } | j                  d||d|D ]  }t        |d         dk(  r|d   d   }	t	        |	d   |      }|	j                  d      }
|
t        |
      nd }|j                  }t        ||	      }|r|j                  |j                  |
       |  y w)NrN   Tmessagesri   choicesr   deltafinish_reason)rp   )messagegeneration_info)chunkr#   )_create_message_dictsr   rj   lenr   r\   dict	__class__r   on_llm_new_tokentext)r%   rm   stopri   rf   message_dictsparamsdefault_chunk_classrs   choicerp   rr   cg_chunks                r&   _streamzChatKonko._stream   s     !% : :8T Jv5F5f5h5,/T// 
"
?E
 	E 5#$)9%a(F3w!4E #JJ7M5B5N=1TX  #(//*H ,,X]](,KN'	s   CC	c                    ||n| j                   }|r! | j                  |f||d|}t        |      S | j                  ||      \  }}	i |	|}	 | j                  d||d|	}
| j                  |
      S )N)rz   ri   rl   r#   )r:   r   r   rt   rj   _create_chat_result)r%   rm   rz   ri   rN   rf   should_streamstream_iterr{   r|   responses              r&   	_generatezChatKonko._generate   s     #)"4$..&$,,#@FK (44 $ : :8T Jv%F%f%-4-- 
"
?E
 ''11r(   c                    | j                   }|d|v rt        d      ||d<   |D cg c]  }t        |       }}||fS c c}w )Nrz   z2`stop` found in both the input and default params.)_client_paramsrG   r   )r%   rm   rz   r|   mr{   s         r&   rt   zChatKonko._create_message_dicts   s\     $$ !UVV!F6N=EF03FFf$$ Gs   ?c                :    i d| j                   i| j                  S )zGet the identifying parameters.
model_name)r.   rO   r$   s    r&   _identifying_paramszChatKonko._identifying_params  s$     F<,E0D0DEEr(   c                    i | j                   S )z-Get the parameters used for the konko client.)rO   r$   s    r&   r   zChatKonko._client_params	  s     ($&&''r(   c                \    d| j                   it        | 	  |      | j                  |S )z,Get the parameters used to invoke the model.r.   )rz   )r.   super_get_invocation_paramsrO   )r%   rz   rf   rw   s      r&   r   z ChatKonko._get_invocation_params  sG    
 TZZ
g,$,7
 ""
 	
 	
r(   c                     y)zReturn type of chat model.z
konko-chatr#   r$   s    r&   	_llm_typezChatKonko._llm_type  s     r(   )rh   zDict[str, str])rh   r)   )rK   r   rh   r   )rh   r5   )r!   Union[str, SecretStr, None]r"   r   r_   r1   rh   zSet[str])N)ri   "Optional[CallbackManagerForLLMRun]rf   r   rh   r   )NN)
rm   List[BaseMessage]rz   Optional[List[str]]ri   r   rf   r   rh   zIterator[ChatGenerationChunk])NNN)rm   r   rz   r   ri   r   rN   zOptional[bool]rf   r   rh   r   )rm   r   rz   r   rh   z+Tuple[List[Dict[str, Any]], Dict[str, Any]])rz   r   rf   r   rh   r5   )rh   r1   )%__name__
__module____qualname____doc__propertyr'   classmethodr,   r-   __annotations__r   DEFAULT_MODELr.   r3   rv   r6   r"   r!   r9   r:   r<   r>   r   rL   rO   staticmethodDEFAULT_API_BASErb   rj   r   r   rt   r   r   r   r   __classcell__)rw   s   @r&   r   r   -   sP    V V   FC}G<E3<K+#(#>L.>V$(NM(#'M='K<It/AsJAJ/# #J 	
 	
 596:..I2.I3.I .I 
	.I .Ib AE0=0PS0	0 %):>	# " 8	
  
'F %):>!%2#2 "2 8	2
 2 2 
2,	%)	%1D	%	4	% F F ( (
 +/	
'	
:=	
		
  r(   r   ).r   
__future__r   loggingrV   rI   typingr   r   r   r   r   r	   r
   r   r   r[   langchain_core.callbacksr   langchain_core.messagesr   r   langchain_core.outputsr   r   langchain_core.utilsr   r   r   pydanticr   r   #langchain_community.adapters.openair   &langchain_community.chat_models.openair   r   r    langchain_community.utils.openair   r   r   	getLoggerr   loggerr   r#   r(   r&   <module>r      s~     "  	 
 
 
  @ B V V % 
 :, 0			8	$o
 or(   