Merge pull request #61 from edenartlab/feat/sentry-config

Add contexts for clients, tools, and middleware for API for error rep…

eve/__init__.py

@@ -17,24 +17,23 @@
 def setup_sentry():
     sentry_dsn = os.getenv("SENTRY_DSN")
     if not sentry_dsn:
+        print("No Sentry DSN found, skipping Sentry setup")
         return
 
+    print(f"Setting up sentry for {db}")
+    # Determine environment
     sentry_env = "production" if db == "PROD" else "staging"
-    if db == "PROD":
-        traces_sample_rate = 0.1
-        profiles_sample_rate = 0.05
-
-    else:
-        traces_sample_rate = 1.0
-        profiles_sample_rate = 1.0
-
-    if sentry_dsn:
-        sentry_sdk.init(
-            dsn=sentry_dsn,
-            traces_sample_rate=traces_sample_rate,
-            profiles_sample_rate=profiles_sample_rate,
-            environment=sentry_env,
-        )
+
+    # Set sampling rates
+    traces_sample_rate = 0.1 if db == "PROD" else 1.0
+    profiles_sample_rate = 0.05 if db == "PROD" else 1.0
+
+    sentry_sdk.init(
+        dsn=sentry_dsn,
+        traces_sample_rate=traces_sample_rate,
+        profiles_sample_rate=profiles_sample_rate,
+        environment=sentry_env,
+    )
 
 
 def load_env(db):
@@ -92,6 +91,8 @@ def verify_env():
 db = os.getenv("DB", "STAGE").upper()
 
 if db not in ["STAGE", "PROD", "WEB3-STAGE", "WEB3-PROD"]:
-    raise Exception(f"Invalid environment: {db}. Must be STAGE, PROD, WEB3-STAGE, or WEB3-PROD")
+    raise Exception(
+        f"Invalid environment: {db}. Must be STAGE, PROD, WEB3-STAGE, or WEB3-PROD"
+    )
 
 load_env(db)

eve/api/api.py

@@ -2,13 +2,16 @@
 import os
 import threading
 import json
+from fastapi.responses import JSONResponse
 import modal
 from fastapi import FastAPI, Depends, BackgroundTasks, Request
 from fastapi.middleware.cors import CORSMiddleware
 from fastapi.security import APIKeyHeader, HTTPBearer
 from apscheduler.schedulers.background import BackgroundScheduler
 from pathlib import Path
 from contextlib import asynccontextmanager
+from starlette.middleware.base import BaseHTTPMiddleware
+import sentry_sdk
 
 from eve import auth, db
 from eve.postprocessing import (
@@ -67,7 +70,34 @@ async def lifespan(app: FastAPI):
             app.state.scheduler.shutdown(wait=True)
 
 
+class SentryContextMiddleware(BaseHTTPMiddleware):
+    async def dispatch(self, request, call_next):
+        with sentry_sdk.configure_scope() as scope:
+            scope.set_tag("package", "eve-api")
+
+            # Extract client context from headers
+            client_platform = request.headers.get("X-Client-Platform")
+            client_agent = request.headers.get("X-Client-Agent")
+
+            if client_platform:
+                scope.set_tag("client_platform", client_platform)
+            if client_agent:
+                scope.set_tag("client_agent", client_agent)
+
+            scope.set_context(
+                "api",
+                {
+                    "endpoint": request.url.path,
+                    "modal_serve": os.getenv("MODAL_SERVE"),
+                    "client_platform": client_platform,
+                    "client_agent": client_agent,
+                },
+            )
+        return await call_next(request)
+
+
 web_app = FastAPI(lifespan=lifespan)
+web_app.add_middleware(SentryContextMiddleware)
 web_app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -167,6 +197,15 @@ async def trigger_delete(
     return await handle_trigger_delete(request, scheduler)
 
 
+@web_app.exception_handler(Exception)
+async def catch_all_exception_handler(request, exc):
+    sentry_sdk.capture_exception(exc)
+    return JSONResponse(
+        status_code=500,
+        content={"message": str(exc)},
+    )
+
+
 # Modal app setup
 app = modal.App(
     name=app_name,
@@ -210,17 +249,24 @@ def fastapi_app():
     image=image, concurrency_limit=1, schedule=modal.Period(minutes=15), timeout=3600
 )
 async def postprocessing():
-    try:
-        await cancel_stuck_tasks()
-    except Exception as e:
-        print(f"Error cancelling stuck tasks: {e}")
-
-    try:
-        await run_nsfw_detection()
-    except Exception as e:
-        print(f"Error running nsfw detection: {e}")
-
-    try:
-        await generate_lora_thumbnails()
-    except Exception as e:
-        print(f"Error generating lora thumbnails: {e}")
+    with sentry_sdk.configure_scope() as scope:
+        scope.set_tag("component", "postprocessing")
+        scope.set_context("function", {"name": "postprocessing"})
+
+        try:
+            await cancel_stuck_tasks()
+        except Exception as e:
+            print(f"Error cancelling stuck tasks: {e}")
+            sentry_sdk.capture_exception(e)
+
+        try:
+            await run_nsfw_detection()
+        except Exception as e:
+            print(f"Error running nsfw detection: {e}")
+            sentry_sdk.capture_exception(e)
+
+        try:
+            await generate_lora_thumbnails()
+        except Exception as e:
+            print(f"Error generating lora thumbnails: {e}")
+            sentry_sdk.capture_exception(e)

eve/clients/common.py

@@ -1,7 +1,9 @@
 import os
 import time
+import asyncio
 
 from eve.models import ClientType
+from .. import sentry_sdk
 
 db = os.getenv("DB", "STAGE")
 
@@ -53,6 +55,25 @@
 day_timestamps = {}
 
 
+def client_context(client_platform: str):
+    """Decorator to add client context to Sentry for all async methods"""
+
+    def decorator(cls):
+        for name, method in cls.__dict__.items():
+            if asyncio.iscoroutinefunction(method):
+
+                async def wrapped_method(self, *args, __method=method, **kwargs):
+                    with sentry_sdk.configure_scope() as scope:
+                        scope.set_tag("client_platform", client_platform)
+                        scope.set_tag("client_agent", self.agent.username)
+                    return await __method(self, *args, **kwargs)
+
+                setattr(cls, name, wrapped_method)
+        return cls
+
+    return decorator
+
+
 def user_over_rate_limits(user):
     user_id = str(user.id)
 

eve/clients/discord/client.py

@@ -40,6 +40,7 @@ def replace_mentions_with_usernames(
     return message_content.strip()
 
 
+@common.client_context("discord")
 class Eden2Cog(commands.Cog):
     def __init__(
         self,
@@ -285,7 +286,11 @@ async def on_message(self, message: discord.Message) -> None:
                 async with session.post(
                     f"{self.api_url}/chat",
                     json=request_data,
-                    headers={"Authorization": f"Bearer {os.getenv('EDEN_ADMIN_KEY')}"},
+                    headers={
+                        "Authorization": f"Bearer {os.getenv('EDEN_ADMIN_KEY')}",
+                        "X-Client-Platform": "discord",
+                        "X-Client-Agent": self.agent.username,
+                    },
                 ) as response:
                     if response.status != 200:
                         error_msg = await response.text()
@@ -394,12 +399,17 @@ def start(
 
         agent_name = os.getenv("EDEN_AGENT_USERNAME")
         agent = Agent.load(agent_name)
-        logger.info(f"Launching Discord bot {agent.username}...")
-
-        bot_token = os.getenv("CLIENT_DISCORD_TOKEN")
-        bot = DiscordBot()
-        bot.add_cog(Eden2Cog(bot, agent, local=local))
-        bot.run(bot_token)
+        with sentry_sdk.configure_scope() as scope:
+            scope.set_tag("package", "eve-clients")
+            scope.set_tag("client_platform", "discord")
+            scope.set_tag("client_agent", agent_name)
+            scope.set_context("discord", {"agent": agent_name, "local": local})
+            logger.info(f"Launching Discord bot {agent.username}...")
+
+            bot_token = os.getenv("CLIENT_DISCORD_TOKEN")
+            bot = DiscordBot()
+            bot.add_cog(Eden2Cog(bot, agent, local=local))
+            bot.run(bot_token)
     except Exception as e:
         logger.error("Failed to start Discord bot", exc_info=True)
         sentry_sdk.capture_exception(e)

eve/clients/farcaster/client.py

@@ -214,7 +214,11 @@ async def process_webhook(
             async with session.post(
                 f"{api_url}/chat",
                 json=request_data,
-                headers={"Authorization": f"Bearer {os.getenv('EDEN_ADMIN_KEY')}"},
+                headers={
+                    "Authorization": f"Bearer {os.getenv('EDEN_ADMIN_KEY')}",
+                    "X-Client-Platform": "farcaster",
+                    "X-Client-Agent": agent.username,
+                },
             ) as response:
                 if response.status != 200:
                     raise Exception("Failed to process request")

eve/clients/telegram/client.py

@@ -4,6 +4,7 @@
 from ably import AblyRealtime
 import aiohttp
 from dotenv import load_dotenv
+import sentry_sdk
 from telegram import Update
 from telegram.constants import ChatAction
 from telegram.ext import (
@@ -16,7 +17,6 @@
 )
 import asyncio
 
-from ... import load_env
 from ...clients import common
 from ...agent import Agent
 from ...llm import UpdateType
@@ -94,10 +94,7 @@ def replace_bot_mentions(message_text: str, bot_username: str, replacement: str)
 
 
 async def send_response(
-    message_type: str, 
-    chat_id: int, 
-    response: list, 
-    context: ContextTypes.DEFAULT_TYPE
+    message_type: str, chat_id: int, response: list, context: ContextTypes.DEFAULT_TYPE
 ):
     """
     Send messages, photos, or videos based on the type of response.
@@ -117,13 +114,9 @@ async def send_response(
             await context.bot.send_message(chat_id=chat_id, text=item)
 
 
+@common.client_context("telegram")
 class EdenTG:
-    def __init__(
-        self, 
-        token: str, 
-        agent: Agent, 
-        local: bool = False
-    ):
+    def __init__(self, token: str, agent: Agent, local: bool = False):
         self.token = token
         self.agent = agent
         self.tools = agent.get_tools()
@@ -132,15 +125,15 @@ def __init__(
         if local:
             self.api_url = "http://localhost:8000"
         else:
-            self.api_url = os.getenv(f"EDEN_API_URL")
+            self.api_url = os.getenv("EDEN_API_URL")
         self.channel_name = common.get_ably_channel_name(
             agent.name, ClientType.TELEGRAM
         )
 
         # Don't initialize Ably here - we'll do it in setup_ably
         self.ably_client = None
         self.channel = None
-        
+
         self.typing_tasks = {}
 
     async def initialize(self, application):
@@ -161,8 +154,12 @@ async def _typing_loop(self, chat_id: int, application: Application):
         """Keep sending typing action until stopped"""
         try:
             while True:
-                await application.bot.send_chat_action(chat_id=chat_id, action=ChatAction.TYPING)
-                await asyncio.sleep(5)  # Telegram typing status expires after ~5 seconds
+                await application.bot.send_chat_action(
+                    chat_id=chat_id, action=ChatAction.TYPING
+                )
+                await asyncio.sleep(
+                    5
+                )  # Telegram typing status expires after ~5 seconds
         except asyncio.CancelledError:
             pass
 
@@ -269,23 +266,18 @@ async def echo(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
 
         if is_direct_message:
             # print author
-            force_reply = False # No DMs
+            force_reply = False  # No DMs
             return
 
         # Lookup thread
         thread_key = f"telegram-{chat_id}"
         if thread_key not in self.known_threads:
-            self.known_threads[thread_key] = self.agent.request_thread(
-                key=thread_key
-            )
+            self.known_threads[thread_key] = self.agent.request_thread(key=thread_key)
         thread = self.known_threads[thread_key]
 
         # Lookup user
         if user_id not in self.known_users:
-            self.known_users[user_id] = User.from_telegram(
-                user_id, 
-                username
-            )
+            self.known_users[user_id] = User.from_telegram(user_id, username)
         user = self.known_users[user_id]
 
         # Check if user rate limits
@@ -333,7 +325,11 @@ async def echo(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
             async with session.post(
                 f"{self.api_url}/chat",
                 json=request_data,
-                headers={"Authorization": f"Bearer {os.getenv('EDEN_ADMIN_KEY')}"},
+                headers={
+                    "Authorization": f"Bearer {os.getenv('EDEN_ADMIN_KEY')}",
+                    "X-Client-Platform": "telegram",
+                    "X-Client-Agent": self.agent.username,
+                },
             ) as response:
                 print(f"Response from {self.api_url}/chat: {response.status}")
                 # json
@@ -348,16 +344,19 @@ async def echo(self, update: Update, context: ContextTypes.DEFAULT_TYPE):
                     return
 
 
-def start(
-    env: str, 
-    local: bool = False
-) -> None:
+def start(env: str, local: bool = False) -> None:
     print("Starting Telegram client...")
     load_dotenv(env)
 
     agent_name = os.getenv("EDEN_AGENT_USERNAME")
     agent = Agent.load(agent_name)
 
+    with sentry_sdk.configure_scope() as scope:
+        scope.set_tag("package", "eve-clients")
+        scope.set_tag("client_platform", "telegram")
+        scope.set_tag("client_agent", agent_name)
+        scope.set_context("telegram", {"agent": agent_name, "local": local})
+
     bot_token = os.getenv("CLIENT_TELEGRAM_TOKEN")
     if not bot_token:
         raise ValueError("CLIENT_TELEGRAM_TOKEN not found in environment variables")