Merge pull request #27 from Cloud-Code-AI/26-setup-observability-with…

…-litellm-to-track-token-usage

feat: added observability for llm check

.env.example

@@ -20,4 +20,8 @@ AZURE_API_VERSION=
 ANTHROPIC_API_KEY=
 
 ## Google Gemini
-GEMINI_API_KEY=
+GEMINI_API_KEY=
+
+## LITELLM OBSERVABILITY
+SUPABASE_URL=
+SUPABASE_KEY=

api/github_helper/pull_requests.py

@@ -32,6 +32,7 @@ def process_pull_request(payload):
         diff_text=diff_text,
         pull_request_title=pr_title,
         pull_request_desc=pr_description,
+        user=repo_name,
     )
     post_pull_request(comment_url, review_body, access_token)
 
@@ -54,6 +55,7 @@ def process_pr_desc(payload):
         diff_text=diff_text,
         pull_request_title=pr_title,
         pull_request_desc=pr_description,
+        user=repo_name,
     )
     patch_pr_body(pr_url, desc, access_token)
 

api/github_helper/utils.py

@@ -5,7 +5,6 @@
 import logging
 import hmac
 import hashlib
-import json
 
 logger = logging.getLogger(__name__)
 
@@ -61,9 +60,3 @@ def is_github_signature_valid(headers, body):
 
     mac = hmac.new(github_secret, msg=body, digestmod=hashlib.sha256)
     return hmac.compare_digest(mac.hexdigest(), signature)
-
-
-def get_config():
-    with open("config.json", "r") as f:
-        config_data = json.loads(f.read())
-    return config_data

api/main.py

@@ -7,7 +7,7 @@
     process_pr_desc,
 )
 from api.github_helper.utils import is_github_signature_valid
-from api.github_helper.constants import CONFIG_DATA
+from cloudcode.utils.config import CONFIG_DATA
 import logging
 
 logging.basicConfig(

cloudcode/actions/reviews.py

@@ -1,4 +1,5 @@
 from cloudcode.helpers import output, parser
+from typing import Optional
 from cloudcode.llms.provider import LLMProvider
 from cloudcode.llms.prompts import (
     CODE_REVIEW_PROMPT,
@@ -14,23 +15,31 @@ def __init__(self):
         self.provider = LLMProvider(system_prompt=CODE_REVIEW_SYSTEM_PROMPT)
 
     def review_pull_request(
-        self, diff_text: str, pull_request_title: str, pull_request_desc: str
+        self,
+        diff_text: str,
+        pull_request_title: str,
+        pull_request_desc: str,
+        user: Optional[str] = None,
     ):
         prompt = CODE_REVIEW_PROMPT.format(
             PULL_REQUEST_TITLE=pull_request_title,
             PULL_REQUEST_DESC=pull_request_desc,
             CODE_DIFF=diff_text,
         )
 
-        resp = self.provider.chat_completion(prompt)
+        resp = self.provider.chat_completion(prompt, user=user)
 
         body = output.create_pr_review_from_json(parser.extract_json(resp))
 
         # Share the review on pull request
         return body
 
     def generate_pull_request_desc(
-        self, diff_text: str, pull_request_title: str, pull_request_desc: str
+        self,
+        diff_text: str,
+        pull_request_title: str,
+        pull_request_desc: str,
+        user: Optional[str] = None,
     ):
         """
         This method generates a AI powered description for a pull request.
@@ -41,7 +50,7 @@ def generate_pull_request_desc(
             CODE_DIFF=diff_text,
         )
 
-        resp = self.provider.chat_completion(prompt)
+        resp = self.provider.chat_completion(prompt, user=user)
         self.logger.debug(f"PROMPT Generate PR Desc RESP: {resp}")
         body = output.create_pr_description(
             parser.extract_json(resp), pull_request_desc

cloudcode/llms/provider.py

@@ -1,5 +1,6 @@
-from litellm import completion
+import litellm
 from cloudcode.llms.prompts import BASIC_SYSTEM_PROMPT
+from cloudcode.utils.config import CONFIG_DATA
 
 
 class LLMProvider:
@@ -18,16 +19,23 @@ def __init__(
         self.model = model
         self.max_tokens = max_tokens
         self.temperature = temperature
+        if CONFIG_DATA.get("language_model", {}).get(
+            "enable_observability_logging", False
+        ):
+            # set callbacks
+            litellm.success_callback = ["supabase"]
+            litellm.failure_callback = ["supabase"]
 
-    def chat_completion(self, prompt):
+    def chat_completion(self, prompt, user: str = None):
         messages = [
             {"role": "system", "content": self.system_prompt},
             {"role": "user", "content": prompt},
         ]
-        response = completion(
+        response = litellm.completion(
             model=self.model,
             messages=messages,
             max_tokens=self.max_tokens,
             temperature=self.temperature,
+            user=user,
         )
         return response["choices"][0]["message"]["content"]

cloudcode/utils/config.py

@@ -0,0 +1,32 @@
+import json
+import os
+
+
+def get_config():
+    with open("config.json", "r") as f:
+        config_data = json.loads(f.read())
+    return config_data
+
+
+def validate_config_settings(config):
+    "Make sure relvant enviorment variables are set"
+    if config.get("github_app", {}).get("check_signature", False):
+        if not os.environ.get("GITHUB_APP_WEBHOOK_SECRET"):
+            raise EnvironmentError(
+                "The environment variable 'GITHUB_APP_WEBHOOK_SECRET' is not set."
+            )
+
+    if config.get("language_model", {}).get("provider", {}) == "litellm":
+        if config.get("language_model", {}).get("enable_observability_logging", False):
+            if not os.environ.get("SUPABASE_URL"):
+                raise EnvironmentError(
+                    "The environment variable 'SUPABASE_URL' is not set."
+                )
+            if not os.environ.get("SUPABASE_KEY"):
+                raise EnvironmentError(
+                    "The environment variable 'SUPABASE_KEY' is not set."
+                )
+    return config
+
+
+CONFIG_DATA = validate_config_settings(get_config())

config.json

@@ -1,6 +1,7 @@
 {
     "language_model": {
-        "provider": "litellm"
+        "provider": "litellm",
+        "enable_observability_logging": true
     },
     "github_app": {
         "check_signature": false

tests/actions/test_review.py

@@ -15,6 +15,7 @@ def test_review_pull_request(valid_review):
         valid_review["input"]["diff"],
         valid_review["input"]["title"],
         valid_review["input"]["description"],
+        user="pytest",
     )
 
     assert fuzz.ratio(result, valid_review["output"]) > 95