import boto3, json, time

REGION = "us-east-1"
BUCKET = "medgemma-async-2026"
ENDPOINT = "medgemma-vllm-endpoint"

sagemaker = boto3.client("sagemaker-runtime", region_name=REGION)
s3 = boto3.client("s3", region_name=REGION)

payload = {
    "inputs": "What are the common symptoms of type 2 diabetes?",
    "parameters": {"max_new_tokens": 200, "temperature": 0.2}
}

input_key = f"input/test_{int(time.time())}.json"
s3.put_object(Bucket=BUCKET, Key=input_key, Body=json.dumps(payload).encode())

response = sagemaker.invoke_endpoint_async(
    EndpointName=ENDPOINT,
    InputLocation=f"s3://{BUCKET}/{input_key}",
    ContentType="application/json",
)
output_uri = response["OutputLocation"]
output_key = output_uri.replace(f"s3://{BUCKET}/", "")

for i in range(60):
    time.sleep(10)
    try:
        obj = s3.get_object(Bucket=BUCKET, Key=output_key)
        print(obj["Body"].read().decode())
        break
    except s3.exceptions.NoSuchKey:
        print(f"  waiting... {(i+1)*10}s")