feat add return types on team/callback

This commit is contained in:
Ishaan Jaff 2024-07-22 20:35:27 -07:00
parent dcd8f7ebf2
commit 5ed82ba5ce
2 changed files with 101 additions and 48 deletions

View file

@ -1702,3 +1702,4 @@ class ProxyErrorTypes(str, enum.Enum):
budget_exceeded = "budget_exceeded" budget_exceeded = "budget_exceeded"
expired_key = "expired_key" expired_key = "expired_key"
auth_error = "auth_error" auth_error = "auth_error"
internal_server_error = "internal_server_error"

View file

@ -20,6 +20,8 @@ from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import ( from litellm.proxy._types import (
AddTeamCallback, AddTeamCallback,
LiteLLM_TeamTable, LiteLLM_TeamTable,
ProxyErrorTypes,
ProxyException,
TeamCallbackMetadata, TeamCallbackMetadata,
UserAPIKeyAuth, UserAPIKeyAuth,
) )
@ -48,6 +50,28 @@ async def add_team_callbacks(
description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability", description="The litellm-changed-by header enables tracking of actions performed by authorized users on behalf of other users, providing an audit trail for accountability",
), ),
): ):
"""
Add a success/failure callback to a team
Use this if if you want different teams to have different success/failure callbacks
Example curl:
```
curl -X POST 'http:/localhost:4000/team/dbe2f686-a686-4896-864a-4c3924458709/callback' \
-H 'Content-Type: application/json' \
-H 'Authorization: Bearer sk-1234' \
-d '{
"callback_name": "langfuse",
"callback_type": "success",
"callback_vars": {"langfuse_public_key": "pk-lf-xxxx1", "langfuse_secret_key": "sk-xxxxx"}
}'
```
This means for the team where team_id = dbe2f686-a686-4896-864a-4c3924458709, all LLM calls will be logged to langfuse using the public key pk-lf-xxxx1 and the secret key sk-xxxxx
"""
try:
from litellm.proxy.proxy_server import ( from litellm.proxy.proxy_server import (
_duration_in_seconds, _duration_in_seconds,
create_audit_log_for_update, create_audit_log_for_update,
@ -101,6 +125,34 @@ async def add_team_callbacks(
team_metadata["callback_settings"] = team_callback_settings_obj_dict team_metadata["callback_settings"] = team_callback_settings_obj_dict
team_metadata_json = json.dumps(team_metadata) # update team_metadata team_metadata_json = json.dumps(team_metadata) # update team_metadata
await prisma_client.db.litellm_teamtable.update( new_team_row = await prisma_client.db.litellm_teamtable.update(
where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore where={"team_id": team_id}, data={"metadata": team_metadata_json} # type: ignore
) )
return {
"status": "success",
"data": new_team_row,
}
except Exception as e:
verbose_proxy_logger.error(
"litellm.proxy.proxy_server.add_team_callbacks(): Exception occured - {}".format(
str(e)
)
)
verbose_proxy_logger.debug(traceback.format_exc())
if isinstance(e, HTTPException):
raise ProxyException(
message=getattr(e, "detail", f"Internal Server Error({str(e)})"),
type=ProxyErrorTypes.internal_server_error.value,
param=getattr(e, "param", "None"),
code=getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR),
)
elif isinstance(e, ProxyException):
raise e
raise ProxyException(
message="Internal Server Error, " + str(e),
type=ProxyErrorTypes.internal_server_error.value,
param=getattr(e, "param", "None"),
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
)