Update app.py

app.py

@@ -1,16 +1,13 @@
-import streamlit as st
-import requests
-import json
 import os
+import requests
+import streamlit as st
 
 # Streamlit app setup
 st.title('Llama Cloud API Document Extraction')
 
-
 # File Upload
 uploaded_file = st.file_uploader('Choose a PDF file', type='pdf')
-
-api_key = os.getenv('api_key')
+api_key = os.getenv('API_KEY')
 
 if uploaded_file is not None:
     # Upload file to Llama Cloud
@@ -20,7 +17,7 @@ if uploaded_file is not None:
     }
     files = {'upload_file': (uploaded_file.name, uploaded_file, 'application/pdf')}
     response = requests.post('https://api.cloud.llamaindex.ai/api/v1/files', headers=headers, files=files)
-    file_id = response.json()['id']
+    file_id = response.json().get('id')
     st.write(f'File uploaded with ID: {file_id}')
 
     # Infer Schema
@@ -29,7 +26,7 @@ if uploaded_file is not None:
         'file_ids': [file_id]
     }
     schema_response = requests.post('https://api.cloud.llamaindex.ai/api/v1/extraction/schemas/infer', headers=headers, json=schema_data)
-    schema_id = schema_response.json()['id']
+    schema_id = schema_response.json().get('id')
     st.write(f'Schema inferred with ID: {schema_id}')
 
     # Start Extraction Job
@@ -38,12 +35,12 @@ if uploaded_file is not None:
         'file_id': file_id
     }
     job_response = requests.post('https://api.cloud.llamaindex.ai/api/v1/extraction/jobs', headers=headers, json=job_data)
-    job_id = job_response.json()['id']
+    job_id = job_response.json().get('id')
     st.write(f'Extraction job started with ID: {job_id}')
 
     # Check Job Status
     status_response = requests.get(f'https://api.cloud.llamaindex.ai/api/v1/extraction/jobs/{job_id}', headers=headers)
-    status = status_response.json()['status']
+    status = status_response.json().get('status')
     st.write(f'Job Status: {status}')
 
     # Display Results
@@ -51,6 +48,6 @@ if uploaded_file is not None:
         results_response = requests.get(f'https://api.cloud.llamaindex.ai/api/v1/extraction/jobs/{job_id}/result', headers=headers)
         results = results_response.json()
         st.write('### Extraction Results:')
-        st.markdown(f'```json\n{json.dumps(results, indent=2)}\n```')
+        st.json(results)
     else:
         st.write('Extraction job is still in progress or has failed.')