Update app.py

app.py

@@ -57,77 +57,81 @@ def query_rag_agent(query: str):
         logger.error("Graph recursion limit reached; query processing failed.")
         return "Graph recursion limit reached. No SQL result generated.", ""
 
-URL = "https://storage.googleapis.com/benchmarks-artifacts/travel-db/travel2.sqlite"
-download_sqlite_db(URL)
-
-with st.sidebar:
-    st.header("About Project")
-    st.markdown(
-        """
-        RAG (Retrieval-Augmented Generation) Agent SQL is an approach that combines retrieval techniques with text generation to create more relevant and contextualised answers from data, 
-        particularly in SQL databases. RAG-Agent SQL uses two main components:
-        - Retrieval: Retrieving relevant information from the database based on a given question or input.
-        - Augmented Generation: Using natural language models (e.g., LLMs such as GPT or LLaMA) to generate more detailed answers, using information from the retrieval results.
-
-        to see the architecture can be seen here [Github](https://github.com/fahmiaziz98/sql_agent/tree/main/002sql-agent-ra)
-        """
-    )
-
-    st.header("Example Question")
-    st.markdown(
-        """
-        - How many different aircraft models are there? And what are the models?
-        - What is the aircraft model with the longest range?
-        - Which airports are located in the city of Basel?
-        - Can you please provide information on what I asked before?
-        - What are the fare conditions available on Boeing 777-300?
-        - What is the total amount of bookings made in April 2024?
-        - What is the scheduled arrival time of flight number QR0051?
-        - Which car rental services are available in Basel?
-        - Which seat was assigned to the boarding pass with ticket number 0060005435212351?
-        - Which trip recommendations are related to history in Basel?
-        - How many tickets were sold for Business class on flight 30625?
-        - Which hotels are located in Zurich?
-        """
-    )
-
-# Main Application Title
-st.title("RAG SQL-Agent")
-
-# Initialize session state for storing chat messages
-if "messages" not in st.session_state:
-    st.session_state.messages = []
-
-# Display conversation history from session state
-for message in st.session_state.messages:
-    role = message.get("role", "assistant")
-    with st.chat_message(role):
-        if "output" in message:
-            st.markdown(message["output"])
-        if "sql_query" in message and message["sql_query"]:
+def main():
+    with st.sidebar:
+        st.header("About Project")
+        st.markdown(
+            """
+            RAG (Retrieval-Augmented Generation) Agent SQL is an approach that combines retrieval techniques with text generation to create more relevant and contextualised answers from data, 
+            particularly in SQL databases. RAG-Agent SQL uses two main components:
+            - Retrieval: Retrieving relevant information from the database based on a given question or input.
+            - Augmented Generation: Using natural language models (e.g., LLMs such as GPT or LLaMA) to generate more detailed answers, using information from the retrieval results.
+    
+            to see the architecture can be seen here [Github](https://github.com/fahmiaziz98/sql_agent/tree/main/002sql-agent-ra)
+            """
+        )
+    
+        st.header("Example Question")
+        st.markdown(
+            """
+            - How many different aircraft models are there? And what are the models?
+            - What is the aircraft model with the longest range?
+            - Which airports are located in the city of Basel?
+            - Can you please provide information on what I asked before?
+            - What are the fare conditions available on Boeing 777-300?
+            - What is the total amount of bookings made in April 2024?
+            - What is the scheduled arrival time of flight number QR0051?
+            - Which car rental services are available in Basel?
+            - Which seat was assigned to the boarding pass with ticket number 0060005435212351?
+            - Which trip recommendations are related to history in Basel?
+            - How many tickets were sold for Business class on flight 30625?
+            - Which hotels are located in Zurich?
+            """
+        )
+
+    # Main Application Title
+    st.title("RAG SQL-Agent")
+    
+    # Initialize session state for storing chat messages
+    if "messages" not in st.session_state:
+        st.session_state.messages = []
+    
+    # Display conversation history from session state
+    for message in st.session_state.messages:
+        role = message.get("role", "assistant")
+        with st.chat_message(role):
+            if "output" in message:
+                st.markdown(message["output"])
+            if "sql_query" in message and message["sql_query"]:
+                with st.expander("SQL Query", expanded=True):
+                    st.code(message["sql_query"])
+    
+    # Input form for user prompt
+    if prompt := st.chat_input("What do you want to know?"):
+        st.chat_message("user").markdown(prompt)
+        st.session_state.messages.append({"role": "user", "output": prompt})
+    
+        # Fetch response from RAG agent function directly
+        with st.spinner("Searching for an answer..."):
+            output_text, sql_query = query_rag_agent(prompt)
+    
+        # Display assistant response and SQL query
+        st.chat_message("assistant").markdown(output_text)
+        if sql_query:
             with st.expander("SQL Query", expanded=True):
-                st.code(message["sql_query"])
-
-# Input form for user prompt
-if prompt := st.chat_input("What do you want to know?"):
-    st.chat_message("user").markdown(prompt)
-    st.session_state.messages.append({"role": "user", "output": prompt})
-
-    # Fetch response from RAG agent function directly
-    with st.spinner("Searching for an answer..."):
-        output_text, sql_query = query_rag_agent(prompt)
-
-    # Display assistant response and SQL query
-    st.chat_message("assistant").markdown(output_text)
-    if sql_query:
-        with st.expander("SQL Query", expanded=True):
-            st.code(sql_query)
-
-    # Append assistant response to session state
-    st.session_state.messages.append(
-        {
-            "role": "assistant",
-            "output": output_text,
-            "sql_query": sql_query,
-        }
-    )
+                st.code(sql_query)
+    
+        # Append assistant response to session state
+        st.session_state.messages.append(
+            {
+                "role": "assistant",
+                "output": output_text,
+                "sql_query": sql_query,
+            }
+        )
+
+if __name__ == "__main__":
+    URL = "https://storage.googleapis.com/benchmarks-artifacts/travel-db/travel2.sqlite"
+    download_sqlite_db(URL)
+    main()
+