o
    1Xxi                     @   s   d dl 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 d d
lmZ d dlmZmZmZ d dl m!Z! d dl"Z"e"#e$Z%G dd deeZ&dS )    N)ProducerConsumerConsumerCallbackFnto_proto_submit)PushLogsRequestPullLogsRequest	LogRecord)LogServiceStub)OtelInterceptor)OperationRecordSeqId)System)OpenTelemetryClientOpenTelemetryGranularitytrace_method)override)SequenceOptionalcast)UUIDc                       s  e Zd ZU dZeed< ejed< eed< e	ed< de
f fddZed	ejed0 fdd	Zedejed0 fddZedejed0 fddZedejeded
dfddZedejededed
efddZedejededee d
ee fddZedeje			d1dededee dee dee d
efddZed ejed!ed
dfd"d Zed
efd#d$Zed
efd%d&Ze ed
e	fd'd(Z!ded)ee d
e	fd*d+Z"ded,e	d-e	d
ee# fd.d/Z$  Z%S )2
LogServicez(
    Distributed Chroma Log Service
    _log_service_stub_channel_log_service_url_log_service_portsystemc                    s8   |j d| _|j d| _|t| _t | d S )Nchroma_logservice_hostchroma_logservice_port)settingsrequirer   r   r   _opentelemetry_clientsuper__init__)selfr   	__class__ ]/home/air/biblejyuku/back/venv/lib/python3.10/site-packages/chromadb/logservice/logservice.pyr#   *   s   zLogService.__init__zLogService.startreturnNc                    sR   t | j d| j | _t g}t j| jg|R  | _t| j| _t	 
  d S )N:)grpcinsecure_channelr   r   r   r   intercept_channelr
   r   r"   start)r$   interceptorsr%   r'   r(   r.   0   s   zLogService.stopc                    s   | j   t   d S N)r   closer"   stopr$   r%   r'   r(   r2   ;   s   
zLogService.reset_statec                    s   t    d S r0   )r"   reset_stater3   r%   r'   r(   r4   A   s   zLogService.delete_logcollection_idc                 C   s   t d)NzNot implemented)NotImplementedError)r$   r5   r'   r'   r(   
delete_logF   s   zLogService.submit_embedding	embeddingc                 C   s    | j std| ||gd S )NComponent not runningr   )_runningRuntimeErrorsubmit_embeddings)r$   r5   r8   r'   r'   r(   submit_embeddingK   s   zLogService.submit_embeddings
embeddingsc                 C   sj   t dt| d|  | jstdt|dkrg S g }dd |D }|| |ttt	 | |S )NzSubmitting z" embeddings to log for collection r9   r   c                 S   s   g | ]}t |qS r'   r   ).0recordr'   r'   r(   
<listcomp>f   s    z0LogService.submit_embeddings.<locals>.<listcomp>)
loggerinfolenr:   r;   append	push_logsr   r   r   )r$   r5   r>   countsprotos_to_submitr'   r'   r(   r<   U   s    	zLogService.subscribe
consume_fnr.   endidc                 C   s   t d| d tddS )NzSubscribing to log for , noop for logservicer   )int)rB   rC   r   )r$   r5   rI   r.   rJ   rK   r'   r'   r(   	subscriber   s   

zLogService.unsubscribesubscription_idc                 C   s   t d| d d S )NzUnsubscribing from rL   )rB   rC   )r$   rO   r'   r'   r(   unsubscribe   s   c                 C      dS )Nr   r'   r3   r'   r'   r(   	min_seqid   s   zLogService.min_seqidc                 C   s   t jS r0   )sysmaxsizer3   r'   r'   r(   	max_seqid   s   zLogService.max_seqidc                 C   rQ   )Nia  r'   r3   r'   r'   r(   max_batch_size   s   zLogService.max_batch_sizerecordsc                 C   s"   t t||d}| j|}|jS )N)r5   rW   )r   strr   PushLogsrecord_count)r$   r5   rW   requestresponser'   r'   r(   rF      s   zLogService.push_logsstart_offset
batch_sizec                 C   s*   t t|||t d}| j|}|jS )N)r5   start_from_offsetr^   end_timestamp)r   rX   timetime_nsr   PullLogsrW   )r$   r5   r]   r^   r[   r\   r'   r'   r(   	pull_logs   s   zLogService.pull_logs)r)   N)NNN)&__name__
__module____qualname____doc__r
   __annotations__r+   ChannelrX   rM   r   r#   r   r   ALLr   r.   r2   r4   r   r7   r   r   r=   r   r<   r   r   rN   rP   rR   rU   propertyrV   rF   r	   rd   __classcell__r'   r'   r%   r(   r       s   
 


	





r   )'rS   r+   ra   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   	overridesr   typingr   r   r   uuidr   logging	getLoggerre   rB   r   r'   r'   r'   r(   <module>   s"    
