make sqs queue configurable

personal/khk/bot-and-runner/sqs-runner.py

@@ -0,0 +1,152 @@
+#
+# requires AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY environment variables
+#
+
+import boto3
+import json
+import subprocess
+import signal
+import os
+import time
+
+from pydantic import BaseModel, ValidationError
+from typing import Optional
+
+from pipecat.transports.services.helpers.daily_rest import DailyRESTHelper, DailyRoomObject, DailyRoomProperties, DailyRoomParams
+
+from dotenv import load_dotenv
+load_dotenv(override=True)
+
+# SQS queue URL
+QUEUE_URL = os.environ.get('SQS_QUEUE_URL')
+
+# The program to be spawned
+SUBPROCESS_PROGRAM = 'your_subprocess_program.py'
+
+
+def escape_bash_arg(s):
+    return "'" + s.replace("'", "'\\''") + "'"
+
+# ------------ Configuration ------------ #
+
+
+MAX_SESSION_TIME = 5 * 60  # 5 minutes
+REQUIRED_ENV_VARS = [
+    'DAILY_API_KEY',
+    'CARTESIA_API_KEY',
+    'SQS_QUEUE_URL',
+    'AWS_ACCESS_KEY_ID',
+    'AWS_SECRET_ACCESS_KEY']
+
+print("DAILY_API_KEY", os.getenv("DAILY_API_KEY"))
+daily_rest_helper = DailyRESTHelper(
+    os.getenv("DAILY_API_KEY", ""),
+    os.getenv("DAILY_API_URL", 'https://api.daily.co/v1'))
+
+# ----------------- API ----------------- #
+
+
+def setup_sqs():
+    """Set up the SQS client."""
+    return boto3.client(
+        'sqs',
+        region_name='us-west-2',
+        # use an iam user instead of role because passing a role into an eks pod requires
+        # adding an add-on and we don't want to change the eks cluster configuration if we
+        # can avoid it
+        aws_access_key_id=os.environ.get('AWS_ACCESS_KEY_ID'),
+        aws_secret_access_key=os.environ.get('AWS_SECRET_ACCESS_KEY')
+    )
+
+
+def receive_message(sqs):
+    """Receive a message from the SQS queue."""
+    response = sqs.receive_message(
+        QueueUrl=QUEUE_URL,
+        MaxNumberOfMessages=1,
+        WaitTimeSeconds=20  # Long polling
+    )
+
+    messages = response.get('Messages', [])
+    if messages:
+        return messages[0]
+    return None
+
+
+def delete_message(sqs, receipt_handle):
+    """Delete a message from the queue after processing."""
+    sqs.delete_message(
+        QueueUrl=QUEUE_URL,
+        ReceiptHandle=receipt_handle
+    )
+
+
+def run_bot(msg):
+    try:
+        print(f"Getting room from URL: {'https://rtvi.daily.co/' + msg['room']}")
+        room: DailyRoomObject = daily_rest_helper.get_room_from_url(
+            'https://rtvi.daily.co/' + msg['room'])
+    except Exception:
+        raise HTTPException(
+            status_code=500, detail=f"Room not found: {msg['room']}")
+
+    # Give the agent a token to join the session
+    token = daily_rest_helper.get_token(room.url, MAX_SESSION_TIME)
+
+    if not room or not token:
+        raise HTTPException(
+            status_code=500, detail=f"Failed to get token for room: {room.url}")
+
+    # Spawn a new agent, and join the user session
+    try:
+        bot_settings_str = json.dumps(msg['config'])
+
+        print(f"Starting bot with settings: {room.url}, {token}, {bot_settings_str}")
+
+        process = subprocess.Popen(
+            [f"python3 -m bot -u {room.url} -t {token} -c {escape_bash_arg(bot_settings_str)}"],
+            shell=True,
+            bufsize=1,
+            cwd=os.path.dirname(os.path.abspath(__file__)))
+
+        start_time = time.time()
+        while time.time() - start_time < (MAX_SESSION_TIME + 10):
+            if process.poll() is not None:
+                # process has finished
+                print("BOT EXITED BEFORE TIMEOUT")
+                return
+            time.sleep(1)
+        # process did not exit. need to kill -9 it
+        print("KILLING BOT PROCESS")
+        os.kill(process.pid, signal.SIGKILL)
+    except Exception as e:
+        raise HTTPException(
+            status_code=500, detail=f"Failed to start subprocess: {e}")
+
+
+def main():
+    print("sqs-runner.py - started.")
+    sqs = setup_sqs()
+
+    while True:
+        print("sqs-runner.py - polling")
+        message = receive_message(sqs)
+        if message:
+            delete_message(sqs, message['ReceiptHandle'])
+
+            message_body = json.loads(message['Body'])
+            print(f"Received message. {message_body}")
+
+            run_bot(message_body)
+
+
+if __name__ == "__main__":
+    # Check environment variables
+    for env_var in REQUIRED_ENV_VARS:
+        if env_var not in os.environ:
+            raise Exception(f"Missing environment variable: {env_var}.")
+
+    try:
+        main()
+    except KeyboardInterrupt:
+        print("Pipecat runner shutting down...")