From c76aabc71eea58acab1412806ebb6a1c4420f3ec Mon Sep 17 00:00:00 2001
From: mrT23 <tal.r@codium.ai>
Date: Sat, 17 Aug 2024 09:15:05 +0300
Subject: [PATCH] Add callback functionality to litellm_ai_handler for enhanced
 logging and metadata capture

---
 pr_agent/agent/pr_agent.py                    |  2 +-
 .../algo/ai_handlers/litellm_ai_handler.py    | 29 +++++++++++++++++++
 pr_agent/settings/configuration.toml          |  1 +
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/pr_agent/agent/pr_agent.py b/pr_agent/agent/pr_agent.py
index d0ac46ca..8bf6cff7 100644
--- a/pr_agent/agent/pr_agent.py
+++ b/pr_agent/agent/pr_agent.py
@@ -79,7 +79,7 @@ class PRAgent:
         if action not in command2class:
             get_logger().debug(f"Unknown command: {action}")
             return False
-        with get_logger().contextualize(command=action):
+        with get_logger().contextualize(command=action, pr_url=pr_url):
             get_logger().info("PR-Agent request handler started", analytics=True)
             if action == "reflect_and_review":
                 get_settings().pr_reviewer.ask_and_reflect = True
diff --git a/pr_agent/algo/ai_handlers/litellm_ai_handler.py b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
index c8b620fe..1aa93b36 100644
--- a/pr_agent/algo/ai_handlers/litellm_ai_handler.py
+++ b/pr_agent/algo/ai_handlers/litellm_ai_handler.py
@@ -89,6 +89,31 @@ class LiteLLMAIHandler(BaseAiHandler):
             response_log['main_pr_language'] = 'unknown'
         return response_log
 
+    def add_callbacks(selfs, kwargs):
+        pr_metadata = []
+
+        def capture_logs(message):
+            # Parsing the log message and context
+            record = message.record
+            log_entry = {}
+            if record.get('extra', {}).get('command', None) is not None:
+                log_entry.update({"command": record['extra']["command"]})
+            if record.get('extra', {}).get('pr_url', None) is not None:
+                log_entry.update({"pr_url": record['extra']["pr_url"]})
+
+            # Append the log entry to the captured_logs list
+            pr_metadata.append(log_entry)
+
+        # Adding the custom sink to Loguru
+        handler_id = get_logger().add(capture_logs)
+        get_logger().debug("Capturing logs for litellm callbacks")
+        get_logger().remove(handler_id)
+
+        # Adding the captured logs to the kwargs
+        kwargs["metadata"] = pr_metadata
+
+        return kwargs
+
     @property
     def deployment_id(self):
         """
@@ -133,6 +158,10 @@ class LiteLLMAIHandler(BaseAiHandler):
                 "force_timeout": get_settings().config.ai_timeout,
                 "api_base": self.api_base,
             }
+
+            if get_settings().litellm.get("enable_callbacks", False):
+                kwargs = self.add_callbacks(kwargs)
+
             seed = get_settings().config.get("seed", -1)
             if temperature > 0 and seed >= 0:
                 raise ValueError(f"Seed ({seed}) is not supported with temperature ({temperature}) > 0")
diff --git a/pr_agent/settings/configuration.toml b/pr_agent/settings/configuration.toml
index 67273c46..79d30f1d 100644
--- a/pr_agent/settings/configuration.toml
+++ b/pr_agent/settings/configuration.toml
@@ -264,6 +264,7 @@ pr_commands = [
 [litellm]
 # use_client = false
 # drop_params = false
+enable_callbacks = false
 
 [pr_similar_issue]
 skip_comments = false