a bge+@sdZddlZddlmZddlmZmZmZmZddl m Z m Z ddl m Z ddlmZmZmZddlmZdd lmZmZmZdd lmZmZeeZGd d d e ZGd dde ZGddde Z dS)z.Tools for interacting with a Power BI dataset.N) perf_counter)AnyDictOptionalTuple)AsyncCallbackManagerForToolRunCallbackManagerForToolRun)BaseTool) ConfigDictFieldmodel_validator)_import_tiktoken)BAD_REQUEST_RESPONSEDEFAULT_FEWSHOT_EXAMPLESRETRY_RESPONSE)PowerBIDataset json_to_mdc@sLeZdZUdZdZeed<dZeed<dZe ed<e dd Z e ed <e Zeeed <e edd Zeee fed <dZeed<dZeed<dZeeed<eddZeddeeedddZeeedddZd'eeee edddZd(eeee eddd Z eee fe!eeee fd!d"d#Z"ee!e#efd$d%d&Z$dS))QueryPowerBIToolz%Tool for querying a Power BI Dataset.Z query_powerbinamea Input to this tool is a detailed question about the dataset, output is a result from the dataset. It will try to answer the question using the dataset, and if it cannot, it will ask for clarification. Example Input: "How many rows are in table1?" descriptionN llm_chainTexcludepowerbiexamples)default_factoryr session_cachemax_iterationsioutput_token_limittiktoken_model_nameZarbitrary_types_allowedbefore)mode)valuesreturncCs0|d}|jjD]}|dvrtd|jjq|S)z8Make sure the LLM chain has the correct input variables.r) tool_inputtablesschemasrzrLLM chain for QueryPowerBITool must have input variables ['tool_input', 'tables', 'schemas', 'examples'], found %s)promptZinput_variables ValueError)clsr$rvarr-t/var/www/html/cobodadashboardai.evdpl.com/venv/lib/python3.9/site-packages/langchain_community/tools/powerbi/tool.py"validate_llm_chain_input_variables.s z3QueryPowerBITool.validate_llm_chain_input_variables)r&r%cCs||jvrdS|j|S)zCheck if the input is present in the cache. If the value is a bad request, overwrite with the escalated version, if not present return None.N)r)selfr&r-r-r. _check_cache=s zQueryPowerBITool._check_cache)r& run_managerkwargsr%c Ks||}r td|||Sz@td||jj||j|j|j |rV| ndd}Wn>t y}z&d||j |<|j |WYd}~Sd}~00|dkr||j |<|j |Std|t }|jj|d}t } td |td | |d ||\} } | dur:d | vr:d |j |<|j |S|dd} | rx| |jkrx|jtj||| d|| ddS| r| n tj| d|j |<|j |S):Execute the query, return the results or an error message.Found cached result for %s: %s%Running PBI Query Tool with input: %sNr&r'r(r callbacksError on call to LLM: I cannot answer thisz PBI Query: %scommandPBI Result: %sPBI Query duration: 0.6f TokenExpiredzCAuthentication token expired or invalid, please try reauthenticate. iterationsrr&queryerrorr&r2rArD)r1loggerdebuginforZpredictrget_table_names get_schemasr get_child Exceptionrrrun _parse_outputgetr_runrformatr r0r&r2r3cacherCexc start_time pbi_resultend_timeresultrDrAr-r-r.rRFsR          zQueryPowerBITool._runc s||}r&td|||dSzFtd||jj||j|j|j |r\| nddIdH}Wn>t y}z&d||j |<|j |WYd}~Sd}~00|dkr||j |<|j |Std|t }|jj|d IdH}t } td |td | |d ||\} } | durVd | vsBd| vrVd|j |<|j |S|dd} | r| |jkr|jtj||| d|| ddIdHS| r| n tj| d|j |<|j |S)r4r5z3, from cache, you have already asked this question.r6Nr7r9r:z PBI Query: %sr;r=r>r?r@ TokenErrorzkAuthentication token expired or invalid, please try to reauthenticate or check the scope of the credential.rArrBrErFrG)r1rHrIrJrZapredictrrKrLrrMrNrrZarunrPrQr_arunrrSrrTr-r-r.r\|sR          zQueryPowerBITool._arun)rXr%cCsd|vrt|ddddd}t|dkr:tddSt|}||\}}|rld|d|jd d fS|d fSd |vrd |d vrd |d d vrd |d d d ddfSd |d fSd |fS)z2Parse the output of the query to a markdown table.resultsrr'rowsz%0 records in result, query was valid.)Nz^0 rows returned, this might be correct, but please validate if all filter values were correct?z[Result too large, please try to be more specific or use the `TOPN` function. The result is z tokens long, the limit is z tokens.NrDz pbi.errordetailsdetail)lenrHrJr_result_too_larger)r0rXr^rZZtoo_longlengthr-r-r.rPs(    zQueryPowerBITool._parse_output)rZr%cCsD|jr@t}||j}t||}td|||jk|fSdS)z!Tokenize the output of the query.zResult length: %s)Fr)r r Zencoding_for_modelraencoderHrJr)r0rZZ tiktoken_encodingrcr-r-r.rbs  z"QueryPowerBITool._result_too_large)N)N)%__name__ __module__ __qualname____doc__rstr__annotations__rrrr rrrrrdictrrrintrr r model_configr classmethodr/r1rrRrr\rrPboolrbr-r-r-r.rsJ         9 7  rc@sveZdZUdZdZeed<dZeed<eddZ e ed<e dd Z dee eed d d Zdee eed ddZd S)InfoPowerBIToolz2Tool for getting metadata about a PowerBI Dataset.Zschema_powerbirz Input to this tool is a comma-separated list of tables, output is the schema and sample rows for those tables. Be sure that the tables actually exist by calling list_tables_powerbi first! Example Input: "table1, table2, table3" rTrrr!Nr&r2r%cCs|j|dS)z4Get the schema for tables in a comma-separated list., )rZget_table_infosplitr0r&r2r-r-r.rRszInfoPowerBITool._runcs|j|dIdHS)Nrs)rZaget_table_infortrur-r-r.r\szInfoPowerBITool._arun)N)Nrfrgrhrirrjrkrr rrr rnrrrRrr\r-r-r-r.rqs&    rqc@s~eZdZUdZdZeed<dZeed<eddZ e ed<e dd Z de ee eed d d Zde ee eed ddZd S)ListPowerBIToolzTool for getting tables names.Zlist_tables_powerbirzUInput is an empty string, output is a comma separated list of tables in the database.rTrrr!NrrcCsd|jSzGet the names of the tables.rsjoinrrKrur-r-r.rRszListPowerBITool._runcsd|jSrxryrur-r-r.r\szListPowerBITool._arun)NN)NNrvr-r-r-r.rws*    rw)!riloggingtimertypingrrrrZlangchain_core.callbacksrrZlangchain_core.toolsr Zpydanticr r r Z&langchain_community.chat_models.openair Z(langchain_community.tools.powerbi.promptrrrZ%langchain_community.utilities.powerbirr getLoggerrfrHrrqrwr-r-r-r.s    C