Ë
    uPêjé  ã                   ó‚   — d dl mZ d dlZd dlmZmZ d dlmZmZ d dl	m
Z
 d dlmZmZ d dlmZ dd	lmZ  G d
„ de«      Zy)é    )ÚFinalN)ÚAllMessageValuesÚBaseLLMException)ÚBaseEmbeddingConfigÚLiteLLMLoggingObj)ÚAllEmbeddingInputValues)ÚEmbeddingResponseÚUsage)Útoken_counteré   ©ÚTritonErrorc                   ól  — e Zd ZdZdd„Zdedefd„Zdededed	e	def
d
„Z
	 	 d dededee   dedededz  dedz  defd„Zdededededef
d„Zdi i i fdedej                   dedededz  dedededefd„Zdededefd„Zdededeej.                  z  defd„Zedee   dee   deee      fd„«       Zy)!ÚTritonEmbeddingConfigzR
    Transformations for triton /embeddings endpoint (This is a trtllm model)
    ÚreturnNc                  ó   — y ©N© )Úselfs    úq/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/llms/triton/embedding/transformation.pyÚ__init__zTritonEmbeddingConfig.__init__   s   € Øó    Úmodelc                 ó   — g S r   r   )r   r   s     r   Úget_supported_openai_paramsz1TritonEmbeddingConfig.get_supported_openai_params   s   € Øˆ	r   Únon_default_paramsÚoptional_paramsÚdrop_paramsc                 ó   — |S )z>
        Map OpenAI params to Triton Embedding params
        r   )r   r   r   r   r   s        r   Úmap_openai_paramsz'TritonEmbeddingConfig.map_openai_params   s
   € ð Ðr   ÚheadersÚmessagesÚlitellm_paramsÚapi_keyÚapi_basec                 ó   — i S r   r   )r   r!   r   r"   r   r#   r$   r%   s           r   Úvalidate_environmentz*TritonEmbeddingConfig.validate_environment(   s	   € ð ˆ	r   Úinputc                 ó*   — ddt        |«      gd|dœgiS )NÚinputsÚ
input_textÚBYTES)ÚnameÚshapeÚdatatypeÚdata)Úlen)r   r   r(   r   r!   s        r   Útransform_embedding_requestz1TritonEmbeddingConfig.transform_embedding_request4   s-   € ð à(Ü! %›j˜\Ø 'Ø!ñ	ðð	
ð 		
r   Úraw_responseÚmodel_responseÚlogging_objÚrequest_datac	                 ó”  — 	 |j                  «       }	g }
|	d   }|D ]F  }|d   }|d   }| j                  ||«      }t        |«      D ]  \  }}|
j                  d||dœ«       Œ ŒH |	j                  dd«      |_	        |
|_
        | j                  ||¬	«      |_        |S # t        $ r" t        |j                  |j                  ¬«      ‚w xY w)
N)ÚmessageÚstatus_codeÚoutputsr.   r0   Ú	embedding)ÚobjectÚindexr;   Ú
model_nameÚNone)r   r6   )ÚjsonÚ	Exceptionr   Útextr9   Úsplit_embedding_by_shapeÚ	enumerateÚappendÚgetr   r0   Ú_build_embedding_usageÚusage)r   r   r3   r4   r5   r$   r6   r   r#   Úraw_response_jsonÚ_embedding_outputÚ_outputsÚoutputÚ_shapeÚ_dataÚ_split_output_dataÚidxr;   s                     r   Útransform_embedding_responsez2TritonEmbeddingConfig.transform_embedding_responseF   sü   € ð	_Ø'3×'8Ñ'8Ó':Ðð $&Ðà+¨IÑ6ˆØò 	ˆFØ˜G‘_ˆFØ˜6‘NˆEØ!%×!>Ñ!>¸uÀfÓ!MÐä"+Ð,>Ó"?ò ‘��YØ!×(Ñ(à"-Ø!$Ø%.ñõñð	ð  1×4Ñ4°\À6ÓJˆÔØ/ˆÔØ#×:Ñ:ÀÐUaÐ:ÓbˆÔØÐøô/ ò 	_Ü l×&7Ñ&7À\×E]ÑE]Ô^Ð^ð	_ús   ‚B Â+Cc                 ó°  — |j                  dg «      }g }|D ]c  }t        |t        «      sŒ|j                  d«      dk(  sŒ)|j                  dg «      }t        |t        «      r|D �cg c]  }t	        |«      ‘Œ }} n d}|D ]  }	|	sŒ	 |t        ||	¬«      z  }Œ t        |d|¬«      S c c}w # t        $ r |t        |	j                  «       «      z  }Y ŒSw xY w)Nr*   r-   r+   r0   r   )r   rB   )Úprompt_tokensÚcompletion_tokensÚtotal_tokens)
rF   Ú
isinstanceÚdictÚlistÚstrr   rA   r1   Úsplitr
   )
r   r   r6   Ú
input_dataÚinput_text_valuesÚitemÚdata_valuesÚvaluerS   rB   s
             r   rG   z,TritonEmbeddingConfig._build_embedding_usagel   sî   € Ø(×,Ñ,¨X°rÓ:ˆ
Ø')ÐØò 	ˆDÜ˜$¤Õ%¨$¯(©(°6Ó*:¸lÓ*JØ"Ÿh™h v¨rÓ2�Ü˜k¬4Ô0ØALÖ(M¸¬¨U­Ð(MÐ%Ð(MÙð	ð ˆØ%ò 	3ˆDÙØð3Ø¤°UÀÔ!FÑF‘ð		3ô Ø'ØØ&ô
ð 	
ùò )Nøô ò 3Ø¤ T§Z¡Z£\Ó!2Ñ2’ð3ús   Á&B(ÂB-Â-%CÃCÚerror_messager9   c                 ó   — t        |||¬«      S )N)r8   r9   r!   r   )r   r`   r9   r!   s       r   Úget_error_classz%TritonEmbeddingConfig.get_error_class…   s   € Ü =¸kÐSZÔ[Ð[r   r0   r.   c                 ó˜   — t        |«      dk7  rt        d«      ‚|d   }t        |d   «      D �cg c]  }| ||z  |dz   |z   ‘Œ c}S c c}w )Nr   zShape must be of length 2.é   r   )r1   Ú
ValueErrorÚrange)r0   r.   Úembedding_sizeÚis       r   rC   z.TritonEmbeddingConfig.split_embedding_by_shapeˆ   sV   € äˆu‹:˜Š?ÜÐ9Ó:Ð:Ø % a¡ˆÜMRÐSXÐYZÑS[Ë_Ö]È��Q˜Ñ'¨1¨q©5°NÑ*BÒCÒ]Ð]ùÒ]s   ¯A)r   N)NN)Ú__name__Ú
__module__Ú__qualname__Ú__doc__r   rY   rX   r   rW   Úboolr    r   r'   r   r2   ÚhttpxÚResponser	   r   rQ   r
   rG   ÚintÚHeadersr   rb   ÚstaticmethodÚfloatrC   r   r   r   r   r      só  „ ñóð°ð ¸ó ð
à ð
ð ð
ð ð	
ð
 ð
ð 
ó
ð& #Ø#ñ
àð
ð ð
ð Ð'Ñ(ð	
ð
 ð
ð ð
ð �t‘ð
ð ˜‘*ð
ð 
ó
ð
àð
ð 'ð
ð ð	
ð
 ð
ð 
ó
ð0 #ØØ "Ø!ñ$àð$ð —n‘nð$ð *ð	$ð
 'ð$ð �t‘ð$ð ð$ð ð$ð ð$ð 
ó$ðL
¨Cð 
¸tð 
Èó 
ð2\¨Sð \¸sð \ÈTÐTY×TaÑTaÑMað \Ðfvó \ð ð^ t¨E¡{ð ^¸4À¹9ð ^ÈÈdÐSXÉkÑIZò ^ó ñ^r   r   )Útypingr   rn   Ú)litellm.llms.base_llm.chat.transformationr   r   Ú.litellm.llms.base_llm.embedding.transformationr   r   Úlitellm.types.llms.openair   Úlitellm.types.utilsr	   r
   Úlitellm.utilsr   Úcommon_utilsr   r   r   r   r   ú<module>r{      s2   ðÝ ã ç X÷õ >ß 8Ý 'å &ô|^Ð/õ |^r   