tjNvdZddlZddlZddlZddlZddlmZddlmZddl m Z m Z m Z m Z mZmZmZddlmZe rddlZddlmZddlmZdd lmZmZmZmZdd lmZdd lm Z m!Z!m"Z"dd l#m$Z$m%Z%m&Z&ej'e(Z)Gd de$Z*Gdde$Z+Gdde$Z,Gdde$Z-Gdde$Z.Gdde$Z/Gdde$Z0Gdde$Z1GddZ2Gdd Z3Gd!d"eZ4Gd#d$e$Z5Gd%d&ee5Z6dS)'z Kinetica SQL generation LLM API.N)version)Path) TYPE_CHECKINGAnyDictListOptionalPatterncast)pre_init)CallbackManagerForLLMRun) BaseChatModel) AIMessage BaseMessage HumanMessage SystemMessage)BaseOutputParser)ChatGeneration ChatResult Generation) BaseModel ConfigDictFieldceZdZUdZeddZeeed<eddZ eeed<egdZ e eed <edd Z ee eed <edd Z eeed <defdZdS)_KdtSuggestContextpydantic API request typeNz Name of tabledefaulttitletablezTable description descriptionzTable columns listcolumnszRules that apply to the table.rulesz)Samples that apply to the entire context.samplesreturnc g}|d|jd|d|jrt|jdkrt dg}|jD]B}|dd}|d|C|d ||d |jr&|d |jd |jd |j rXt|j dkr@|d|jd|j D]}|d|d|}|S)Nz CREATE TABLE z AS(rzcolumns list can't be null."z z, z);zCOMMENT ON TABLE z IS 'z';z-- When querying table z the following rules apply:z-- *  ) appendr r"len ValueErrorreplacestripjoinr!r#)selflinesr"columnruleresults f/var/www/html/CA-Chatbot/venv/lib/python3.11/site-packages/langchain_community/chat_models/kinetica.py to_system_strz _KdtSuggestContext.to_system_str1s 4TZ444555 S| 6s4<00A55 4 5 5 5l + +F^^C,,2244F NN>>> * * * * UZZ(())) T   T LLRTZRRd>NRRR S S S : -#dj//A-- LLQ$*QQQ     - - ^T^^,,,,5!! )__name__ __module__ __qualname____doc__rr r str__annotations__r!r"rr#r$rr7r8r6rr$s ## 5_EEEE8C=EEE!&t;N!O!O!OK#OOOr1EFFFGT#YFFF!&<"""E8DI $eGGXd^sr8rcteZdZUdZdZeeed<ee ed<defdZ dee fdZ de fdZ dS) _KdtSuggestPayloadrNquestioncontextr%cg}|jD]3}|j |}||4d|S)Nz )rCr r7r+r0)r1r2 table_context context_strs r6get_system_strz!_KdtSuggestPayload.get_system_strTs\!\ & &M"*'5577K LL % % % %{{5!!!r8c,g}|jD]}|j |jD]e\}}|dd}|t d|pd|t d|f|S)Nz'''userr)rolecontent assistant)rCr$itemsr.r+dict)r1messagesrCrBanswers r6 get_messagesz_KdtSuggestPayload.get_messages]s| H HG&$+O$9$9$;$; H H &c22&(.b I I IJJJ+v F F FGGGG  H r8c:g}|td||||td|jpdt|}|S)NsystemrKrJr)rQ)r+rPrGextendrSrB)r1rQresponses r6 to_completionz _KdtSuggestPayload.to_completionjs(D4G4G4I4IJJJKKK))++,,,&$-2E2FFFGGG***r8)r9r:r;r<rBr r=r>rrrGrrSrYr?r8r6rArANs##"Hhsm""" $ %%%%""""" d4j    tr8rAceZdZUdZeed<dS)_KdtoSuggestRequestrpayloadN)r9r:r;r<rAr>r?r8r6r[r[ss$## r8r[cBeZdZUdZeddZeed<eed<dS) _KdtMessagepydantic API response typer)zOne of [user|assistant|system]rrLrMN)r9r:r;r<rrLr=r>r?r8r6r^r^ys?$$b(HIIID#III LLLLLr8r^cXeZdZUdZeed<eddZee ed<e ed<dS) _KdtChoicer_indexNzThe generated SQLrmessage finish_reason) r9r:r;r<intr>rrcr r^r=r?r8r6rarasR$$ JJJ%*U4?R%S%S%SGXk "SSSr8rac2eZdZUdZeed<eed<eed<dS) _KdtUsager_ prompt_tokenscompletion_tokens total_tokensN)r9r:r;r<rer>r?r8r6rgrgs<$$r8rgceZdZUdZeed<eed<eed<eed<eeed<e ed<e dd Z eed <d S) _KdtSqlResponser_idobjectcreatedmodelchoicesusager)zThe input questionrpromptN) r9r:r;r<r=r>rerrargrrsr?r8r6rlrlsy$$ GGG KKK LLL JJJ *  %*>???FC?????r8rlc(eZdZUdZeed<eed<dS)_KdtCompletionResponser_statusdataN)r9r:r;r<r=r>rlr?r8r6rurus-$$ KKK r8ruceZdZUdZejdejZee d<e de de de fdZ e de jdefd Ze de d e defd Zd S) _KineticaLlmFileContextParserz*Parser for Kinetica LLM context datafiles.z'^<\|(?P\w+)\|>\W*(?P.*)$PARSERtextsuffixr%cd|r-||r|dt| S|S)N)endswithr,)clsr{r|s r6 _removesuffixz+_KineticaLlmFileContextParser._removesuffixs9  (dmmF++ (3v;;,' ' r8 input_filect|}||jd}t|}|||S)Nz.txt)rrnameopenreadparse_dialogue)rrpathschemar2s r6parse_dialogue_filez1_KineticaLlmFileContextParser.parse_dialogue_filesUJ""49f55Z  %%''!!%000r8rcg}d}|d}d}t|D]\}}|}t|dkr.|j|} | t d|| } | d} | dkr|t d|| d}| dkr|t d || }| d krA|td || || | d}t d | |||d S)Nz<|end|>rz"Could not find starting token in: rLrUz"Only one system token allowed in: rMrJz*Found user token without assistant token: rNz*Found assistant token without user token: zUnknown token: )rrUrQ) split enumerater/r,rzmatchr- groupdict Exceptionr+) rr{rrQrUr2 user_messageidxlinerrrLs r6rz,_KineticaLlmFileContextParser.parse_dialogues 9%% "5)) ; ;IC::< classmethodr=rosPathLikerrrr?r8r6ryrys44!bj!KRYWWFGWWWcc[ 1R[1T111[1'J#'Js'Jt'J'J'J['J'J'Jr8ryc eZdZdZe d deedeedeeddfdZed ed eedefd ZdS) KineticaUtilzKinetica utility functions.NurlrJpasswdr%z gpudb.GPUdbc ddl}n#t$rtdwxYw|d|}|d|}|d|}|j}||_||_d|_d|_ d|_ |j|| }t d |td |j|S) aCreate a connectica connection object and verify connectivity. If None is passed for one or more of the parameters then an attempt will be made to retrieve the value from the related environment variable. Args: url: The Kinetica URL or ``KINETICA_URL`` if None. user: The Kinetica user or ``KINETICA_USER`` if None. passwd: The Kinetica password or ``KINETICA_PASSWD`` if None. Returns: The Kinetica connection object. rNzUCould not import Kinetica python package. Please install it with `pip install gpudb`. KINETICA_URL KINETICA_USERKINETICA_PASSWDTINFO)hostoptionsz.Connected to Kinetica: {}. (api={}, server={})gpudb)rModuleNotFoundError ImportError_get_envGPUdbOptionsusernamepasswordskip_ssl_cert_verificationdisable_failover logging_levelLOGinfoformatget_urlrserver_version)rrrJrrrkdbcs r6 create_kdbczKineticaUtil.create_kdbcs*  LLLL"   >   ll>3//||OT22/88+%%''!-1*#' &u{W555  < C C  0 0$2E      s!rrc^||Stj|}||Std|)z-Get an environment variable or use a default.Nz;Parameter was not passed and not found in the environment: )rgetenvr-)rrrr5s r6rzKineticaUtil._get_envsD  N4  M P$ P P   r8)NNN) r9r:r;r<rr r=rrr?r8r6rrs%%"" $ .. c].sm. .  ...[.`  C  (3-  C    [    r8rc eZdZUdZedZeed< ede de fdZ e de fdZ e de e effd Z dd eed eee d eededef dZde defdZd ee defdZde de fdZededeefdZedede fdZede defdZede deefdZd S) ChatKineticaaKKinetica LLM Chat Model API. Prerequisites for using this API: * The ``gpudb`` and ``typeguard`` packages installed. * A Kinetica DB instance. * Kinetica host specified in ``KINETICA_URL`` * Kinetica login specified ``KINETICA_USER``, and ``KINETICA_PASSWD``. * An LLM context that specifies the tables and samples to use for inferencing. This API is intended to interact with the Kinetica SqlAssist LLM that supports generation of SQL from natural language. In the Kinetica LLM workflow you create an LLM context in the database that provides information needed for infefencing that includes tables, annotations, rules, and samples. Invoking ``load_messages_from_context()`` will retrieve the contxt information from the database so that it can be used to create a chat prompt. The chat prompt consists of a ``SystemMessage`` and pairs of ``HumanMessage``/``AIMessage`` that contain the samples which are question/SQL pairs. You can append pairs samples to this list but it is not intended to facilitate a typical natural language conversation. When you create a chain from the chat prompt and execute it, the Kinetica LLM will generate SQL from the input. Optionally you can use ``KineticaSqlOutputParser`` to execute the SQL and return the result as a dataframe. The following example creates an LLM using the environment variables for the Kinetica connection. This will fail if the API is unable to connect to the database. Example: .. code-block:: python from langchain_community.chat_models.kinetica import KineticaChatLLM kinetica_llm = KineticaChatLLM() If you prefer to pass connection information directly then you can create a connection using ``KineticaUtil.create_kdbc()``. Example: .. code-block:: python from langchain_community.chat_models.kinetica import ( KineticaChatLLM, KineticaUtil) kdbc = KineticaUtil._create_kdbc(url=url, user=user, passwd=passwd) kinetica_llm = KineticaChatLLM(kdbc=kdbc) Texcludervaluesr%cr|dd}|t}||d<|S)zPydantic object validator.rN)getrr)rrrs r6validate_environmentz!ChatKinetica.validate_environmentZs;zz&$'' <++--D!F6N r8cdS)Nzkinetica-sqlassistr?r1s r6 _llm_typezChatKinetica._llm_typeds##r8cltt|jjt dS)Nr)kinetica_version api_version)rPr=rrrrs r6_identifying_paramsz ChatKinetica._identifying_paramshs4 !9::PWHXHX    r8NrQstop run_managerkwargsc |tdfd|D}|}tt|jdj}|}|} t|j j |j j |j } tt| g| S)Nzstop kwargs are not permitted.c:g|]}|Sr?)_convert_message_to_dict.0mr1s r6 z*ChatKinetica._generate..xs'LLLa66q99LLLr8r) input_tokens output_tokens model_name)rc) generations llm_output)r-_submit_completionr r^rqrc model_dump_convert_message_from_dictrPrrrhrirprr) r1rQrrr dict_messages sql_responseresponse_messagegenerated_dictgenerated_messagers ` r6 _generatezChatKinetica._generatens  =>> >LLLL8LLL ..}==  \-A!-D-LMM)4466 ;;NKK%+9&,>#)   '0ABBBC!    r8 context_namecd|d}|}|d}tj|}t|}|j}g}|td|| | fd|D} | S)aRLoad a lanchain prompt from a Kinetica context. A Kinetica Context is an object created with the Kinetica Workbench UI or with SQL syntax. This function will convert the data in the context to a list of messages that can be used as a prompt. The messages will contain a ``SystemMessage`` followed by pairs of ``HumanMessage``/``AIMessage`` that contain the samples. Args: context_name: The name of an LLM context in the database. Returns: A list of messages containing the information from the context. z/GENERATE PROMPT WITH OPTIONS (CONTEXT_NAMES = 'z')PromptrUrKc:g|]}|Sr?rrs r6rz;ChatKinetica.load_messages_from_context..s'NNN1D33A66NNNr8) _execute_sqljsonloadsr[model_validater\r+rPrGrWrS) r1rsqlr5rs prompt_jsonrequestr\rrQs ` r6load_messages_from_contextz'ChatKinetica.load_messages_from_contexts"Q PPP""3''!j(( &44[AA/ Tx9O9O9Q9QRRRSSSW1133444NNNN NNNr8c0t|}tj|}|jd|}tj|}|d}|dkrs|d}t jd}||} | 1| d} tj| }|d}t||d } t | } | j dkrtd | jS) z/Submit a /chat/completions request to Kinetica.rVz/chat/completionsrvOKrczresponse:({.*})NrwzSQL Generation failed)rPrdumpsr_GPUdb__submit_request_jsonrrrsearchgroupr-rurrvrw) r1rQr request_json response_raw response_jsonrvrc match_respr5rXrws r6rzChatKinetica._submit_completions )))z'** y<<    <00 x( T>>#I.G$677J&&w//F!!<<?? $ 8 4 4 ' 2W%% %V$)88>> ?d " "455 5}r8rc<|j|dd}|d}|ddkr|d}t||d}t|dkrtd |d }i}|D] \}} | ||< |S) z+Execute an SQL query and return the result.rF)limitget_column_major status_inforvrrcrecordszNo records returned.r)rexecute_sql_and_decoder-r,rO) r1rrXrrcrrecord response_dictcolvals r6rzChatKinetica._execute_sqls933 q54  }- x D ( (!),GW%% %9% w<<1  344 4   % %HC!$M#  r8 sa_datafilecdt|}||}|S)z8Load a lanchain prompt from a Kinetica context datafile.)ryr_convert_dict_to_messages)rr datafile_dictrQs r6load_messages_from_datafilez(ChatKinetica.load_messages_from_datafiles/6II+VV 00??r8rcctt|j}t|trd}nBt|t rd}n*t|t rd}ntd|t||}|S)z*Convert a single message to a BaseMessage.rJrNrUzGot unsupported message type: rK) r r=rM isinstancerrrr-rP)rrcrMrLresult_messages r6rz%ChatKinetica._convert_message_to_dictssGO,, g| , , IDD  + + IDD  / / IDDGgGGHH H4999r8c|d}|d}|dkrt|S|dkrt|S|dkrt|Std|)z,Convert a single message from a BaseMessage.rLrMrJrMrNrUzGot unsupported role: )rrrr-)rrcrLrMs r6rz'ChatKinetica._convert_message_from_dictsv)$ 6>>000 0 [ W--- - X   111 1.s'PPP!C::1==PPPr8)rrr+rrW)rr rrUrQ result_lists` r6rz&ChatKinetica._convert_dict_to_messagess"":& 999:::)+ =888999PPPPxPPPQQQr8)NN)r9r:r;r<rrrr>r rrpropertyr=rrrrr r rrrrlrrrrrrrrr?r8r6rr&sW..`d###D##### $4X$3$$$X$ T#s(^   X %):>   {# tCy! 67         8"s"t""""H4:/8,dtK?P[ {t[  > >+ > > >[ >  k9J   [   r8rcteZdZUdZedZeed< edZe ed< e dZ dS) KineticaSqlResponsezResponse containing SQL and the fetched data. This object is returned by a chain with ``KineticaSqlOutputParser`` and it contains the generated SQL and related Pandas Dataframe fetched from the database. r))rrN dataframeTarbitrary_types_allowed) r9r:r;r<rrr=r>rrr model_configr?r8r6rrsw uR   C   U4(((Is(((;: $LLLr8rceZdZUdZedZeed< edZ de de fdZ d d d e ed ede fd Zede fdZdS)KineticaSqlOutputParseraFetch and return data from the Kinetica LLM. This object is used as the last element of a chain to execute generated SQL and it will output a ``KineticaSqlResponse`` containing the SQL and a pandas dataframe with the fetched data. Example: .. code-block:: python from langchain_community.chat_models.kinetica import ( KineticaChatLLM, KineticaSqlOutputParser) kinetica_llm = KineticaChatLLM() # create chain ctx_messages = kinetica_llm.load_messages_from_context(self.context_name) ctx_messages.append(("human", "{input}")) prompt_template = ChatPromptTemplate.from_messages(ctx_messages) chain = ( prompt_template | kinetica_llm | KineticaSqlOutputParser(kdbc=kinetica_llm.kdbc) ) sql_response: KineticaSqlResponse = chain.invoke( {"input": "What are the female users ordered by username?"} ) assert isinstance(sql_response, KineticaSqlResponse) LOG.info(f"SQL Response: {sql_response.sql}") assert isinstance(sql_response.dataframe, pd.DataFrame) Trrrr{r%cX|j|}t||S)N)rr)rto_dfr)r1r{dfs r6parsezKineticaSqlOutputParser.parsePs( Y__T " ""tr::::r8F)partialr5rcB||djS)Nr)rr{)r1r5rs r6 parse_resultz$KineticaSqlOutputParser.parse_resultTszz&).)))r8cdS)Nkinetica_sql_output_parserr?rs r6_typezKineticaSqlOutputParser._typeYs++r8N)r9r:r;r<rrrr>rrr=rrrrboolrrr!r?r8r6rr)s>d###D#####: $L;#;"5;;;; r/s'&  &&&&&&JJJJJJJJJJJJJJJJJJ))))))LLL======DDDDDD EDDDDDIIIIIIIIII1111111111g!! ''''''''T""""""""J     )   )  @ @ @ @ @i @ @ @Y=J=J=J=J=J=J=J=J@@ @ @ @ @ @ @ @ Fmmmmm=mmm`)&2,2,2,2,2,./BC2,2,2,2,2,r8