debug logs shown only in localhost

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import os
 import pandas as pd
 from typing import Literal, TypedDict
-from sqlalchemy import create_engine, inspect
+from sqlalchemy import create_engine, inspect, text
 import json
 from transformers import AutoTokenizer
 from utils import pprint
@@ -155,7 +155,58 @@ def get_table_schema(table_name):
     
     inspector = inspect(st.session_state.engine)
     columns = inspector.get_columns(table_name)
-    return {col['name']: str(col['type']) for col in columns}
+    schema = {col['name']: str(col['type']) for col in columns}
+    
+    # Get table comment
+    table_comment_query = """
+        SELECT obj_description(c.oid) as table_comment
+        FROM pg_class c
+        JOIN pg_namespace n ON n.oid = c.relnamespace
+        WHERE c.relname = :table_name
+        AND n.nspname = 'public'
+    """
+    
+    # Get column comments
+    column_comments_query = """
+        SELECT 
+            cols.column_name,
+            (
+                SELECT pg_catalog.col_description(c.oid, cols.ordinal_position::int)
+                FROM pg_catalog.pg_class c
+                WHERE c.oid = (SELECT ('"' || cols.table_name || '"')::regclass::oid)
+                    AND c.relname = cols.table_name
+            ) as column_comment
+        FROM information_schema.columns cols
+        WHERE cols.table_name = :table_name
+        AND cols.table_schema = 'public'
+    """
+    
+    try:
+        with st.session_state.engine.connect() as conn:
+            # Get table comment
+            table_comment_result = conn.execute(text(table_comment_query), {"table_name": table_name}).fetchone()
+            table_comment = table_comment_result[0] if table_comment_result else None
+            
+            # Get column comments
+            column_comments_result = conn.execute(text(column_comments_query), {"table_name": table_name}).fetchall()
+            column_comments = {row[0]: row[1] for row in column_comments_result}
+            
+            # Create enhanced schema dictionary
+            enhanced_schema = {
+                "table_comment": table_comment,
+                "columns": {
+                    col_name: {
+                        "type": schema[col_name],
+                        "comment": column_comments.get(col_name)
+                    }
+                    for col_name in schema
+                }
+            }
+            
+            return enhanced_schema
+    except Exception as e:
+        st.error(f"Error fetching schema comments: {str(e)}")
+        return schema  # Fallback to basic schema if comment retrieval fails
 
 
 def get_sample_data(table_name):
@@ -223,10 +274,12 @@ User Query: {user_query}
     pprint(f"\n[{MODEL}] Prompt tokens for SQL generation: {prompt_tokens}")
 
     # Debug prompt in a Streamlit expander for better organization
-    with st.expander("Debug: Prompt Generation"):
-        st.write(f"\nUser Query: {user_query}")
-        st.write("\nFull Prompt:")
-        st.code(prompt, language="text")
+    # Check if running locally based on Streamlit's origin header
+    if 'localhost' in st.context.headers.get("Origin", ""):
+        with st.expander("Debug: Prompt Generation"):
+            st.write(f"\nUser Query: {user_query}")
+            st.write("\nFull Prompt:")
+            st.code(prompt, language="text")
 
     start_time = time.time()
     with st.spinner(f"Generating SQL query using {MODEL}..."):