
    \jw                     *   d dl Z d dlZd dl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mZmZmZmZmZmZ d dlZd dlmZmZmZmZ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&m'Z'm(Z(m)Z) d dl*m+Z+ erd dl,m-Z- d dl	m.Z. d dl%m/Z/m0Z0m1Z1m2Z2m3Z3m4Z4m5Z5 dZ6	 d dl7m8Z9 dZ6e)jv                  Z< G d de      Z=	 	 	 	 dBdZ>dCdZ?ddde@fdZAdDdZBdEdZCdEd ZDdFd!ZEd"ddeFfd#ZGdGd$ZHdHd&ZI	 	 	 	 dId'ZJ	 	 	 	 dJd(ZK	 	 	 	 	 	 	 	 dKd+ZLd,e@d-d%dd.fd/ZMdLd0ZN	 	 	 	 dMd1ZO	 	 	 	 dNd2ZP	 	 	 	 dNd3ZQd4d5d6eFdd7fd8ZR	 	 	 	 dOd9ZSd:d.d;dde@ddd)d*ddfd<ZT	 	 	 	 	 	 	 	 dPd=ZU	 	 	 	 	 	 dQd>ZV	 	 	 	 	 	 dRd?ZW	 	 	 	 	 	 	 	 	 	 dSd@ZX	 	 	 	 	 	 	 	 dTdAZYy# e:$ r Y w xY w)U    N)wraps   )ORIGINTOOL_ATTRIBUTES_MAPGEN_AI_SYSTEM)BLOB_DATA_SUBSTITUTE)	TYPE_CHECKINGIterableAnyCallableListOptionalUnion	TypedDictDict)set_data_normalizedtruncate_and_annotate_messagesnormalize_message_rolestransform_google_content_partget_modality_from_mime_type)OPSPANDATA)should_send_default_pii)capture_internal_exceptionsevent_from_exceptionsafe_serialize)GenerateContentConfigPartContentPartDict)chain)Span)TextPart)GenerateContentResponseContentListUnionContentUnionDictToolModelEmbedContentResponseContentUnionF)ImageTc                   D    e Zd ZU dZeed<   eed<   eed<   eed<   eed<   y)	UsageDatazStructure for token usage data.input_tokensinput_tokens_cachedoutput_tokensoutput_tokens_reasoningtotal_tokensN)__name__
__module____qualname____doc__int__annotations__     T/root/env/lib/python3.12/site-packages/sentry_sdk/integrations/google_genai/utils.pyr-   r-   ?   s#    )  r:   r-   returnc                    t        ddddd      }t        | t              r| j                  di       }|s|S |j                  dd      xs d}|j                  dd      xs d}||z   |d<   |j                  dd      xs d}||d<   |j                  d	d      xs d}||d
<   |j                  dd      xs d}||z   |d<   |j                  dd      xs d}||d<   |S t	        | d      s|S | j
                  }t        |dd      xs d}t        |dd      xs d}||z   |d<   t        |dd      xs d}||d<   t        |d	d      xs d}||d
<   t        |dd      xs d}||z   |d<   t        |dd      xs d}||d<   |S )a@  Extract usage data from response into a structured format.

    Args:
        response: The GenerateContentResponse object or dictionary containing usage metadata

    Returns:
        UsageData: Dictionary with input_tokens, input_tokens_cached,
                   output_tokens, and output_tokens_reasoning fields
    r   )r.   r/   r0   r1   r2   usage_metadataprompt_token_counttool_use_prompt_token_countr.   cached_content_token_countr/   thoughts_token_countr1   candidates_token_countr0   total_token_countr2   )r-   
isinstancedictgethasattrr>   getattr)	response
usage_datausageprompt_tokenstool_use_prompt_tokenscached_tokensreasoning_tokenscandidates_tokensr2   s	            r;   extract_usage_datarR   I   s     !J (D!-r2		"6:?a!&+H!!L!QPQ%25K%K
>"		">BGa,9
() 99%;Q?D10@
,-!II&>BGa '8:J&J
?#yy!4a8=A%1
>"8-.##E E#7;@qM$U,I1MRQR!.1G!GJ~ E#?CHqM(5J$% u&<a@EA,<J()  '?CHq"36F"FJ5"5q9>QL!-J~r:   r   c                     t        | t        j                         j                  ddd      \  }}t        j                  ||       y)z.Capture exception with Google GenAI mechanism.google_genaiF)typehandled)client_options	mechanism)hintN)r   
sentry_sdk
get_clientoptionscapture_event)exceventrY   s      r;   _capture_exceptionr`      s?    &!,,.66)e<KE4
 U.r:   modelzUnion[str, Model]c                 ~    t        | t              r| S t        | d      rt        | j                        S t        |       S )z(Extract model name from model parameter.name)rE   strrH   rc   )ra   s    r;   get_model_namere      s3    %uf5::u:r:   contentsr%   c                    | g S g }t        | t              rd| dgS t        | t              rg| r@t        d | D              r.g }| D ]!  }t	        |      }||j                  |       # d|dgS | D ]  }t        |      }|j                  |         |S t        | t              r| j                  dd      }| j                  d      }|rg }g }|D ]W  }t        |      }	|	t        |	t              r&|	j                  d      dk(  r|j                  |	       G|j                  |	       Y |r!|dk(  rd	n|xs d}
|j                  |
|d       |j                  |       |S d
| v r|j                  || d
   d
dgd       |S d| v r-|j                  d| d   j                  dd      t        di       |S t        | d      r| j                  rt        | dd      xs d}g }g }| j                  D ]W  }t        |      }	|	t        |	t              r&|	j                  d      dk(  r|j                  |	       G|j                  |	       Y |r|dk(  rd	n|}
|j                  |
|d       |j                  |       |S t        |       }	|	r.t        |	t              r|	j                  d      dk(  r|	gS d|	gdgS t        r.t        | t         j"                        rt%        |       }|rd|gdgS t        | d      rCt        | d      r7t        | dd      }t        | dd      }||d}dt'        |      ||d}d|gdgS t        | d
      r)| j(                  rdt        | j(                        d
dgdgS g S )a$  Extract messages from contents parameter which can have various formats.

    Returns a list of message dictionaries in the format:
    - System: {"role": "system", "content": "string"}
    - User/Assistant: {"role": "user"|"assistant", "content": [{"text": "...", "type": "text"}, ...]}
    Nuserrolecontentc              3   2   K   | ]  }t        |        y wN)_is_part_like).0items     r;   	<genexpr>z,extract_contents_messages.<locals>.<genexpr>   s     EDM$/Es   rj   partstoolra   	assistanttextru   rU   inline_data	mime_type rx   dataurirU   modalityrx   r|   )rE   rd   listall_extract_part_from_itemappendextract_contents_messagesextendrF   rG   _extract_part_contentr   rH   rr   rI   _is_PIL_availablePILImager+   _extract_pil_imager   ru   )rf   messagescontent_partsrp   partitem_messagesrj   rr   tool_messagespart_resultnormalized_role	blob_partfile_urirx   s                 r;   r   r      s    	H (C H566 (D!EHEEM  /.t4#!((./
 $>?? ! / 9$ ?./ O (D!||FF+W%MM 
63D9&k40[__V5LPV5V!((5 "((5
6 15+dnf] ST OOM*( ' xOO )1&)96 JK$  h& OO!%-m%<%@%@b%Q 4$  x!hnnx.8&NN 	2D/5K"+t,1HF1R$$[1$$[1	2 -1W_k$OOO_OP& (1Kk4([__V-D-N= #>?? Z(..A&x0	#<== xGHk$B8UD1Hk48	 	 7	B&	I $<== x X]]#hmm2Df)U(VW
 	
 Ir:   c                    | yt        | t              rEd| v rt        |       S | j                  d      r| d   ddS t	        |       }|d| v r	t
        |d<   |S yt        | d      r| j                  rt        |       S t        | d      r| j                  r| j                  ddS t        | d      rH| j                  r<| j                  }t        |dd      }t        |d	d      }||d
}dt        |      ||dS t        | d      rL| j                  r@| j                  }t        |dd      }t        |d	d      }||d
}dt        |      |t
        dS y)zExtract content from a Part object or dict.

    Returns:
        - dict for content part (text/blob) or tool message
        - None if part should be skipped
    Nfunction_responseru   rv   rw   rk   	file_datar   rx   ry   r|   r}   r{   blobrU   r~   rx   rk   )rE   rF   _extract_tool_message_from_partrG   r   r   rH   r   ru   r   rI   r   rw   )r   resultr   r   rx   rw   r{   s          r;   r   r   =  s    | $$&248888F L&99 /t4$ %9y!M t()d.D.D.t44 tV		622 t[!dnnNN	9j$7I{D9	 	 7	B&	  t]#(8(8&&{FD1Kd;	 	 7	B&/	  r:   c                 8   d}t        | t              r| j                  d      }nt        | d      r| j                  }|syd}d}d}t        |t              rG|j                  d      }|j                  d      }|j                  di       }|j                  d|      }nPt        |dd      }t        |dd      }t        |dd      }|i }t        |t              r|j                  d|      }n|}|syd|rt        |      ndt        |      |t        |      nddd	S )
zExtract tool message from a Part with function_response.

    Returns:
        {"role": "tool", "content": {"toolCallId": "...", "toolName": "...", "output": "..."}}
        or None if not a valid tool message
    Nr   idrc   rJ   outputrs   )
toolCallIdtoolNamer   ri   )rE   rF   rG   rH   r   rI   rd   r   )r   r   tool_call_id	tool_namer   response_dictresponse_objs          r;   r   r     s8    $ HH%89	*	+ 22 LIF#T*(,,T2%))&1	)--j"=""8]; 0$=-vt<	0*dCLlD)!%%h=F!F /;#l+I060BnV,
 r:   c                     t         rt        | t        j                        sy| j                  xs d}|j                         }d| }dt        |      |t        dS )z'Extract blob part from PIL.Image.Image.NJPEGzimage/r   r   )r   rE   r   r+   formatlowerr   r   )image
format_strsuffixrx   s       r;   r   r     s\    Juhnn$E 'JF!I /	:'	 r:   rp   c                 .   t        | t        t        f      ryt        | t        t        f      ryt        | t
              r)d| v sd| v ryt        t        | j                         z        S t        | d      ryt        rt        | t        j                        ryy)zZCheck if item is a part-like value (PartUnionDict) rather than a Content/multi-turn entry.TFrj   rr   r|   )rE   rd   r   r   r   rF   bool_PART_DICT_KEYSkeysrH   r   r   r+   )rp   s    r;   rn   rn     sw    $d$$w($T>W_Odiik122tUZhnn=r:   c                    t        | t              r| ddS t        | t              r"d| v rd| d   j                  dd      t        diS t        |       }||S t        r%t        | t        j                        rt        |       S t        | d      r<t        | d      r0t        | dd      }t        | dd      xs d}|dt        |      ||d	S y)
z7Convert a single part-like item to a content part dict.ru   rv   rw   rx   ry   rz   Nr|   r}   )rE   rd   rF   rG   r   r   r   r   r+   r   rH   rI   r   )rp   r   r   rx   s       r;   r   r     s    $f-- $-4"7!-044["E,
 	
 #4(F Zhnn=!$'' tUk :4-D+t4:	7	B&	  r:   Optional[str]c                 ~   t        |       }|syg }|D ]  }|j                  d      }t        |t              r|j	                  |       6t        |t
              sG|D ]I  }t        |t              s|j                  d      dk(  s)|j	                  |j                  dd             K  |rdj                  |      S dS )zExtract text from contents parameter which can have various formats.

    This is a compatibility function that extracts text from messages.
    For new code, use extract_contents_messages instead.
    Nrk   rU   ru   ry    )r   rG   rE   rd   r   r   rF   join)rf   r   textsmessagerk   r   s         r;   extract_contents_textr     s     )2HE 7++i(gs#LL!& 7dD)dhhv.>&.HLL&"!5677 $388E?--r:   c           
         g }| D ]  }t        |      r+|j                  t        |dd      t        |dd      d       9t        |d      rH|j                  <|j                  D ],  }|j                  t        |dd      t        |dd      d       . t        j                         D ](  \  }}t        ||d      s|j                  ||d          |r|S dS )	z%Format tools parameter for span data.r3   unknownr6   N)rc   descriptionfunction_declarationsrc   r   )callabler   rI   rH   r   r   items)toolsformatted_toolsrs   	func_decl	attr_namer   s         r;   _format_tools_for_spanr   $  s     O  D>""#D*i@#*4D#A D12**6 "77 	&& '	64 @'.y-'N +>*C*C*E &	;4D1#**$-+6 1 D .?747r:   c                    g }t        | dg       r| j                  D ]  }t        |d      rt        |j                  dg       s'|j                  j                  D ]b  }t        |dd      s|j
                  }t        |dd      dd}t        |dd      rt        |j                        |d	<   |j                  |       d  t        | d
d      r| j                  D ]  }t        |dd      st        |dg       D ]a  }t        |dd      s|j
                  }t        |dd      dd}t        |d      rt        |j                        |d	<   |j                  |       c  |r|S dS )z\Extract tool/function calls from response candidates and automatic function calling history.
candidatesrk   rr   function_callNrc   )rc   rU   args	arguments"automatic_function_calling_history)
rI   r   rH   rk   rr   r   r   r   r   r   )rJ   
tool_calls	candidater   r   	tool_callrk   s          r;   extract_tool_callsr   N  s~   
 J xr*!,, 	1I9i0!!7B9 !))// 14$7$($6$6M 'vt D /!I }fd;1?@R@R1S	+.%%i01	1, x=tDBB 	1G7GT2"5 14$7$($6$6M 'vt D /!I }f51?@R@R1S	+.%%i01		1$ $:--r:   kwargsdict[str, Any]c                 (   |r|j                         ni }| rf	 t        j                  |      }t        |j                  j                               }t        |       D ]  \  }}|t        |      k  s||||   <    	 |S |S # t        $ r	 | |d<   Y |S w xY w)z(Capture tool input from args and kwargs.r   )	copyinspect	signaturer   
parametersr   	enumeratelen	Exception)r   r   rs   
tool_inputsigparam_namesiargs           r;   _capture_tool_inputr     s     #)bJ 	&##D)Cs~~2245K#D/ 53s;''14J{1~.5 :	  	&!%Jv		&s   AA? 0
A? ?BBr   tool_docr"   c                     t        j                  t        j                  d|  t              }|j                  t        j                  |        |r |j                  t        j                  |       |S )z!Create a span for tool execution.zexecute_tool )oprc   origin)	rZ   
start_spanr   GEN_AI_EXECUTE_TOOLr   set_datar   GEN_AI_TOOL_NAMEGEN_AI_TOOL_DESCRIPTION)r   r   spans      r;   _create_tool_spanr     sZ      !!YK(D
 	MM(++Y7h66AKr:   c                      t               s S t         dd       j                  t        j                         rt               d fd       }|S t               d fd       }|S )z3Wrap a tool to emit execute_tool spans when called.r3   r   c                    K   t              5 }t        | |      }t               5  |j                  t        j
                  t        |             d d d        	  | i | d {   }t               5  |j                  t        j                  t        |             d d d        |cd d d        S # 1 sw Y   axY w7 W# 1 sw Y   "xY w# t        $ r}t        |        d }~ww xY w# 1 sw Y   y xY wwrm   
r   r   r   r   r   GEN_AI_TOOL_INPUTr   GEN_AI_TOOL_OUTPUTr   r`   	r   r   r   r   r   r^   rs   r   r   s	         r;   async_wrappedz#wrapped_tool.<locals>.async_wrapped  s     "9h7 40vtD
02 MM 22N:4N
#'#8#88F 56 $779O
 "#   9  ! &s+% s{   C5C)*B2C)C&B>'C5*C 	C(
C52B;	7C)>C C		C	C&C!!C&&C))C2.C5c                     t              5 }t        | |      }t               5  |j                  t        j
                  t        |             d d d        	  | i |}t               5  |j                  t        j                  t        |             d d d        |cd d d        S # 1 sw Y   YxY w# 1 sw Y    xY w# t        $ r}t        |        d }~ww xY w# 1 sw Y   y xY wrm   r   r   s	         r;   sync_wrappedz"wrapped_tool.<locals>.sync_wrapped  s    "9h7 40vtD
02 MM 22N:4N
!4262F 56 $779O
 "#    ! &s+% sX   C*B(CC +*B4	C (B1	-C4B=	9C  	C	CCCC&)r   r   r   r   r<   r   )r   rI   r6   r   iscoroutinefunctionr   )rs   r   r   r   r   s   `  @@r;   wrapped_toolr     sr    D>j)4I||H""4(	t	 
	.  
t	 
	. r:   c                     | rt        | dd      s| S t        j                  |       }| j                  D cg c]  }t        |       c}|_        |S c c}w )zWrap tools in config to emit execute_tool spans. Tools are sometimes passed directly as
    callable functions as a part of the config object.r   N)rI   r   r   r   )configr   rs   s      r;   wrapped_config_with_toolsr     sK     $7YYvF39<<@4L&@FLM As   Ac                 h   | rt        | dg       syg }| j                  D ]  }t        |d      rt        |j                  d      s&|j                  |j                  j                  I|j                  j                  D ]+  }t        |dd      s|j                  |j                         -  |r|S dS )z&Extract text from response candidates.r   Nrk   rr   ru   )rI   r   rH   rk   rr   r   ru   )rJ   r   r   r   s       r;   _extract_response_textr     s    
 78\2>E(( 	(	y),GI<M<Mw4W$	(9(9(?(?(G%%++ 	(DtVT*TYY'	(	( 5#t#r:   c                     | rt        | dg       syg }| j                  D ]N  }t        |dd      st        |j                        }d|v r|j	                  d      d   }|j                  |       P |r|S dS )z0Extract finish reasons from response candidates.r   Nfinish_reason.)rI   r   rd   r   splitr   )rJ   finish_reasonsr   reasons       r;   extract_finish_reasonsr     s     78\2>N(( *	9ot4001Ff}c*2.!!&)* ,>55r:   system_instructions%Union[ContentUnionDict, ContentUnion]can_be_contentlist[TextPart]c                 N   g }t        | t              rd| dgS t        | t              r| j                  rd| j                  dgS |rvt        | t              rft        | j
                  t              rJ| j
                  D ];  }t        |j                  t              s|j                  d|j                  d       = |S t        | t              r| j                  d      r	d| d   dgS |rt        | t              r| j                  dg       }|D ]  }t        |t              r9t        |j                  t              r|j                  d|j                  d       Lt        |t              s]t        |j                  d      t              s}|j                  d|d   d        |S |S )Nru   )rU   rk   rr   )
rE   rd   r   ru   r   rr   r   r   rF   rG   )r   r   
text_partsr   rr   s        r;   '_transform_system_instruction_one_levelr     sz    $&J%s+,?@AA%t,1D1I1I,?,D,DEFF*%8'B)//6+11 Ndii-%%v$))&LMN %t,1D1H1H1P,?,GHII	J':DA#''4 	MD$%*TYY*D!!6dii"HID$'Jtxx7G,M!!6d6l"KL		M
 r:   c                     g }t        | t              r't        t        j                  d | D                    }|S t	        | d      S )Nc              3   6   K   | ]  }t        |d         yw)Fr   N)r  )ro   instructionss     r;   rq   z1_transform_system_instructions.<locals>.<genexpr>I  s)        ! 8   s   Tr  )rE   r   r!   from_iterabler  )r   r  s     r;   _transform_system_instructionsr	  B  sU     $&J%t,   %8	  

 2D r:   r   integrationc                    | j                  t        j                  t               | j                  t        j                  |       |j	                  dd      r | j                  t        j
                  d       |j	                  d      }t               r|j                  rg }d}|rt        |d      r|j                  }n%t        |t              rd|v r|j	                  d      }|<| j                  t        j                  t        j                  t        |                   t!        |      }|j#                  |       |rKt%        |      }	t'        j(                         }
t+        |	| |
      }|t-        | t        j.                  |d       dt        j0                  fd	t        j2                  fd
t        j4                  fdt        j6                  fdt        j8                  fdt        j:                  fdt        j<                  ffD ]3  \  }}t        ||      st?        ||      }|"| j                  ||       5 |Ht        |d      r;|j@                  }|r,tC        |      }|rt-        | t        jD                  |d       yyyyy)zSet span data for the request.streamFTr   Nsystem_instructionunpacktemperaturetop_ptop_kmax_output_tokenspresence_penaltyfrequency_penaltyseedr   )#r   r   r   GEN_AI_REQUEST_MODELrG   GEN_AI_RESPONSE_STREAMINGr   include_promptsrH   r  rE   rF   GEN_AI_SYSTEM_INSTRUCTIONSjsondumpsr	  r   r   r   rZ   get_current_scoper   r   GEN_AI_REQUEST_MESSAGESGEN_AI_REQUEST_TEMPERATUREGEN_AI_REQUEST_TOP_PGEN_AI_REQUEST_TOP_KGEN_AI_REQUEST_MAX_TOKENSGEN_AI_REQUEST_PRESENCE_PENALTY GEN_AI_REQUEST_FREQUENCY_PENALTYGEN_AI_REQUEST_SEEDrI   r   r   GEN_AI_REQUEST_AVAILABLE_TOOLS)r   r
  ra   rf   r   r   r   r   contents_messagesnormalized_messagesscopemessages_dataparamspan_keyvaluer   r   s                    r;   set_span_data_for_requestr.  X  sF    	MM(((-8MM(//7zz(E"h88$?06

80DF  [%@%@ #gf&:;"(";";%*>&*H"(**-A"B*MM33

9:MNO 6h?)*"9("C002E:#T5M (#44! 	 
;;<	(//0	(//0	h@@A	XEEF	hGGH	--. /x 65!FE*E h./ gfg64U;O#;;# 	   7r:   c                    |syt               rB|j                  r6t        |      }|r)| j                  t        j
                  t        |             t        |      }|r)| j                  t        j                  t        |             t        |      }|rt        | t        j                  |       t        |dd      r*| j                  t        j                  |j                         t        |dd      r*| j                  t        j                  |j                          t#        |      }|d   r#| j                  t        j$                  |d          |d   r#| j                  t        j&                  |d          |d   r#| j                  t        j(                  |d          |d   r#| j                  t        j*                  |d          |d   r$| j                  t        j,                  |d          yy)	zSet span data for the response.Nresponse_idmodel_versionr.   r/   r0   r1   r2   )r   r  r   r   r   GEN_AI_RESPONSE_TEXTr   r   GEN_AI_RESPONSE_TOOL_CALLSr   r   GEN_AI_RESPONSE_FINISH_REASONSrI   GEN_AI_RESPONSE_IDr0  GEN_AI_RESPONSE_MODELr1  rR   GEN_AI_USAGE_INPUT_TOKENS GEN_AI_USAGE_INPUT_TOKENS_CACHEDGEN_AI_USAGE_OUTPUT_TOKENS$GEN_AI_USAGE_OUTPUT_TOKENS_REASONINGGEN_AI_USAGE_TOTAL_TOKENS)r   r
  rJ   response_textsr   r   rK   s          r;   set_span_data_for_responser=    s     [%@%@/9MM(779WX#H-Jh99>*;UV+H5N(99>	
 x-h1183G3GHx$/h44h6L6LM#H-J.!h88*^:TU'(55,-	

 /"h99:o;VW+,9901	

 .!h88*^:TU "r:   c                     | r| d   n|j                  dd      }t        |       dkD  r| d   n|j                  d      }t        |      }|j                  d      }t        |      }||ur||d<   |||fS )zBExtract and prepare common arguments for generate_content methods.r   ra   r   r   rf   r   )rG   r   re   r   )r   r   ra   rf   
model_namer   wrapped_configs          r;   prepare_generate_content_argsrA    s{     DGGY!?Ed)a-tAwVZZ
-CH&JZZ!F.v6NV#)x(J&&r:   c                 f    |j                  dd      }|j                  d      }t        |      }||fS )zxExtract and prepare common arguments for embed_content methods.

    Returns:
        tuple: (model_name, contents)
    ra   r   rf   )rG   re   )r   r   ra   rf   r?  s        r;   prepare_embed_content_argsrC    s7     JJw	*Ezz*%H&Jxr:   c                 4   t               r|j                  r|r~g }t        |t              r|g}nGt        |t              r'|D ]!  }t        |      }|s|j                  |       # nt        |      }|r|g}|rt        | t        j                  |d       yyyyy)z$Set span data for embedding request.Fr  N)
r   r  rE   rd   r   r   r   r   r   GEN_AI_EMBEDDINGS_INPUT)r   r
  rf   r   input_textsrp   ru   s          r;   set_span_data_for_embed_requestrG    s    
  [%@%@K (C('jHd+$ 1D06D#**401
 -X6#'&K#44 	 #  &A r:   c                 4   |syt        |d      r|j                  r|d}|j                  D ]D  }t        |d      s|j                  st        |j                  dd      }|7|t	        |      z  }F |dkD  r!| j                  t        j                  |       yyyy)z%Set span data for embedding response.N
embeddingsr   
statisticstoken_count)rH   rI  rJ  rI   r7   r   r   r7  )r   r
  rJ   r2   	embeddingrK  s         r;    set_span_data_for_embed_responserM    s      x&8+>+>!,, 	5Iy,/I4H4H%i&:&:M4P* C$44L		5 !MM(<<lK  ,?&r:   )rJ   z.Union[GenerateContentResponse, dict[str, Any]]r<   r-   )r^   r   r<   N)rf   r%   r<   zList[Dict[str, Any]])r   r   r<   Optional[dict[str, Any]])r   r   r<   rN  )rp   r   r<   rN  )rf   r%   r<   r   )r   z#Iterable[Tool | Callable[..., Any]]r<   Optional[List[dict[str, Any]]])rJ   r$   r<   rO  )r   tuple[Any, ...]r   r   rs   r'   r<   r   )rs   Tool | Callable[..., Any]r<   rQ  )r   r   r<   r   )rJ   r$   r<   zOptional[List[str]])r   r   r<   r   )r   r"   r
  r   rJ   r$   r<   N)r   rP  r   r   r<   ztuple[Any, Any, str])r   rP  r   r   r<   ztuple[str, Any])
r   r"   r
  r   rf   r   r   r   r<   N)r   r"   r
  r   rJ   r)   r<   N)Zr   r  r   	functoolsr   constsr   r   r   sentry_sdk._typesr   typingr	   r
   r   r   r   r   r   r   r   rZ   sentry_sdk.ai.utilsr   r   r   r   r   sentry_sdk.constsr   r   sentry_sdk.scoper   sentry_sdk.utilsr   r   r   google.genai.typesr   r   r   r    	itertoolsr!   sentry_sdk.tracingr"   r#   r$   r%   r&   r'   r(   r)   r*   r   PILr+   r   ImportError__optional_keys__r   r-   rR   r`   rd   re   r   r   r   r   r   rn   r   r   r   r   r   r   r   r   r   r   r  r	  r.  r=  rA  rC  rG  rM  r9   r:   r;   <module>r`     s       > > 2
 
 
   + 4 
 N M '*    	% ,,	 H>HHV/- # RjL^2j$ $ (#L..'80'8%'8T2.'2.%2.j
%5=C*
 
 
F 
>B#$'$$.6'66(@ D@,M
MM M !	M
 M 
M`3V
3V$3V0I3V	3Vl'
'%5''  
 %5  
$05?O	@L
L$L0FL	LM  		s   F
 
FF