diff --git a/api-example.py b/api-example.py
index 4bf4f0d..d6053fd 100644
--- a/api-example.py
+++ b/api-example.py
@@ -7,7 +7,7 @@ URI = f'http://{HOST}/api/v1/generate'
 # For reverse-proxied streaming, the remote will likely host with ssl - https://
 # URI = 'https://your-uri-here.trycloudflare.com/api/v1/generate'
 
-def run(context):
+def run(prompt):
     request = {
         'prompt': prompt,
         'max_new_tokens': 250,