
    ug                        d dl Z d dlmZ d dlZd dlZd dlmZmZmZ d dl	m
Z
 d dlmZ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 d d
lmZmZmZmZ d dlmZ d dlm Z m!Z!m"Z" d dl#m$Z$ d dl%Z% e%jL                  e'      Z( G d dee      Z)y)    N) RetryOnRpcErrorClientInterceptor)ProducerConsumerConsumerCallbackFn)to_proto_submit)PushLogsRequestPullLogsRequest	LogRecord)LogServiceStub)OtelInterceptor)OperationRecordSeqId)System)OpenTelemetryClientOpenTelemetryGranularityadd_attributes_to_current_spantrace_method)override)SequenceOptionalcast)UUIDc                       e Zd ZU dZeed<   eed<   ej                  ed<   e	ed<   eed<   de
f fdZ ed	ej                        ed- fd              Z edej                        ed- fd              Z edej                        ed- fd              Z edej                        eded
dfd              Z edej                        eded
dfd              Z edej                        ededed
efd              Z edej                        ededee   d
ee   fd              Z edej                        e	 	 	 d.dededee   dee   d ee   d
efd!              Z ed"ej                        ed#ed
dfd$              Zed
efd%       Zed
efd&       Z e!ed
efd'              Z"ded(ee   d
efd)Z#ded*ed+ed
ee$   fd,Z% xZ&S )/
LogServicez(
    Distributed Chroma Log Service
    _log_service_stub_request_timeout_seconds_channel_log_service_url_log_service_portsystemc                    |j                   j                  d      | _        |j                   j                  d      | _        |j                   j                  d      | _        |j                  t
              | _        t        | !  |       y )Nchroma_logservice_hostchroma_logservice_port)chroma_logservice_request_timeout_seconds)	settingsrequirer   r   r   r   _opentelemetry_clientsuper__init__)selfr    	__class__s     S/var/www/openai/venv/lib/python3.12/site-packages/chromadb/logservice/logservice.pyr)   zLogService.__init__-   sn     & 7 78P Q!'!8!89Q!R(.(?(?7)
% &,^^4G%H"     zLogService.startreturnNc                 .   t        j                  | j                   d| j                         | _        t               t               g}t        j                  | j                  g| | _        t        | j                        | _	        t        | -          y )N:)grpcinsecure_channelr   r   r   r   r   intercept_channelr   r   r(   start)r*   interceptorsr+   s     r,   r4   zLogService.start6   s{     --$$%Qt'='=&>?
 ()+K+MN..t}}L|L!/!>r-   zLogService.stopc                 V    | j                   j                          t        |           y N)r   closer(   stopr*   r+   s    r,   r9   zLogService.stopA   s     	r-   zLogService.reset_statec                 "    t         |           y r7   )r(   reset_stater:   s    r,   r<   zLogService.reset_stateG   s     	r-   zLogService.delete_logcollection_idc                     t        d      NzNot implementedNotImplementedErrorr*   r=   s     r,   
delete_logzLogService.delete_logL        ""344r-   zLogService.purge_logc                     t        d      r?   r@   rB   s     r,   	purge_logzLogService.purge_logQ   rD   r-   zLogService.submit_embedding	embeddingc                 \    | j                   st        d      | j                  ||g      d   S )NComponent not runningr   )_runningRuntimeErrorsubmit_embeddings)r*   r=   rG   s      r,   submit_embeddingzLogService.submit_embeddingV   s1    
 }}677%%mi[A!DDr-   zLogService.submit_embeddings
embeddingsc           	      r   t         j                  dt        |       d|        t        dt        |      i       | j                  st        d      t        |      dk(  rg S g }|D cg c]  }t        |       }}|j                  | j                  |t        t        t           |                   |S c c}w )NzSubmitting z" embeddings to log for collection records_countrI   r   )loggerinfolenr   rJ   rK   r   append	push_logsr   r   r   )r*   r=   rN   countsrecordprotos_to_submits         r,   rL   zLogService.submit_embeddings`   s    
 	#j/**L]O\	
 	'Z	
 }}677z?aI BLM*OF3*MNNXo.0@A	
  Ns   (B4zLogService.subscribe
consume_fnr4   endidc                 L    t         j                  d| d       t        d      S )NzSubscribing to log for , noop for logservicer   )int)rQ   rR   r   )r*   r=   rY   r4   rZ   r[   s         r,   	subscribezLogService.subscribe   s%     	-m_<QRS{r-   zLogService.unsubscribesubscription_idc                 6    t         j                  d| d       y )NzUnsubscribing from r]   )rQ   rR   )r*   r`   s     r,   unsubscribezLogService.unsubscribe   s     	)/)::OPQr-   c                      y)Nr    r*   s    r,   	min_seqidzLogService.min_seqid   s    r-   c                 "    t         j                  S r7   )sysmaxsizere   s    r,   	max_seqidzLogService.max_seqid   s    {{r-   c                      y)Nd   rd   re   s    r,   max_batch_sizezLogService.max_batch_size   s     r-   recordsc                     t        t        |      |      }| j                  j                  || j                        }|j
                  S )N)r=   rn   timeout)r   strr   PushLogsr   record_count)r*   r=   rn   requestresponses        r,   rU   zLogService.push_logs   sG    !M0BGT))22T:: 3 
 $$$r-   start_offset
batch_sizec                     t        t        |      ||t        j                               }| j                  j                  || j                        }|j                  S )N)r=   start_from_offsetrx   end_timestamprp   )r	   rr   timetime_nsr   PullLogsr   rn   )r*   r=   rw   rx   ru   rv   s         r,   	pull_logszLogService.pull_logs   s[     "m,*!,,.	
 ))22T:: 3 
 r-   )r.   N)NNN)'__name__
__module____qualname____doc__r   __annotations__r^   r1   Channelrr   r   r)   r   r   ALLr   r4   r9   r<   r   rC   rF   r   r   rM   r   rL   r   r   r_   rb   rf   rj   propertyrm   rU   r
   r   __classcell__)r+   s   @r,   r   r   "   s%    &%!!ll!v ! $&>&B&BC  D #%=%A%AB  C *,D,H,HI  J )+C+G+GH5 5 5  I5 (*B*F*FG5t 5 5  H5 /1I1M1MNE!E.=E	E  OE 02J2N2NO!/7/H	%  PB (*B*F*FG
 "&#!		 '	 		
 e_	 TN	 
	  H	 *,D,H,HIR4 RD R  JR 5   5      %t %h6O %TW % ! 14 BE 	)	 r-   r   )*rh   chromadb.proto.utilsr   r1   r|   chromadb.ingestr   r   r   chromadb.proto.convertr   chromadb.proto.logservice_pb2r   r	   r
   "chromadb.proto.logservice_pb2_grpcr   %chromadb.telemetry.opentelemetry.grpcr   chromadb.typesr   r   chromadb.configr    chromadb.telemetry.opentelemetryr   r   r   r   	overridesr   typingr   r   r   uuidr   logging	getLoggerr   rQ   r   rd   r-   r,   <module>r      ss    
 A   
 3 U U = A #   + +  			8	$S 8 S r-   