
    Zi

                     t    d dl mZmZ d dlZd dlmZmZmZmZ d dl	m
Z
  G d dej        ej                  ZdS )    )OptionalSetN)retrystop_after_attemptwait_exponential_jitterretry_if_result)Spanc                       e Zd ZU dZeed<   eej                 ed<   d e	ej        j
        ej        j        g          fdedeej                 ddfdZd Zd	 Zd
 Zd Zd ZdS ) RetryOnRpcErrorClientInterceptorz
    A gRPC client interceptor that retries RPCs on specific status codes. By default, it retries on UNAVAILABLE and UNKNOWN status codes.

    This interceptor should be placed after the OpenTelemetry interceptor in the interceptor list.
    max_attemptsretryable_status_codes   returnNc                 "    || _         || _        d S N)r   r   )selfr   r   s      Q/root/projects/butler/venv/lib64/python3.11/site-packages/chromadb/proto/utils.py__init__z)RetryOnRpcErrorClientInterceptor.__init__   s     )&<###    c                      d fd}t          t          dd          t           j                  t	           fd          |          fd            } |||          S )Nc                 F    ddl m} ||                    d          d S d S )Nr   )tracerzWaiting to retry RPC) chromadb.telemetry.opentelemetryr   
start_span)_r   
sleep_spans     r   before_sleepzFRetryOnRpcErrorClientInterceptor._intercept_call.<locals>.before_sleep    s=    ?????? !#../EFF


 "!r   g?)jitterc                 :    |                                  j        v S r   )coder   )xr   s    r   <lambda>zBRetryOnRpcErrorClientInterceptor._intercept_call.<locals>.<lambda>*   s    AFFHH8S,S r   )waitstopr   r   c                  D                                      d  | i |S r   )end)argskwargscontinuationr   s     r   wrappedzARetryOnRpcErrorClientInterceptor._intercept_call.<locals>.wrapped'   s5     %   !
<0000r   )r   r   r   r   r   )r   r)   client_call_detailsrequest_or_iteratorr   r*   r   s   ``    @r   _intercept_callz0RetryOnRpcErrorClientInterceptor._intercept_call   s    %)
	G 	G 	G 	G 	G 
(S999#D$566!"S"S"S"STT%	

 

 

	1 	1 	1 	1 	1

 

	1 w*,?@@@r   c                 0    |                      |||          S r   r-   r   r)   r+   requests       r   intercept_unary_unaryz6RetryOnRpcErrorClientInterceptor.intercept_unary_unary6       ##L2EwOOOr   c                 0    |                      |||          S r   r/   r0   s       r   intercept_unary_streamz7RetryOnRpcErrorClientInterceptor.intercept_unary_stream9   r3   r   c                 0    |                      |||          S r   r/   r   r)   r+   request_iterators       r   intercept_stream_unaryz7RetryOnRpcErrorClientInterceptor.intercept_stream_unary<        ##L2EGWXXXr   c                 0    |                      |||          S r   r/   r7   s       r   intercept_stream_streamz8RetryOnRpcErrorClientInterceptor.intercept_stream_streamA   r:   r   )__name__
__module____qualname____doc__int__annotations__r   grpc
StatusCodesetUNAVAILABLEUNKNOWNr   r-   r2   r5   r9   r<    r   r   r   r      s           0000 7:s_($/*AB8
 8
= == !$DO 4= 
= = = =A A A2P P PP P PY Y Y
Y Y Y Y Yr   r   )typingr   r   rC   tenacityr   r   r   r   opentelemetry.tracer	   UnaryUnaryClientInterceptorUnaryStreamClientInterceptorr   rH   r   r   <module>rN      s                     X X X X X X X X X X X X $ $ $ $ $ $=Y =Y =Y =Y =Y$d&G=Y =Y =Y =Y =Yr   