o
    -i%                  	   @   s  U 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
 d dlZd dlmZ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 d d	lmZ d d
lmZmZ d dlmZ d dl m!Z! d dl"m#Z# d dl$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l0m1Z1m2Z2 d dl3m4Z4 d dl5m6Z6m7Z7m8Z8m9Z9 d dl:m;Z;m<Z< d dl=m>Z> eZ?eege#dB f Z@ee?egee f ZAeeeffeeeffe/e-e,ffe*e'e(ffe<e9e6ffe;e8e7ffe4e2e1ffgZBeCeDe?eDe@eAf f  eEd< dd eBD ZFdefddZGdS )    N)	AwaitableCallable)
HTTPStatus)Any)	APIRouterDependsHTTPExceptionRequest)JSONResponseResponse)base)chatcreate_chat_completion)ChatCompletionRequest)
completioncreate_completion)CompletionRequest)ErrorResponse)OpenAIServing)validate_json_request)classifycreate_classify)ClassificationRequest)create_embedding	embedding)EmbeddingRequest)create_poolingpooling)PoolingRequest)create_score	do_rerankrerankscore)RerankRequestScoreRequesthealthINVOCATION_TYPESc                 C   s&   g | ]\}\}}t |||ffqS  )pydanticTypeAdapter).0request_typeget_handlerendpointr(   r(   ^/home/ubuntu/veenaModal/venv/lib/python3.10/site-packages/vllm/entrypoints/sagemaker/routes.py
<listcomp>D   s    
r0   routerc              
   C   s   | j dtd| jdtdtjdtdtfdd}| j dttgtj	j
dtitjj
dtitjj
dtiid	tjt tjdd
dtfdd}| S )Nz/ping)response_classraw_requestreturnc                    s   t | I dH S )z+Ping check. Endpoint required for SageMakerNr%   r3   r(   r(   r/   pingK   s   z'register_sagemaker_routes.<locals>.pingz/invocationsmodel)dependencies	responses)adapter_pathc           
   
      s   z	   I dH }W n t jy# } zttjjd| d|d}~ww  fddtD }|D ]\}}z||}W n
 tj	yD   Y q/w || I dH   S fdd|D }d| }t
 j|d}	t|	 |	jjd	S )
z9For SageMaker, routes requests based on the request type.NzJSON decode error: )status_codedetailc                    s(   g | ]\}\}}| d ur||fqS )Nr(   )r+   	validatorr-   r.   r5   r(   r/   r0   h   s
    
zBregister_sagemaker_routes.<locals>.invocations.<locals>.<listcomp>c                    s.   g | ]\}}t |j  tr jnt qS r(   )
isinstance_typetype__name__str)r+   r=   _)tr(   r/   r0   v   s    z;Cannot find suitable handler for request. Expected one of: )message)contentr;   )jsonJSONDecodeErrorr   r   BAD_REQUESTvalueINVOCATION_VALIDATORSvalidate_pythonr)   ValidationErrorr   create_error_responser
   
model_dumperrorcode)
r3   bodyevalid_endpointsrequest_validatorr.   request
type_namesmsgresr(   )r3   rD   r/   invocationsR   s8   


z.register_sagemaker_routes.<locals>.invocations)postr   getsagemaker_standardsregister_ping_handlerr	   r   r   r   rI   rJ   r   UNSUPPORTED_MEDIA_TYPEINTERNAL_SERVER_ERRORregister_invocation_handlerstateful_session_managerinject_adapter_id)r1   r6   rZ   r(   r(   r/   register_sagemaker_routesJ   s"   	
 rd   )HrG   collections.abcr   r   httpr   typingr   +model_hosting_container_standards.sagemaker	sagemakerr]   r)   fastapir   r   r   r	   fastapi.responsesr
   r   "vllm.entrypoints.openai.api_serverr   2vllm.entrypoints.openai.chat_completion.api_routerr   r   0vllm.entrypoints.openai.chat_completion.protocolr   -vllm.entrypoints.openai.completion.api_routerr   r   +vllm.entrypoints.openai.completion.protocolr   'vllm.entrypoints.openai.engine.protocolr   &vllm.entrypoints.openai.engine.servingr   vllm.entrypoints.openai.utilsr   ,vllm.entrypoints.pooling.classify.api_routerr   r   *vllm.entrypoints.pooling.classify.protocolr   )vllm.entrypoints.pooling.embed.api_routerr   r   'vllm.entrypoints.pooling.embed.protocolr   +vllm.entrypoints.pooling.pooling.api_routerr   r   )vllm.entrypoints.pooling.pooling.protocolr   )vllm.entrypoints.pooling.score.api_routerr   r    r!   r"   'vllm.entrypoints.pooling.score.protocolr#   r$   ,vllm.entrypoints.serve.instrumentator.healthr&   RequestTypeGetHandlerFn
EndpointFnr'   listtuple__annotations__rK   rd   r(   r(   r(   r/   <module>   sP   







 