; idZddlmZddlZddlmZmZmZddlm Z ddl m Z ddl m Z ddl m Z Gd d eZGd d ejZGd dee jZGddee jZd dZd!dZd"dZd#dZd$dZdS)%z2Utilities for working with post-processing tokens.) annotationsN)AnyCallableSequence)py_utils) llm_function)llmfn_output_row)llmfn_post_processceZdZdZdS)PostProcessParseErrorz,An error parsing the post-processing tokens.N)__name__ __module__ __qualname____doc__/var/www/development/aibuddy-work/election-extract/venv/lib/python3.11/site-packages/google/generativeai/notebook/post_process_utils.pyr r s6666rr cVeZdZdZejd dZejd dZdS) ParsedPostProcessExprz:A post-processing expression parsed from the command line.returnstrcdS)z$Returns the name of this expression.Nrselfs rnamezParsedPostProcessExpr.name"rllm_fnllm_function.LLMFunctioncdS)zEAdds this parsed expression to `llm_fn` as a post-processing command.Nrrrs radd_to_llm_functionz)ParsedPostProcessExpr.add_to_llm_function&rrNrrrrrr)r rrrabcabstractmethodrr!rrrrrsfDD3333 TTTTTTrrc2eZdZdZddZddZdd ZddZdS)_ParsedPostProcessAddExprzEAn expression that returns the value of a new column to add to a row.rrfnCallable[[str], Any]c"||_||_dS)aConstructor. Args: name: The name of the expression. The name of the new column will be derived from this. fn: A function that takes the result of a row and returns a new value to add as a new column in the row. N_name_fnrrr(s r__init__z"_ParsedPostProcessAddExpr.__init__0s rrc|jSNr,rs rrz_ParsedPostProcessAddExpr.name< zrrows-Sequence[llmfn_output_row.LLMFnOutputRowView] Sequence[Any]c fd|DS)Nc^g|])}|*Srr- result_value.0rowrs r z6_ParsedPostProcessAddExpr.__call__..@1===))++,,===rrrr4s` r__call__z"_ParsedPostProcessAddExpr.__call__?========rrrc:||j|SNrr()add_post_process_add_fnr,r s rr!z-_ParsedPostProcessAddExpr.add_to_llm_functionBs--4:$-GGGrN)rrr(r)r")r4r5rr6r#r rrrr/rrAr!rrrr'r'+ssPO    >>>>HHHHHHrr'c2eZdZdZddZddZdd ZddZdS)_ParsedPostProcessReplaceExprz:An expression that returns the new result value for a row.rrr(Callable[[str], str]c"||_||_dS)zConstructor. Args: name: The name of the expression. fn: A function that takes the result of a row and returns the new result. Nr+r.s rr/z&_ParsedPostProcessReplaceExpr.__init__Ks rrc|jSr1r2rs rrz"_ParsedPostProcessReplaceExpr.nameUr3rr4r5 Sequence[str]c fd|DS)Nc^g|])}|*Srr9r;s rr>z:_ParsedPostProcessReplaceExpr.__call__..Yr?rrr@s` rrAz&_ParsedPostProcessReplaceExpr.__call__XrBrrrc:||j|SrD)add_post_process_replace_fnr,r s rr!z1_ParsedPostProcessReplaceExpr.add_to_llm_function[s11tzd1KKKrN)rrr(rJr")r4r5rrMr#rGrrrrIrIFssED>>>>LLLLLLrrIr(r)c.t|j|SrD)r'r r(s rpost_process_add_fnrT`s $"+" = = ==rrJc.t|j|SrD)rIr rSs rpost_process_replace_fnrVds (bkb A A AArtokensrMrNonecl|stdt|dkrtddS)Nz,Cannot have empty post-processing expressionz3Post-processing expression should be a single token)r len)rWs r'validate_one_post_processing_expressionr\hsA T#$RSSS 6{{Q#$YZZZrtuple[str, Any]c t||dd}tj}t |D]`\}}t |}||vr@t dd|d|dz||}ad||fS)z)Returns name and the resolved expression.r.zUnable to resolve "{}"NrZ ) r\splitrget_main_module enumeratevarsr formatjoin)rW token_partscurrent_modulepart_numpartcurrent_module_varss r'_resolve_one_post_processing_expressionrlqs,F333)//#&&K-//N#K0033$">22 * * *'(//^xRS|^9T0U0UVV -T2 HHV  n --rSequence[Sequence[str]]Sequence[ParsedPostProcessExpr]cDg}|D]}t|\}}t|tr||?t|tr%|t ||yt d||S)a Resolves post-processing tokens into ParsedPostProcessExprs. E.g. Given [["add_length"], ["to_upper"]] as input, this function will return a sequence of ParsedPostProcessExprs that will execute add_length() and to_upper() on each entry of the LLM output as post-processing operations. Raises: PostProcessParseError: An error parsing or resolving the tokens. Args: tokens: A sequence of post-processing tokens after splitting. Returns: A sequence of ParsedPostProcessExprs. rEz{} is not callable)rl isinstancerappendrr'r re)rWresults expression expr_name expr_values rresolve_post_processing_tokensrvs$,.GPP G S S : j"7 8 8 P NN: & & & &  H - - P NN4) SSS T T T T'(<(C(CI(N(NOO O Nr)r(r))r(rJ)rWrMrrX)rWrMrr])rWrmrrn)r __future__rr$typingrrrgoogle.generativeai.notebookr google.generativeai.notebook.librr r RuntimeErrorr ABCrLLMFnPostProcessBatchAddFnr'LLMFnPostProcessBatchReplaceFnrIrTrVr\rlrvrrrrs98"""""" **********111111999999======??????77777L777 T T T T TCG T T THHHHH-HHHH6LLLLL-LLLL4>>>>BBBB[[[[....*r