o
    ưi+                     @   s   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
 d dlZd dl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 erQd dlmZ eZneZG d	d
 d
eZdS )    N)TYPE_CHECKINGAnyDictListLiteralOptionalUnioncast)HeadersResponse)BaseBatchesConfig)BaseLLMException)AllMessageValuesCreateBatchRequest)LiteLLMBatchLlmProvidersModelResponse)Loggingc                   @   s  e Zd Zdd ZedefddZ		d,deded	e	e
 d
ededee dee defddZdee dee ded
edededefddZdeded
ededeeeeeef f f
ddZdee dejdededef
ddZdee ded
ededef
ddZded
ededeeeeeef f fddZdee dejdededef
dd Zd!ed"edeeef dd#fd$d%Z		d,deded&eded'ed	e	e
 d
eded(edee d)ee  defd*d+Z!dS )-AnthropicBatchesConfigc                 C   s,   ddl m} ddlm} | | _| | _d S )N   )AnthropicConfig)AnthropicModelInfo)Zchat.transformationr   common_utilsr   anthropic_chat_configanthropic_model_info)selfr   r    r   d/home/app/Keep/.python/lib/python3.10/site-packages/litellm/llms/anthropic/batches/transformation.py__init__   s   zAnthropicBatchesConfig.__init__returnc                 C   s   t jS )z4Return the LLM provider type for this configuration.)r   Z	ANTHROPIC)r   r   r   r   custom_llm_provider   s   z*AnthropicBatchesConfig.custom_llm_providerNheadersmodelmessagesoptional_paramslitellm_paramsapi_keyapi_basec           	      C   sJ   |p| j  }|du rtdddd|d}d|vrd|d< || |S )zAValidate and prepare environment-specific headers and parameters.NzMissing Anthropic API Key - A call is being made to anthropic but no key is set either in the environment variables or via paramszapplication/jsonz
2023-06-01)acceptzanthropic-versionzcontent-typez	x-api-keyzanthropic-betazmessage-batches-2024-09-24)r   Zget_api_key
ValueErrorupdate)	r   r!   r"   r#   r$   r%   r&   r'   _headersr   r   r   validate_environment"   s   
z+AnthropicBatchesConfig.validate_environmentdatac                 C   s.   |p| j |}|ds|d d}|S )z0Get the complete URL for batch creation request.z/v1/messages/batches/)r   get_api_baseendswithrstrip)r   r'   r&   r"   r$   r%   r-   r   r   r   get_complete_batch_url?   s   

z-AnthropicBatchesConfig.get_complete_batch_urlcreate_batch_datac                 C      t d)z
        Transform the batch creation request to Anthropic format.
        
        Not currently implemented - placeholder to satisfy abstract base class.
        0Batch creation not yet implemented for AnthropicNotImplementedError)r   r"   r3   r$   r%   r   r   r   transform_create_batch_requestN      z5AnthropicBatchesConfig.transform_create_batch_requestraw_responselogging_objc                 C   r4   )z
        Transform Anthropic MessageBatch creation response to LiteLLM format.
        
        Not currently implemented - placeholder to satisfy abstract base class.
        r5   r6   )r   r"   r:   r;   r%   r   r   r   transform_create_batch_response\   r9   z6AnthropicBatchesConfig.transform_create_batch_responsebatch_idc                 C   s$   |p| j |}|d d| S )a  
        Get the complete URL for batch retrieval request.
        
        Args:
            api_base: Base API URL (optional, will use default if not provided)
            batch_id: Batch ID to retrieve
            optional_params: Optional parameters
            litellm_params: LiteLLM parameters
            
        Returns:
            Complete URL for Anthropic batch retrieval: {api_base}/v1/messages/batches/{batch_id}
        r.   z/v1/messages/batches/)r   r/   r1   )r   r'   r=   r$   r%   r   r   r   get_retrieve_batch_urlj   s   z-AnthropicBatchesConfig.get_retrieve_batch_urlc                 C   s   i S )z
        Transform batch retrieval request for Anthropic.
        
        For Anthropic, the URL is constructed by get_retrieve_batch_url(),
        so this method returns an empty dict (no additional request params needed).
        r   )r   r=   r$   r%   r   r   r    transform_retrieve_batch_request   s   z7AnthropicBatchesConfig.transform_retrieve_batch_requestc              
   C   s  z|  }W n ty } ztd| d}~ww |dd}|dd}dddd	}	|	|d}
d
tt dtt fdd}||d}||d}||d}||d}||d}|di }ddlm} |t	|dd|dd|dd|dd|ddg|dd|ddd}t
d2i d|ddddd dd!d"d#d$d%|
d&|d'dd|ptt d(|dkr|ndd|d)dd*|d+kr|ndd,dd-|r|ndd.|d/kr|ndd0|d/kr|r|n
dd|d1i S d|d1i S )3zFTransform Anthropic MessageBatch retrieval response to LiteLLM format.z*Failed to parse Anthropic batch response: Nid processing_statusin_progress
cancelling	completed)rC   	cancelingendedts_strr   c                 S   sJ   | sd S zddl m } || dd}t| W S  ty$   Y d S w )Nr   )datetimeZz+00:00)rI   fromisoformatreplaceint	timestamp	Exception)rH   rI   dtr   r   r   parse_timestamp   s   zQAnthropicBatchesConfig.transform_retrieve_batch_response.<locals>.parse_timestamp
created_atended_at
expires_atcancel_initiated_atarchived_atrequest_countsr   )BatchRequestCounts
processingZ	succeededZerroredcanceledZexpired)totalrE   failedobjectbatchZendpointz/v1/messageserrorsZinput_file_idNoneZcompletion_windowZ24hstatusZoutput_file_idZerror_file_idZin_progress_atZfinalizing_atZcompleted_atrG   Z	failed_atZ
expired_atZcancelling_atrF   Zcancelled_atmetadatar   )jsonrO   r)   getr   strrM   Zopenai.types.batchrX   sumr   time)r   r"   r:   r;   r%   Zresponse_dataer=   rB   Zstatus_mappingZopenai_statusrQ   rR   rS   rT   rU   rV   Zrequest_counts_datarX   rW   r   r   r   !transform_retrieve_batch_response   s   







	

z8AnthropicBatchesConfig.transform_retrieve_batch_responseerror_messagestatus_coder   c                 C   s@   ddl m} t|trt|}n	t|tr|nd}||||dS )z.Get the appropriate error class for Anthropic.r   )AnthropicErrorN)rk   messager!   )r   rl   
isinstancedictr
   )r   rj   rk   r!   rl   Zheaders_objr   r   r   get_error_class   s
   

z&AnthropicBatchesConfig.get_error_classmodel_responserequest_dataencoding	json_modec              
   C   s   ddl m} ddlm} |j }g }zO|d}|D ]9}| }|s$qz&t|}|d d }| j	j
|||d}t|dd }|rI|t|| W q tjyT   Y qw ||}t|d| |W S  typ } z|d }~ww )	Nr   )BaseTokenUsageProcessor)Usage
resultrm   )completion_responser:   rq   usage)Zlitellm.cost_calculatorru   litellm.types.utilsrv   textstripsplitrc   loadsr   Ztransform_parsed_responsegetattrappendr	   JSONDecodeErrorZcombine_usage_objectssetattrrO   )r   r"   r:   rq   r;   rr   r#   r$   r%   rs   r&   rt   ru   rv   Zresponse_textZ	all_usagelineslineZresponse_jsonry   Ztransformed_responseZtransformed_response_usageZcombined_usagerh   r   r   r   transform_response   sF   



z)AnthropicBatchesConfig.transform_response)NN)"__name__
__module____qualname__r   propertyr   r    ro   re   r   r   r   r,   r   r   r2   r   bytesr   r8   httpxr   LoggingClassr   r<   r>   r?   ri   rM   r
   rp   r   boolr   r   r   r   r   r      s    	






P

	
r   )rc   rg   typingr   r   r   r   r   r   r   r	   r   r
   r   Z,litellm.llms.base_llm.batches.transformationr   Z)litellm.llms.base_llm.chat.transformationr   Zlitellm.types.llms.openair   r   r{   r   r   r   Z*litellm.litellm_core_utils.litellm_loggingr   ZLiteLLMLoggingObjr   r   r   r   r   r   <module>   s    (