Ë
    uPêjà  ã                   ó$  — d 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	 ddl
mZ ddlmZmZ de	d	z  d
ed	z  deeef   fd„Z edd¬«       G d„ d«      «       Z edd¬«       G d„ d«      «       Z edd¬«       G d„ d«      «       Z edd¬«       G d„ d«      «       Z edd¬«       G d„ d«      «       Z edd¬«       G d„ d«      «       Z edd¬«       G d„ d«      «       Zedeg ef   deed	f   fd„«       Zedeg ef   deed	f   fd„«       Zy	) aù  
Units of work over a single Prisma batch.

Each context manager here opens one ``db.batch_()`` and binds a typed write
repository per table to it, so every update queued through the yielded object
lands in the same transaction. The batch commits when the block exits cleanly
and is abandoned, writing nothing, when the block raises.

``spend_reset_unit_of_work`` covers the per-row key/user/team resets;
``budget_cascade_unit_of_work`` covers a budget tier's reset, where the
dependent spend and the tier's next window have to move together.

Each write repository queues narrow ``{spend}`` / ``{budget_reset_at}`` updates
instead of full-model writes, which trip ``prisma.errors.DataError`` on rows
carrying fields the update input type rejects (see #27730).
é    )ÚAsyncGeneratorÚCallableÚMapping)Úasynccontextmanager)Ú	dataclass)Údatetime)ÚFinal)Ú
BatchTableÚPrismaBatchÚbudget_reset_atNÚspend_decrementÚreturnc                 ó   — |�d|ind}|| dœS )NÚ	decrementr   )Úspendr   © )r   r   r   s      úf/home/oem/CyberBrain26-unified/.venv/lib/python3.12/site-packages/litellm/repositories/unit_of_work.pyÚ_spend_reset_datar      s)   € ð Ð&ð 
�oÑ&àð 
ð
 ¨Ñ?Ð?ó    T)ÚfrozenÚslotsc            	       ó@   — e Zd ZU eed<   	 ddededz  dedz  ddfd„Zy)	ÚKeySpendResetWritesÚtableNÚtokenr   r   r   c                 óV   — | j                   j                  d|it        ||«      ¬«       y )Nr   ©ÚwhereÚdata©r   Úupdater   )Úselfr   r   r   s       r   Úqueue_spend_resetz%KeySpendResetWrites.queue_spend_reset(   s-   € ð 	�
‰
×ÑØ˜EÐ"Ü" ?°OÓDð 	õ 	
r   ©N©	Ú__name__Ú
__module__Ú__qualname__r
   Ú__annotations__Ústrr   Úfloatr#   r   r   r   r   r   $   s=   … àÓð ]añ
Øð
Ø+3°d©?ð
ØMRÐUYÉ\ð
à	ô
r   r   c            	       ó@   — e Zd ZU eed<   	 ddededz  dedz  ddfd„Zy)	ÚUserSpendResetWritesr   NÚuser_idr   r   r   c                 óV   — | j                   j                  d|it        ||«      ¬«       y )Nr.   r   r    )r"   r.   r   r   s       r   r#   z&UserSpendResetWrites.queue_spend_reset5   ó-   € ð 	�
‰
×ÑØ˜gÐ&Ü" ?°OÓDð 	õ 	
r   r$   r%   r   r   r   r-   r-   1   ó=   … àÓð _cñ
Øð
Ø-5¸©_ð
ØOTÐW[É|ð
à	ô
r   r-   c            	       ó@   — e Zd ZU eed<   	 ddededz  dedz  ddfd„Zy)	ÚTeamSpendResetWritesr   NÚteam_idr   r   r   c                 óV   — | j                   j                  d|it        ||«      ¬«       y )Nr4   r   r    )r"   r4   r   r   s       r   r#   z&TeamSpendResetWrites.queue_spend_resetB   r0   r   r$   r%   r   r   r   r3   r3   >   r1   r   r3   c                   óP   — e Zd ZU eed<   deeef   ddfd„Zdeeef   de	ddfd„Z
y)ÚLinkedSpendResetWritesr   r   r   Nc                 óB   — | j                   j                  |ddi¬«       y )Nr   r   r   ©r   Úupdate_many)r"   r   s     r   Úqueue_spend_zeroz'LinkedSpendResetWrites.queue_spend_zeroO   s   € Ø�
‰
×Ñ U°'¸1°ÐÕ>r   Úamountc                 óF   — | j                   j                  |dd|ii¬«       y)z™``decrement`` rather than a read-then-set, so spend written between the
        cascade's read and its commit survives the reset instead of being erased.r   r   r   Nr9   )r"   r   r<   s      r   Úqueue_spend_decrementz,LinkedSpendResetWrites.queue_spend_decrementR   s+   € ð 	�
‰
×ÑØØ˜K¨Ð0Ð1ð 	õ 	
r   )r&   r'   r(   r
   r)   r   r*   Úobjectr;   r+   r>   r   r   r   r7   r7   K   sH   … àÓð? g¨c°6¨kÑ&:ð ?¸tó ?ð
¨7°3¸°;Ñ+?ð 
Èð 
ÐSWô 
r   r7   c                   ó,   — e Zd ZU eed<   dededdfd„Zy)ÚBudgetWindowWritesr   Ú	budget_idr   r   Nc                 óF   — | j                   j                  d|id|i¬«       y)zŽ``update_many`` so a tier deleted between the read and the commit is a
        no-op row count instead of a P2025 that aborts the whole chunk.rB   r   r   Nr9   )r"   rB   r   s      r   Úqueue_window_advancez'BudgetWindowWrites.queue_window_advance_   s'   € ð 	�
‰
×Ñ k°9Ð%=ÐEVÐXgÐDhÐÕir   )r&   r'   r(   r
   r)   r*   r   rD   r   r   r   rA   rA   [   s'   … àÓðj¨cð jÀHð jÐQUô jr   rA   c                   ó,   — e Zd ZU eed<   eed<   eed<   y)ÚSpendResetUnitOfWorkÚkeysÚusersÚteamsN)r&   r'   r(   r   r)   r-   r3   r   r   r   rF   rF   e   s   … à
ÓØÓØÔr   rF   c                   óX   — e Zd ZU dZeed<   eed<   eed<   eed<   eed<   eed<   eed<   y	)
ÚBudgetCascadeUnitOfWorka&  Every write a budget-tier reset performs, bound to one batch.

    The dependent spend rows and the budget rows' ``budget_reset_at`` advance
    must land together: advancing the window without zeroing the spend it
    gates leaves the dependents pinned at their cap until the next window.
    Úteam_membershipsrG   ÚorganizationsÚtagsÚmodel_access_groupsÚendusersÚbudgetsN)r&   r'   r(   Ú__doc__r7   r)   rA   r   r   r   rK   rK   l   s3   … ñð -Ó,Ø
 Ó Ø)Ó)Ø
 Ó Ø/Ó/Ø$Ó$ØÔr   rK   Ú	new_batchc                óè   K  —  | «       }t        t        |j                  ¬«      t        |j                  ¬«      t        |j                  ¬«      ¬«      ­–— |j                  «       ƒ d {  –—†  y 7 Œ­w)N©r   )rG   rH   rI   )rF   r   Úlitellm_verificationtokenr-   Úlitellm_usertabler3   Úlitellm_teamtableÚcommit©rS   Úbatchs     r   Úspend_reset_unit_of_workr\   ~   sU   è ø€ á‹K€EÜ
Ü  u×'FÑ'FÔGÜ"¨×)@Ñ)@ÔAÜ"¨×)@Ñ)@ÔAôó ð
 �,‰,‹.×Òús   ‚A(A2Á*A0Á+A2c                ó�  K  —  | «       }t        t        |j                  ¬«      t        |j                  ¬«      t        |j                  ¬«      t        |j
                  ¬«      t        |j                  ¬«      t        |j                  ¬«      t        |j                  ¬«      ¬«      ­–— |j                  «       ƒ d {  –—†  y 7 Œ­w)NrU   )rL   rG   rM   rN   rO   rP   rQ   )rK   r7   Úlitellm_teammembershiprV   Úlitellm_organizationtableÚlitellm_tagtableÚ#litellm_modelaccessgroupbudgettableÚlitellm_endusertablerA   Úlitellm_budgettablerY   rZ   s     r   Úbudget_cascade_unit_of_workrd   ‰   s�   è ø€ ñ ‹K€EÜ
!Ü/°e×6RÑ6RÔSÜ#¨%×*IÑ*IÔJÜ,°5×3RÑ3RÔSÜ#¨%×*@Ñ*@ÔAÜ2¸×9bÑ9bÔcÜ'¨e×.HÑ.HÔIÜ"¨×)BÑ)BÔCôó ð �,‰,‹.×Òús   ‚B<CÂ>CÂ?C)rR   Úcollections.abcr   r   r   Ú
contextlibr   Údataclassesr   r   Útypingr	   Ú%litellm.repositories.prisma_protocolsr
   r   r+   r*   r?   r   r   r-   r3   r7   rA   rF   rK   r\   rd   r   r   r   ú<module>rj      s§  ðñ÷" >Ñ =Ý *Ý !Ý Ý ç Ið@ x°$¡ð @ÈÐQUÉð @ÐZaÐbeÐgmÐbmÑZnó @ñ �$˜dÔ#÷	
ð 	
ó $ð	
ñ �$˜dÔ#÷	
ð 	
ó $ð	
ñ �$˜dÔ#÷	
ð 	
ó $ð	
ñ �$˜dÔ#÷
ð 
ó $ð
ñ �$˜dÔ#÷jð jó $ðjñ �$˜dÔ#÷ ð  ó $ð ñ �$˜dÔ#÷ ð  ó $ð ð" ð¨h°r¸;°Ñ.Gð ÈNÐ[oÐquÐ[uÑLvò ó ðð ðØ˜˜K˜Ñ(ðàÐ+¨TÐ1Ñ2òó ñr   