fix filtering

app/main.py

@@ -130,12 +130,15 @@ if uploaded_file:
         few_shot_examples = None
         remaining_data = df
 
-    # Limit rows to 20 to control costs
-    num_rows_to_send = st.number_input('Select number of rows to send to OpenAI ($$)', value=20)
+    # Limit rows based on user input to control costs
+    num_rows_to_send = st.number_input('Select number of rows to send to OpenAI ($$)',
+                                       min_value=1, max_value=len(remaining_data),
+                                       value=min(20, len(remaining_data)))
     if len(remaining_data) > num_rows_to_send:
         st.warning(f"Only the first {num_rows_to_send} rows of the remaining dataset will be sent to OpenAI to minimize costs.")
 
-    limited_data = remaining_data.head(20)
+    # Apply the limit correctly
+    limited_data = remaining_data.head(num_rows_to_send)
 
     # Prepare Few-Shot Examples for Prompting
     example_rows = []