Spaces:
Running
Running
| from typing import List | |
| import pandas as pd | |
| import streamlit as st | |
| from langchain.schema import Document | |
| from langchain_experimental.retrievers.vector_sql_database import VectorSQLDatabaseChainRetriever | |
| from backend.chains.retrieval_qa_with_sources import CustomRetrievalQAWithSourcesChain | |
| from backend.constants.myscale_tables import MYSCALE_TABLES | |
| from backend.constants.variables import CHAINS_RETRIEVERS_MAPPING, DIVIDER_HTML, RetrieverButtons | |
| from backend.callbacks.vector_sql_callbacks import VectorSQLSearchDBCallBackHandler, VectorSQLSearchLLMCallBackHandler | |
| from ui.utils import display | |
| from logger import logger | |
| def process_sql_query(selected_table: str, query_type: str): | |
| place_holder = st.empty() | |
| logger.info( | |
| f"button-1: {st.session_state[RetrieverButtons.vector_sql_query_from_db]}, " | |
| f"button-2: {st.session_state[RetrieverButtons.vector_sql_query_with_llm]}, " | |
| f"table: {selected_table}, " | |
| f"content: {st.session_state.query_sql}" | |
| ) | |
| with place_holder.expander('🪵 Query Log', expanded=True): | |
| try: | |
| if query_type == RetrieverButtons.vector_sql_query_from_db: | |
| callback = VectorSQLSearchDBCallBackHandler() | |
| vector_sql_retriever: VectorSQLDatabaseChainRetriever = \ | |
| st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["sql_retriever"] | |
| relevant_docs: List[Document] = vector_sql_retriever.get_relevant_documents( | |
| query=st.session_state.query_sql, | |
| callbacks=[callback] | |
| ) | |
| callback.progress_bar.progress( | |
| value=1.0, | |
| text="[Question -> LLM -> SQL Statement -> MyScaleDB -> Results] Done! ✅" | |
| ) | |
| st.markdown(f"### Vector Search Results from `{selected_table}` \n" | |
| f"> Here we get documents from MyScaleDB with given sql statement \n\n") | |
| display( | |
| pd.DataFrame( | |
| [{**d.metadata, 'abstract': d.page_content} for d in relevant_docs] | |
| ) | |
| ) | |
| elif query_type == RetrieverButtons.vector_sql_query_with_llm: | |
| callback = VectorSQLSearchLLMCallBackHandler(table=selected_table) | |
| vector_sql_chain: CustomRetrievalQAWithSourcesChain = \ | |
| st.session_state[CHAINS_RETRIEVERS_MAPPING][selected_table]["sql_chain"] | |
| chain_results = vector_sql_chain( | |
| inputs=st.session_state.query_sql, | |
| callbacks=[callback] | |
| ) | |
| callback.progress_bar.progress( | |
| value=1.0, | |
| text="[Question -> LLM -> SQL Statement -> MyScaleDB -> " | |
| "(Question,Results) -> LLM -> Results] Done! ✅" | |
| ) | |
| documents_reference: List[Document] = chain_results["source_documents"] | |
| st.markdown(f"### Vector Search Results from `{selected_table}` \n" | |
| f"> Here we get documents from MyScaleDB with given sql statement \n\n") | |
| display( | |
| pd.DataFrame( | |
| [{**d.metadata, 'abstract': d.page_content} for d in documents_reference] | |
| ) | |
| ) | |
| st.markdown( | |
| f"### Answer from LLM \n" | |
| f"> The response of the LLM when given the vector search results. \n\n" | |
| ) | |
| st.write(chain_results['answer']) | |
| st.markdown( | |
| f"### References from `{selected_table}`\n" | |
| f"> Here shows that which documents used by LLM \n\n" | |
| ) | |
| if len(chain_results['sources']) == 0: | |
| st.write("No documents is used by LLM.") | |
| else: | |
| display( | |
| dataframe=pd.DataFrame( | |
| [{**d.metadata, 'abstract': d.page_content} for d in chain_results['sources']] | |
| ), | |
| columns_=['ref_id'] + MYSCALE_TABLES[selected_table].must_have_col_names, | |
| index='ref_id' | |
| ) | |
| else: | |
| raise NotImplementedError(f"Unsupported query type: {query_type}") | |
| st.markdown(DIVIDER_HTML, unsafe_allow_html=True) | |
| except Exception as e: | |
| st.write('Oops 😵 Something bad happened...') | |
| raise e | |