formatting logs

agentops/client.py

@@ -296,7 +296,7 @@ def start_session(self, tags: Optional[List[str]] = None, config: Optional[Confi
             self._session = None
             return logger.warning("Cannot start session - No server response")
 
-        logger.info('View info on this session at https://app.agentops.ai/drilldown?session_id=%s',
+        logger.info('\x1b[34mView info on this session at https://app.agentops.ai/drilldown?session_id=%s\x1b[0m',
                     self._session.session_id)
 
         return self._session.session_id

agentops/log_config.py

@@ -1,10 +1,28 @@
 import logging
 
+class AgentOpsFormatter(logging.Formatter):
+    blue = "\x1b[34m"
+    bold_red = "\x1b[31;1m"
+    reset = "\x1b[0m"
+    format = "🖇 AgentOps: %(message)s"
+
+    FORMATS = {
+        logging.DEBUG: "(DEBUG) " + format,
+        logging.INFO: format,
+        logging.WARNING: format,
+        logging.ERROR: format,
+        logging.CRITICAL: bold_red + format + reset,
+    }
+
+    def format(self, record):
+        log_fmt = self.FORMATS.get(record.levelno)
+        formatter = logging.Formatter(log_fmt)
+        return formatter.format(record)
+
 logger = logging.getLogger("agentops")
 logger.setLevel(logging.CRITICAL)
 
 handler = logging.StreamHandler()
 handler.setLevel(logging.DEBUG)
-formatter = logging.Formatter('🖇 AgentOps: %(message)s')
-handler.setFormatter(formatter)
+handler.setFormatter(AgentOpsFormatter())
 logger.addHandler(handler)