xiaxinhang/pr-agent
1
0
Fork 0
mirror of https://github.com/qodo-ai/pr-agent.git synced 2025-07-02 03:40:38 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #1145 from MarkRx/feature/litellm-logging-observability

Browse Source 
Add and document abilty to use LiteLLM Logging Observability tools
This commit is contained in:
Tal
2024-08-22 09:58:53 +03:00
committed by GitHub
parent 771d0b8c60 8aa76a0ac5
commit 745e955d1f
4 changed files with 91 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

66
pr_agent/algo/ai_handlers/litellm_ai_handler.py
View File

@ -1,10 +1,10 @@
import os
import requests
import boto3
import litellm
import openai
from litellm import acompletion
from tenacity import retry, retry_if_exception_type, stop_after_attempt
from pr_agent.algo.ai_handlers.base_ai_handler import BaseAiHandler
from pr_agent.config_loader import get_settings
from pr_agent.log import get_logger
@ -44,6 +44,12 @@ class LiteLLMAIHandler(BaseAiHandler):
litellm.use_client = True
if get_settings().get("LITELLM.DROP_PARAMS", None):
litellm.drop_params = get_settings().litellm.drop_params
if get_settings().get("LITELLM.SUCCESS_CALLBACK", None):
litellm.success_callback = get_settings().litellm.success_callback
if get_settings().get("LITELLM.FAILURE_CALLBACK", None):
litellm.failure_callback = get_settings().litellm.failure_callback
if get_settings().get("LITELLM.SERVICE_CALLBACK", None):
litellm.service_callback = get_settings().litellm.service_callback
if get_settings().get("OPENAI.ORG", None):
litellm.organization = get_settings().openai.org
if get_settings().get("OPENAI.API_TYPE", None):
@ -89,6 +95,60 @@ class LiteLLMAIHandler(BaseAiHandler):
response_log['main_pr_language'] = 'unknown'
return response_log
def add_litellm_callbacks(selfs, kwargs) -> dict:
captured_extra = []
def capture_logs(message):
# Parsing the log message and context
record = message.record
log_entry = {}
if record.get('extra', None).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
captured_extra.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)
context = captured_extra[0] if len(captured_extra) > 0 else None
command = context.get("command", "unknown")
pr_url = context.get("pr_url", "unknown")
git_provider = get_settings().config.git_provider
metadata = dict()
callbacks = litellm.success_callback + litellm.failure_callback + litellm.service_callback
if "langfuse" in callbacks:
metadata.update({
"trace_name": command,
"tags": [git_provider, command],
"trace_metadata": {
"command": command,
"pr_url": pr_url,
},
})
if "langsmith" in callbacks:
metadata.update({
"run_name": command,
"tags": [git_provider, command],
"extra": {
"metadata": {
"command": command,
"pr_url": pr_url,
}
},
})
# Adding the captured logs to the kwargs
kwargs["metadata"] = metadata
return kwargs
@property
def deployment_id(self):
"""
@ -133,6 +193,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_litellm_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")