
    Cj2K                         S r SSKrSSKrSSKrSSKJrJrJrJr  SSK	J
r
  SSKJr  SSKJr  SSKJr  SSKJrJr  SS	KJrJrJr  SS
KJrJr  SSKJrJr   " S S\5      rg)ao  
PostHog Integration - sends LLM analytics events to PostHog

Follows PostHog's LLM Analytics format: https://posthog.com/docs/llm-analytics/manual-capture

async_log_success_event: stores batch of events in memory and flushes to PostHog
async_log_failure_event: logs failed LLM calls with error information

For batching specific details see CustomBatchLogger class
    N)AnyDictOptionalTuple)verbose_logger)uuid)CustomBatchLogger)
safe_dumps)should_use_posthog_mockcreate_mock_posthog_client)_get_httpx_clientget_async_httpx_clienthttpxSpecialProvider)POSTHOG_MAX_BATCH_SIZEPostHogEventPayload)StandardCallbackDynamicParamsStandardLoggingPayloadc            
         ^  \ rS rSrU 4S jrS rS rS r SS jrS\	\
\4   S\4S	 jrS
\S\	\
\4   S\
S\	\
\4   4S jrS\	\
\4   S\	\
\4   4S jrS\	\
\4   S\	\
\4   4S jrS
\S\	\
\4   S\
4S jrS\	\
\4   S\\\
   \\
   4   4S jrS rS rS\	\
\4   S\	\
\4   4S jrS\
4S jrS\S\
S\	\
\4   4S jrS S\S\
S\S\4S jjrS rSrU =r$ )!PostHogLogger%   c                   >  [         R                  " S5        [        5       U l        U R                  (       a   [	        5         [         R                  " S5        [
        R                  " SS5      c  [        S5      e[        [        R                  S9U l        [        5       U l        [
        R                  " S5      U l        [
        R                  " SS5      nUR                  S	5      U l        U R                    S
3U l        SU l        SU l        / U l        [*        R,                  " U R.                  5        [0        TU ]d  " S0 UDS[4        S.D6  g! [         a)  n[         R6                  " S[9        U5       35        UeSnAff = f)z
Initializes the PostHog logger, checks if the correct env variables are set

Required environment variables:
`POSTHOG_API_KEY` - your PostHog API key
`POSTHOG_API_URL` - your PostHog API URL (defaults to https://app.posthog.com)
zPostHog: in init posthog loggerz6[POSTHOG MOCK] PostHog logger initialized in mock modePOSTHOG_API_KEYNz4POSTHOG_API_KEY is not set, set 'POSTHOG_API_KEY=<>')llm_providerPOSTHOG_API_URLzhttps://us.i.posthog.com//batch/F)
flush_lock
batch_sizez.PostHog: Got exception on init PostHog client  )r   debugr   is_mock_moder   osgetenv	Exceptionr   r   LoggingCallbackasync_clientr   sync_clientr   rstripposthog_hostcapture_url_async_initializedr   	log_queueatexitregister_flush_on_exitsuper__init__r   	exceptionstr)selfkwargsposthog_api_urle	__class__s       n/home/rurouni/.local/share/pipx/venvs/strix-agent/lib/python3.13/site-packages/litellm/integrations/posthog.pyr1   PostHogLogger.__init__&   sQ   &	  !BC 7 9D  *,$$L yy*D19 VWW 61AA!D  12D#%99->#?D  ii(9;UVO / 6 6s ;D"&"3"3!4G<D&+D#"DODN OOD//0G %)6L  	$$@QI G		s   EE 
F$F  Fc                     [         R                  " SU5        U R                  U5      u  pVUb  Uc  [        S5      eU R	                  U5      nSS0nU R                  U/U5      n	UR                  S5       S3n
U R                  R                  U
[        U	5      US9nUR                  5         UR                  S:w  a%  [        S	UR                   S
UR                   35      eU R                  (       a  [         R                  " S5        g [         R                  " S5        g ! [         a,  n[         R                  " S[        U5       35         S nAg S nAff = f)Nz<PostHog: Sync logging - Enters logging function for model %sz'PostHog credentials not found in kwargsContent-Typeapplication/jsonr   r   urlcontentheaders   'Response from PostHog API status_code: , text: z-[POSTHOG MOCK] Sync event successfully mockedz%PostHog: Sync event successfully sentzPostHog Sync Layer Error - )r   r    _get_credentials_for_requestr$   create_posthog_event_payload_create_posthog_payloadr(   r'   postr
   raise_for_statusstatus_codetextr!   r2   r3   )r4   r5   response_obj
start_timeend_timeapi_keyapi_urlevent_payloadrA   payloadr*   responser7   s                r9   log_success_eventPostHogLogger.log_success_eventV   s[   #	M  NPV  $@@HG'/ IJJ ==fEM  2G 22M?GLG$^^C019K'',,"7+ - H
 %%'##s*=h>R>R=SS[\d\i\i[jk    $$%TU$$%LM 	M$$'B3q6(%KLL	Ms   DD! 
D! !
E+"EEc                   #     [         R                  " SU5        U R                  5         U R                  XX45      I S h  vN   g  N! [         a,  n[         R
                  " S[        U5       35         S nAg S nAff = f7fNz=PostHog: Async logging - Enters logging function for model %szPostHog Layer Error - r   r    _ensure_async_setup_log_async_eventr$   r2   r3   r4   r5   rL   rM   rN   r7   s         r9   async_log_success_event%PostHogLogger.async_log_success_event|   n     	  OQW $$&''jSSS 	$$'=c!fX%FG	@   B<A  AA BA 
A>"A94B9A>>Bc                   #     [         R                  " SU5        U R                  5         U R                  XX45      I S h  vN   g  N! [         a,  n[         R
                  " S[        U5       35         S nAg S nAff = f7frW   rX   r[   s         r9   async_log_failure_event%PostHogLogger.async_log_failure_event   r^   r_   c                 T  #    U R                  U5      u  pVU R                  U5      nU R                  R                  XuUS.5        [        R
                  " SU R                   S35        [        U R                  5      U R                  :  a  U R                  5       I S h  vN   g g  N7f)N)eventrO   rP   z-PostHog, event added to queue. Will flush in z seconds...)
rE   rF   r,   appendr   r    flush_intervallenr   flush_queue)r4   r5   rL   rM   rN   rO   rP   rQ   s           r9   rZ   PostHogLogger._log_async_event   s       <<VD99&A 	#GL	
 	;D<O<O;PP[\	
 t~~$//1""$$$ 2$s   BB(B& B(r5   returnc                     UR                  SS5      nUc  [        S5      eUR                  SS5      nUS:X  a  SOSnU R                  UUUS	9nU R                  X!5      n[	        UUUS
9$ )z
Helper function to create a PostHog event payload for logging

Args:
    kwargs (Dict[str, Any]): request kwargs containing standard_logging_object

Returns:
    PostHogEventPayload: defined in types.py
standard_logging_objectNz+standard_logging_object not found in kwargs	call_type 	embeddingz$ai_embedding$ai_generation)rl   r5   
event_name)rd   
propertiesdistinct_id)get
ValueError_create_posthog_properties_get_distinct_idr   )r4   r5   rl   rm   rq   rr   rs   s          r9   rF   *PostHogLogger.create_posthog_event_payload   s     EKJJ%tE
 #*JKK+//R@	(1[(@_FV
44$;! 5 

 ++,CL"!#
 	
    rl   rq   c                 @   0 nU R                  USS5      US'   U R                  USS5      US'   U R                  US5      nUb  XTS'   US:X  a  U R                  US	5      nUb  XdS
'   U R                  USS5      US'   US:X  a  U R                  USS5      US'   U R                  US5      nUb  XtS'   U R                  USS5      US'   U R                  US5      S:X  a  SUS'   U R                  US5      nUb  XS'   U R                  XB5        U R                  XB5        U$ )z6Create PostHog properties following LLM Analytics specmodelrn   z	$ai_modelcustom_llm_providerz$ai_providermessagesz	$ai_inputrp   rS   z$ai_output_choicesprompt_tokensr   z$ai_input_tokenscompletion_tokensz$ai_output_tokensresponse_costz$ai_total_cost_usdresponse_time        z$ai_latencystatusfailureTz$ai_is_error	error_strz	$ai_error)	_safe_get_add_trace_properties_add_custom_metadata_properties)	r4   rl   r5   rq   rr   r}   rS   r   r   s	            r9   rv   (PostHogLogger._create_posthog_properties   sn    
 #'..1H'SU"V
;%)^^#%:B&

>"
 >>"9:F&.{#))~~&=zJH#3;/0 *.#_a*

%& )).2nn')<a/J*+
 '>P$/<+,$(NN#_c%

=!
 >>18<	I)-J~&'>LI$*3;' 	"":6 	,,Z@ry   rr   c                 J   U R                  US0 5      nU R                  USU R                  5       5      nXAS'   U R                  USU R                  5       5      nXQS'   U R                  U5      nUR                  S5      =(       d    UR                  S5      nU(       a  XqS'   g g )	Nrl   trace_idz$ai_trace_ididz$ai_span_idparent_run_id	parent_idz$ai_parent_id)r   
_safe_uuid_extract_metadatart   )r4   rr   r5   rl   r   span_idmetadatar   s           r9   r   #PostHogLogger._add_trace_properties  s    "&..9RTV"W>>#Z1B
 &.>"..!8$@QR$+=!))&1LL1NX\\+5N	*3' ry   c                     U R                  U5      n[        U[        5      (       d  g1 SknUR                  5        H  u  pVXT;  d  M  XaU'   M     g)z0Add custom metadata fields to PostHog propertiesN>    rA   user_idapi_baseendpointr   
deployment
model_infomodel_groupuser_api_keyhidden_paramsr   caching_groupsmodel_group_sizeuser_api_key_hashuser_api_key_aliasuser_api_key_spendlitellm_api_versionuser_api_key_org_idrequester_ip_addressuser_api_key_team_iduser_api_key_user_iduser_api_key_metadatauser_api_key_max_budgetuser_api_key_team_aliasuser_api_key_team_spenduser_api_key_user_emaillitellm_parent_otel_spanuser_api_key_end_user_idglobal_max_parallel_requestsuser_api_end_user_max_budgetuser_api_key_team_max_budgetuser_api_key_model_max_budget)r   
isinstancedictitems)r4   rr   r5   r   litellm_internal_fieldskeyvalues          r9   r   -PostHogLogger._add_custom_metadata_properties  sO     ))&1(D))"#
H #..*JC1"'3 +ry   c                    U R                  U5      nU R                  US5      nU(       a  [        U5      $ U R                  US5      nU(       a  [        U5      $ U R                  US5      nU(       a  [        U5      $ U R                  5       $ )Nr   end_userr   )r   r   r3   r   )r4   rl   r5   r   r   r   r   s          r9   rw   PostHogLogger._get_distinct_idA  sz     ))&1..95w<>>"9:Fx= >>"9:Fx=   ry   c                     UR                  SS5      nUbK  UR                  S5      =(       d    U R                  nUR                  S5      =(       d    U R                  nX44$ U R                  nU R                  nX44$ )a0  
Get PostHog credentials for this request.

Checks for per-request credentials in standard_callback_dynamic_params,
falls back to instance defaults from environment variables.

Args:
    kwargs: Request kwargs containing standard_callback_dynamic_params

Returns:
    tuple[str, str]: (api_key, api_url)
 standard_callback_dynamic_paramsNposthog_api_keyr6   )rt   r   r)   )r4   r5   r   rO   rP   s        r9   rE   *PostHogLogger._get_credentials_for_requestQ  s      JJ94@ 	) ,70445FG ('' 
 1445FG %$$   **G''Gry   c                   #     U R                   (       d  g[        R                  " S[        U R                   5       S35        U R                  (       a  [        R                  " S5        0 nU R                    H,  nUS   US   4nX1;  a  / X'   X   R                  US   5        M.     UR                  5        H  u  u  pEnSS	0nU R                  Xd5      nUR                  S
5       S3n	U R                  R                  U	[        U5      US9I Sh  vN n
U
R                  5         U
R                  S:w  d  M  [        SU
R                   SU
R                   35      e   U R                  (       a.  [        R                  " S[        U R                   5       S35        g[        R                  " S[        U R                   5       S35        g N! [         a,  n[        R                   " S[#        U5       35         SnAgSnAff = f7f)z~
Sends the in memory logs queue to PostHog API

Raises:
    Raises a NON Blocking verbose_logger.exception if an error occurs
NzPostHog: Sending batch of z eventsz@[POSTHOG MOCK] Mock mode enabled - API calls will be interceptedrO   rP   rd   r<   r=   r   r   r>   rB   rC   rD   z[POSTHOG MOCK] Batch of z events successfully mockedzPostHog: Batch of z events successfully sentz"PostHog Error sending batch API - )r,   r   r    rg   r!   re   r   rG   r(   r&   rH   r
   rI   rJ   r$   rK   r2   r3   r4   batches_by_credentialsitemr   rO   rP   eventsrA   rR   r*   rS   r7   s               r9   async_send_batchPostHogLogger.async_send_batchs  s    3	T>>  ,S-@,AI   $$V
 CE"IY8424*/&+224=A	 ' /E.J.J.L*"F"$6 66vG!(!4 5W=!%!2!2!7!7#&w/# "8 " 
 ))+''3.#A(BVBVAWW_`h`m`m_no ! /M(   $$.s4>>/B.CC^_ $$(T^^)<(==VW#(  	T$$'I#a&%RSS	Ts^   HG HC:G G"G 7A%G H-G 
HG 
H"G>9H>HHc                 J   U R                   (       d]   [        R                  " 5       U l        [        R                  " U R                  5       5        SU l         [        R                  " S5        g g ! [         a(  n[        R                  " S[        U5       35        e S nAff = f)NTz%PostHog: Async components initializedz0PostHog: Failed to initialize async components: )r+   asyncioLockr   create_taskperiodic_flushr   r    r$   errorr3   )r4   r7   s     r9   rY   !PostHogLogger._ensure_async_setup  s    &&	"),,.##D$7$7$9:*.'$$%LM '  $$Fs1vhO 	s   AA0 0
B":#BB"c                 n    UR                  S0 5      =(       d    0 nUR                  S0 5      =(       d    0 $ )Nlitellm_paramsr   )rt   )r4   r5   r   s      r9   r   PostHogLogger._extract_metadata  s0    $4b9?R!!*b17R7ry   c                 >    [        [        R                  " 5       5      $ N)r3   r   uuid4)r4   s    r9   r   PostHogLogger._safe_uuid  s    4::<  ry   r   rO   c                 
    X!S.$ )N)rO   batchr   )r4   r   rO   s      r9   rG   %PostHogLogger._create_posthog_payload  s    "44ry   objr   defaultc                 P    Ub  [        US5      (       d  U$ UR                  X#5      $ )Nrt   )hasattrrt   )r4   r   r   r   s       r9   r   PostHogLogger._safe_get  s'    ;gc511Nwws$$ry   c                    U R                   (       d  g[        R                  " S[        U R                   5       S35         0 nU R                    H,  nUS   US   4nX1;  a  / X'   X   R	                  US   5        M.     UR                  5        H  u  u  pEnSS0nU R                  Xd5      nUR                  S	5       S
3n	U R                  R                  U	[        U5      US9n
U
R                  5         U
R                  S:w  d  Mw  [        R                  " SU
R                   35        M     U R                  (       a.  [        R                  " S[        U R                   5       S35        O-[        R                  " S[        U R                   5       S35        U R                   R                  5         g! [          a,  n[        R                  " S[#        U5       35         SnAgSnAff = f)ai  
Flush remaining events from internal log_queue before process exit.
Called automatically via atexit handler.

This works in conjunction with GLOBAL_LOGGING_WORKER's atexit handler:
1. GLOBAL_LOGGING_WORKER atexit invokes pending callbacks
2. Callbacks add events to this logger's internal log_queue
3. This atexit handler flushes the internal queue to PostHog
NzPostHog: Flushing z remaining events on exitrO   rP   rd   r<   r=   r   r   r>   rB   z*PostHog: Failed to flush on exit - status z$[POSTHOG MOCK] Successfully flushed z events on exitzPostHog: Successfully flushed z(PostHog: Error flushing events on exit: )r,   r   r    rg   re   r   rG   r(   r'   rH   r
   rI   rJ   r   r!   clearr$   r3   r   s               r9   r/   PostHogLogger._flush_on_exit  s    ~~ T^^!4 55NO	
)	VBD"IY8424*/&+224=A	 ' /E.J.J.L*"F"$6 66vG!(!4 5W=++00#&w/# 1 
 ))+''3."((DXEYEYDZ[! /M(   $$:3t~~;N:O_ $$4S5H4IY NN  " 	V  #KCPQF8!TUU	Vs    CF4 B-F4 4
G*>"G%%G*)	r   r+   r&   r*   r   r!   r,   r)   r'   )Nr   r   r   ) __name__
__module____qualname____firstlineno__r1   rT   r\   ra   rZ   r   r3   r   r   rF   r   rv   r   r   rw   r   r   rE   r   rY   r   r   listrG   r   r/   __static_attributes____classcell__)r8   s   @r9   r   r   %   s   .`$ML		 CF%$!
38n!
	!
F8!78 S#X8 	8
 
c3h8t4S#X 4SRUX 4 .(sCx..(26sCx..(`!'=!GKCQTH~!	!   38n  	x}hsm+	,  D:Tx8S#X 84S> 8!C !5d 5S 5T#s(^ 5%S %s %S %C %
:V :Vry   r   )__doc__r   r-   r"   typingr   r   r   r   litellm._loggingr   litellm._uuidr   (litellm.integrations.custom_batch_loggerr	   *litellm.litellm_core_utils.safe_json_dumpsr
   (litellm.integrations.posthog_mock_clientr   r   &litellm.llms.custom_httpx.http_handlerr   r   r   "litellm.types.integrations.posthogr   r   litellm.types.utilsr   r   r   r   ry   r9   <module>r      sR   	   	 - - +  F A 
 V`V% `Vry   