From ce0a5e00996b8e47e33a89adda3f82db4a56d20c Mon Sep 17 00:00:00 2001
From: Vishwanath Martur <64204611+vishwamartur@users.noreply.github.com>
Date: Thu, 26 Dec 2024 22:37:01 +0530
Subject: [PATCH] Integrate Watchtower logging into swarms

Related to #702

Integrate Watchtower for optional logging to AWS CloudWatch Logs.

* **Watchtower Integration:**
  - Import `watchtower` in `swarms/utils/loguru_logger.py`.
  - Add `initialize_watchtower_logger` function to configure Watchtower logging.
  - Update `initialize_logger` to optionally use Watchtower logging if configured.

* **Agent Logging Updates:**
  - Import `initialize_watchtower_logger` in `swarms/agents/ape_agent.py`, `swarms/agents/tool_agent.py`, and `swarms/structs/agent.py`.
  - Update logger initialization in `swarms/agents/ape_agent.py` and `swarms/agents/tool_agent.py` to use `initialize_watchtower_logger`.

* **Terraform Configuration:**
  - Add `terraform/logging.tf` to create log groups and policies for AWS CloudWatch Logs.
---
 swarms/agents/ape_agent.py    |  4 +--
 swarms/agents/tool_agent.py   |  4 +--
 swarms/structs/agent.py       |  1 +
 swarms/utils/loguru_logger.py | 11 +++++++-
 terraform/logging.tf          | 53 +++++++++++++++++++++++++++++++++++
 5 files changed, 68 insertions(+), 5 deletions(-)
 create mode 100644 terraform/logging.tf

diff --git a/swarms/agents/ape_agent.py b/swarms/agents/ape_agent.py
index 420b7aaa1..9712a03b2 100644
--- a/swarms/agents/ape_agent.py
+++ b/swarms/agents/ape_agent.py
@@ -8,9 +8,9 @@
 from swarms.prompts.prompt_generator_optimizer import (
     prompt_generator_sys_prompt,
 )
-from swarms.utils.loguru_logger import initialize_logger
+from swarms.utils.loguru_logger import initialize_logger, initialize_watchtower_logger
 
-logger = initialize_logger(log_folder="ape_agent")
+logger = initialize_watchtower_logger(log_folder="ape_agent")
 
 
 @retry(
diff --git a/swarms/agents/tool_agent.py b/swarms/agents/tool_agent.py
index b686f3b0a..2714f2e0f 100644
--- a/swarms/agents/tool_agent.py
+++ b/swarms/agents/tool_agent.py
@@ -1,9 +1,9 @@
 from typing import Any, Optional, Callable
 from swarms.tools.json_former import Jsonformer
-from swarms.utils.loguru_logger import initialize_logger
+from swarms.utils.loguru_logger import initialize_logger, initialize_watchtower_logger
 from swarms.utils.lazy_loader import lazy_import_decorator
 
-logger = initialize_logger(log_folder="tool_agent")
+logger = initialize_watchtower_logger(log_folder="tool_agent")
 
 
 @lazy_import_decorator
diff --git a/swarms/structs/agent.py b/swarms/structs/agent.py
index d6caed666..207ebb4bb 100644
--- a/swarms/structs/agent.py
+++ b/swarms/structs/agent.py
@@ -53,6 +53,7 @@
 from swarms.utils.wrapper_clusterop import (
     exec_callable_with_clusterops,
 )
+from swarms.utils.loguru_logger import initialize_watchtower_logger
 
 
 # Utils
diff --git a/swarms/utils/loguru_logger.py b/swarms/utils/loguru_logger.py
index af5c72399..e32068780 100644
--- a/swarms/utils/loguru_logger.py
+++ b/swarms/utils/loguru_logger.py
@@ -1,9 +1,14 @@
 import os
 import uuid
 from loguru import logger
+import watchtower
 
+def initialize_watchtower_logger(log_group: str = "swarms_logs"):
+    handler = watchtower.CloudWatchLogHandler(log_group=log_group)
+    logger.add(handler)
+    return logger
 
-def initialize_logger(log_folder: str = "logs"):
+def initialize_logger(log_folder: str = "logs", use_watchtower: bool = False, log_group: str = "swarms_logs"):
 
     AGENT_WORKSPACE = "agent_workspace"
 
@@ -34,4 +39,8 @@ def initialize_logger(log_folder: str = "logs"):
         retention="10 days",
         # compression="zip",
     )
+
+    if use_watchtower:
+        initialize_watchtower_logger(log_group)
+
     return logger
diff --git a/terraform/logging.tf b/terraform/logging.tf
new file mode 100644
index 000000000..8f9255a72
--- /dev/null
+++ b/terraform/logging.tf
@@ -0,0 +1,53 @@
+provider "aws" {
+  region = "us-west-2"
+}
+
+resource "aws_cloudwatch_log_group" "swarms_log_group" {
+  name              = "swarms-log-group"
+  retention_in_days = 14
+}
+
+resource "aws_iam_role" "swarms_logging_role" {
+  name = "swarms-logging-role"
+
+  assume_role_policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action = "sts:AssumeRole"
+        Effect = "Allow"
+        Principal = {
+          Service = "ec2.amazonaws.com"
+        }
+      }
+    ]
+  })
+}
+
+resource "aws_iam_policy" "swarms_logging_policy" {
+  name        = "swarms-logging-policy"
+  description = "Policy for allowing swarms to create and manage CloudWatch log groups"
+  policy      = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action = [
+          "logs:CreateLogGroup",
+          "logs:CreateLogStream",
+          "logs:PutLogEvents",
+          "logs:DescribeLogGroups",
+          "logs:DescribeLogStreams",
+          "logs:PutRetentionPolicy",
+          "logs:TagLogGroup"
+        ]
+        Effect   = "Allow"
+        Resource = "*"
+      }
+    ]
+  })
+}
+
+resource "aws_iam_role_policy_attachment" "swarms_logging_role_attachment" {
+  role       = aws_iam_role.swarms_logging_role.name
+  policy_arn = aws_iam_policy.swarms_logging_policy.arn
+}