TCjlqBUdZddlZddlZddlZddlZddlZddlZddlZddlZddlm Z m Z m Z m Z m Z mZmZmZmZddlZddlmZddlmZddlmZddlmZdd lmZdd lmZdd lmZejd krdd lmZnej Zej!r ddl"m#Z$ddl%m&Z'ne Z$ej e(d<e Z'ej e(d<eddZ)dZ*ej+dZ,dej-dej.fdZ/dee0ej1fde e0fdZ2dee0ej1fde e0fdZ3 dJdee0ej4fde ej5dej5fdZ6 dKde ej-de7de8e9fd Z:de ej-d!e8e9ddfd"Z; dLde ej-d#e edej>fd(Z?d'eed0e d%e fdej>fd1ZDd'ej>d2e d%e fde fd3ZEd'ej>d2e d%e fde fd4ZFd5ejGd6eZQd?eejRejSfd@ejGdeejRejSffdAZT dMdBee0ejUe0ejVfdCe ejWdDe e0dePejXe9e0ffdEZY dJde ej-de7fdFZZde ee!  D8#u+,,33 [ 23;q>>d2 r'ct|tr|dr|St|tr|dr|dSt|t jr|jr|jSdS)z8Extracts the BigQuery URI from the source, if available.zbq:// bigquery_uriN)r"r+r,r#r-rr.r2r/s r%_get_bigquery_urir3QsScnnW55 J#t!8!8 ~ #u+,,1A   r'c`|tj}d}|jstj}d||_|jskt |}t|}|r(|dr|ddd|_n#|r!|ptj}|d||_|S)zBFormats the destination uri based on the source uri for Vertex AI.Ngenai_batch_job_z.jsonliz/dest_dest_) rCreateBatchJobConfig display_namer timestamped_unique_namedestr0r3endswith)r(r unique_namegcs_source_uribigquery_source_uris r%format_destinationr?^s  ^  ' ) )F+  ;133K:[::F @!#&&N+C00@.11(;;@%crc*111fkk @D7#B#D#Dk*??+??fk -r'Fis_agent_platformcL|sgSt|}g}|r|js|St|jD]n\}}t|tjs t |ddr||t |ddr|s||o|S)z9Checks if the config contains any AFC incompatible tools.function_declarationsN mcp_servers)r&tools enumerater"rrgetattrappend)rr@ config_modelincompatible_tools_indexesindextools r%"find_afc_incompatible_tool_indexesrLxs  I6v>>,*, &\/& %%|122//keT dEJ ' 't,d33/ ''...t]D))/2C/ ''... ##r'rIc|sdSd}t|tjr|jrt |j}n=r'mcp_to_genai_tool_adaptersis_caller_method_async.ci}|s|St|}|jr^|jD]V}t|rEtj|r'|s%t jd|jd|jd|||j<W|rm|st jd|D]-\}}| |rtd|d.| ||S)z'Returns a function map from the config.z Function zh is a coroutine function, which is not supported for automatic function calling. Please manually invoke z to get the function response.z3MCP tools are not supported in synchronous methods.Tool $ is already defined for the request.) r&rDcallableinspectiscoroutinefunctionrUnsupportedFunctionError__name__itemsr- ValueErrorupdate)rrVrW function_maprHrK tool_name_s r%get_function_maprfs^OQ,  6v>>, +"++ $+  &t , , 5K /G$-GGGGG   '+ T]# 4 !  + ?  388::  1   ) $ $  CI C C C    2333 r'argsc>d|DS)z:Converts float values in dict with no decimal to integers.c4i|]\}}|t|Sr!,convert_number_values_for_function_call_args.0keyvalues r% zEconvert_number_values_for_dict_function_call_args..s7    #u 7 > >   r')r`rgs r%1convert_number_values_for_dict_function_call_argsrrs+      r'ct|tr#|rt|St|trd|DSt|t r d|DS|S)z2Converts float values with no decimal to integers.c4i|]\}}|t|Sr!rjrls r%rpz@convert_number_values_for_function_call_args..s7    C 9% @ @   r'c,g|]}t|Sr!rj)rmros r% z@convert_number_values_for_function_call_args..s.   @E4U;;   r')r"float is_integerintr#r`listrqs r%rkrkse!2!2 t99d  **,,   d  IM    +r' annotationc tj|ot|tjS#t $rYdSwxYwNF)r\isclass issubclasspydantic BaseModel TypeError)r{s r%is_annotation_pydantic_modelrsV ?: & & :H&,,  55s -0 >>ro param_name func_namec t|trRt|rC |d i|S#tj$r(}t jddd|d|d| d}~wwxYwt|tr=t|tkr%t|dfd|DSt|trLt|tkr4t|\}fd| DSt|ttfvrt|D]}t|rt|tus4t||s$t|tr9t|r* t||cS#tj$rYwxYwt jddd|d |d t|tr |tur|St||s!t jddd|d|d |S) NzFailed to parse parameter z for function z> from function call part because function call argument value z- is not compatible with parameter annotation z, due to error rc4g|]}t|Sr!convert_if_exist_pydantic_model)rmitemr item_typers r%rvz3convert_if_exist_pydantic_model..s7     (iYOO   r'c <i|]\}}|t|Sr!r)rmkvrr value_types r%rpz3convert_if_exist_pydantic_model.. s>    Aq *1j*i P P   r'z- cannot be converted to parameter annotation .r!)r"r#rrValidationErrorr UnknownFunctionCallArgumentErrorrzr r r`r rrryrw) ror{rrereargrrs `` @@r%rrsgt !=j!I!I  Z %   #  3 -z - - - - - - - -*+ - -  tJ!7!74!?!?$$Q'I         tJ!7!74!?!?Z((MAz      KKMM     y111 ##   C== '__44 s # #5%%5*Fs*K*K5 0S*i     '    (   1 Z             s e 3 3 L E: & &  1 CZ C Cy C CDI C C5? C C C   ,s'2A)#A$$A)'F;;G  G functionctj|}|j}i}|jD])\}}||vr t |||j||||<*|SN)r\ signaturer_ parametersr`rr{)rgrrrconverted_argsrparams r%convert_argument_from_functionr6s)))).$/5577j%T#B z      $$nZ r'function_to_invokec t||} |di|S#t$r*}tjd|jd|d|d|d}~wwxYwNzFailed to invoke function z with converted arguments z, from model returned function call argument z because of error r!r ExceptionrFunctionInvocationErrorr_rgrrrs r%invoke_function_from_dict_argsrGs2$8JKK.   / / / //   ( 6%7%@ 6 6 . 6 6 6 623 6 6  s A%A  Ac Kt||} |di|d{VS#t$r*}tjd|jd|d|d|d}~wwxYwrrrs r%$invoke_function_from_dict_args_asyncrUs2$8JKK.##55n55 5 5 5 5 5 55   ( 6%7%@ 6 6 . 6 6 6 623 6 6  s " A%AAresponserccdg}|j%t|jdjtjr|jdjj|jdjjD]}|js |jj}||jj||}t|jj} t|tsdt||i}n(#t$r}dt|i}Yd}~nd}~wwxYwtj||} || |S)6Returns the function response parts from the response.Nrresulterrornamer) candidatesr"contentrContentparts function_callrrgrrrrrr+Partfrom_function_responserG) rrcfunc_response_partspartrfuncrg func_responserfunc_response_parts r%get_function_response_partsrcsZ % X(+3U] C C&  a ( . :#A&.477  $)i  4#5#:#FI&@   #   ,D"788 8tDDM , , ,"CFF+------ ,"Z>>]?   ""#5666 s)'C C6C11C6c Kg}|jt|jdjtjr|jdjj|jdjjD]k}|js |jj}|Q||}|jj |jjni}t|} t|tr[| tj ||d{V}t|dt|dd} | rd|i} nUd|i} nPtj|rdt!||d{Vi} n#dt#jt&||d{Vi} n(#t($r} dt+| i} Yd} ~ nd} ~ wwxYwtj|| } || m|S) rNr)rrgis_errorisErrorFrrr)rr"rrrrrrrgrrr call_tool FunctionCallrFr\r]rasyncio to_threadrrr+rrrG) rrcrrrrraw_argsrgmcp_tool_responserrrrs r%!get_function_response_parts_asyncrsf % X(+3U] C C&  a ( . :#A&.4*7*7  $)i  I&!&2   # #  AJJ , 3 4 4 &*nn" ===''!!!!!! !)9e<<H  <&(9:mm'):;mm*400  DT4 P PPPPPPPMM  12D$!!M  , , ,"CFF+------ ,"Z>>]?   ""#5666 s3CE44 F>FFc|sdSt|}|r[|jrT|jj Gt|jjdkr*td|jjddS|jr |jjdS|jjr`|jj Sd|jjvrEt|jjdkr(td|jjd |jjS) z6Returns whether automatic function calling is enabled.FNrz6max_remote_calls in automatic_function_calling_config zu is less than or equal to 0. Disabling automatic function calling. Please set max_remote_calls to a positive integer.Tmaximum_remote_callsz`automatic_function_calling.disable` is set to `True`. And `automatic_function_calling.maximum_remote_calls` is a positive number a9. Disabling automatic function calling. If you want to enable automatic function calling, please set `automatic_function_calling.disable` to `False` or leave it unset, and set `automatic_function_calling.maximum_remote_calls` to a positive integer or leave `automatic_function_calling.maximum_remote_calls` unset.)r&automatic_function_callingrryrQrRdisablemodel_fields_setrrHs r%should_disable_afcrsR  56v>>,   1   1 F   l5J K Kq P P NN >  3 H > > > 4  1  0 8 @ 5-5  1 F   !  0 A B B l5J K Ka O O NN D  3 H D D D     0 88r'c|stS t|rtdt|}|jr |jjtSt |jjS)NzUautomatic function calling is not enabled, but SDK is trying to get max remote calls.)_DEFAULT_MAX_REMOTE_CALLS_AFCrrar&rrryrs r%get_max_remote_calls_afcrs| ) ((J    7v>>,  1)  0 E M (( \ 4 I J JJr'c|r|jr |jjsdS|j}|r|jnd}|jjj}|r|st ddSdS)z9Raises an error if the config is not compatible with AFC.NFavRunning in streaming mode with stream_function_call_arguments enabled, this feature is not compatible with automatic function calling (AFC). Please set config.automatic_function_calling.disable to True to disable AFC or leave config.tool_config. function_calling_config.stream_function_call_arguments to be empty or set to False to disable streaming function call arguments.) tool_configfunction_calling_configrrstream_function_call_argumentsra)r afc_configdisable_afc_configstream_function_calls r%'raise_error_for_afc_incompatible_configr s      7  F0*-7Bz))U 0O"4  I  r'cT|sdSt|}|jsdS|jj S)NT)r&rignore_call_historyrs r%should_append_afc_historyr#s=  46v>>,  0 4  4 H HHr'cb|sdSt|}|ddi}|j|_|jrptj|jrW|jt ji|_|jj i|j_tj |jj|S)zVReturns a parsed config with an appended MCP header if MCP tools or sessions are used.NrDrb)headers) r& model_copyrDrhas_mcp_tool_usage http_optionsr HttpOptionsrset_mcp_usage_header)rrHconfig_model_copys r%parse_config_for_mcp_usager.s  46v>>,#--gt_-EE(.LJ9,:LMML%-','8'D'D'D$%-5/1$,#$5$B$JKKK r'cKi}t|}|sd|fS|ddi}|jr@g|_tjs!|j|jn ddlm}n!#t$rtddi}YnwxYw|jD]}t||rt|| d{V|}|j|j|jD]V}|j rM|j D]E} | j<|| jrt!d | jd ||| j<FW|j|||fS) zReturns a parsed config with MCP sessions converted to GenAI tools. Also returns a map of MCP tools to GenAI tool adapters to be used for AFC. NrDrrr DummySessionr!)r@rYrZ)rrrDr_is_mcp_loadedextendmcpr ImportErrortyper"r list_toolsrBrr-rarG) rr@rV parsed_configparsed_config_copy_McpClientSessionrKmcp_to_genai_tool_adapter genai_toolfunction_declarations r%parse_config_for_mcp_sessionsrCsBD,V44- , + ++%///GG 0!  $ & &0%%m&9::::9::::::: 999 R889 %00$ d- . . 0&;$//++++++++?P' ' ' #  " ) )*C*I J J J5;  j/ *4*J  &',8/334H4MNN$, 4 9,,, 0--A-FG   " ) )$ / / / / 7 77s6A==BBcontentschunkc|Y|jR|jdj}tj|}t |t r||||S)zFAppends the contents of the chunk to the contents list and returns it.Nr)rrt t_contentsr"rzrG)rr chunk_contents r%append_chunk_contentsrzs_  5+7$Q'/M|H%%H(D!!%m&?oom$$$ /r'fileuser_http_optionsuser_mime_typec d}|}t|tjr|tdt |drd|jvrtd|}|dtj ||z }||tj ntj |}|rtj |st|dtj |}|t!j|\}}|td|r?t|t$rt'jdi|}|}d |_d d d ||d |_nt'jd d d d ||d }t|t.tjfr5|ji|_tj ||jd<|||fS)a-Prepares the HTTP options, file bytes size and mime type for a resumable upload. This function inspects a file (from a path or an in-memory object) to determine its size and MIME type. It then constructs the necessary HTTP headers and options required to initiate a resumable upload session. NzfUnknown mime type: Could not determine the mimetype for your file please set the `mime_type` argumentmodebz'The file must be opened in binary mode.rz is not a valid file path.ziUnknown mime type: Could not determine the mimetype for your file please set the `mime_type` argumentzapplication/json resumablestart)z Content-TypezX-Goog-Upload-ProtocolzX-Goog-Upload-Commandz#X-Goog-Upload-Header-Content-Lengthz!X-Goog-Upload-Header-Content-Type) api_versionrzX-Goog-Upload-File-Namer!)r"ioIOBaserahasattrrtellseekosSEEK_ENDSEEK_SETfspathpathisfileFileNotFoundErrorgetsize mimetypes guess_typer#rrrrr+PathLikebasename) rrr size_bytes mime_typeoffsetfs_pathrers r%prepare_resumable_uploadrsD*)bi    8  tVD DI  BCCC YY[[FIIav%JIIfbk""""iooG C"'..11C AAA B BB))J)'22li  ;   #T**A+@@.?@@$L!L*"-!(2<09^ L$.&1%,6@?4=      LsBK())M#! 68g6F6Ft6L6LL23 z9 ,,r'cz|sdSt|}|jsdS|jD]}t|ddrdSdS)zBChecks whether the configuration contains any MCP server requests.FrCNT)r&rDrF)rrHrKs r%has_agent_platform_mcp_serversrse  U8@@L   U" }d + +tt 5r' config_clsusagecndtjd|}|s |}n t|tr |di|}n|}t |ds|St |dd}|$t j}t|d|tj t j|}|j pi}dD]z}||vrod|||vra|||vrWdtj||vr-|| dtj|||<a||xxd|z cc<u|||<{||_ |S)z(Returns the usage header for the config.zgoogle-genai-sdk/+rN)z user-agentzx-goog-api-client r!) public_version __version__r"r#rrFrrsetattrtypingcastrreplace)rrr usage_headerrHrexisting_headers header_keys r%get_usage_headerr*sJ^%?II%II, :<>> = : < < + , ,*:*)E)M)M>."<>> * * : & & : & & &*r?s UTT  \\\\\\\\\\\\\\\\\\\\\\ '''''',,,,,,w') 333333'''''''"%FK%%%'6; GC*+++ "  0 1 1  -    sE.. /  c]      sE.. /  c]    48 sE(( ) U/ 0 6;?#$$ U6 7$$ #Y$$$$2 U6 7 $S  8;? #( !! U6 7! ( S' '(!! ! !  #uXc3h')>>? ?@ !!!!H    S&[!4<7 8 4V d6lF 23$ S T    ; ;;-0;=@;;;;;|  (0c(: "    2:382D         2:382D      + sE(38"46K"KLLM  %*    F6+6sE(38"46K"KLLM6 %*6666t;?4949 U6 749 49494949p;?KK U6 7KKKKK(HU=X4Y 4;?II U6 7I IIII;? U6 7 e)*,;?#4848 U6 74848  U ()# #$&48484848n E*E,FFG   (  5 !5#= =>    <@$(F-F- R[%ry0 1F- 78F-SMF-    F-F-F-F-T;?   U6 7       + U4S>1,- .+<@G+LO+++++++r'