feat(endpoints.py): support writing credentials to db

This commit is contained in:
Krrish Dholakia 2025-03-10 18:27:43 -07:00
parent f1cdc26967
commit a962a97fcb
4 changed files with 54 additions and 9 deletions

View file

@ -6,12 +6,13 @@ import asyncio
import traceback
from typing import Optional
from fastapi import APIRouter, Depends, Request, Response
from fastapi import APIRouter, Depends, HTTPException, Request, Response
import litellm
from litellm.proxy._types import UserAPIKeyAuth
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
from litellm.proxy.utils import handle_exception_on_proxy
from litellm.proxy.utils import handle_exception_on_proxy, jsonify_object
from litellm.types.utils import CredentialItem
router = APIRouter()
@ -28,11 +29,33 @@ async def create_credential(
credential: CredentialItem,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
):
"""
Stores credential in DB.
Reloads credentials in memory.
"""
from litellm.proxy.proxy_server import prisma_client
try:
litellm.credential_list.append(credential)
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
credentials_dict = credential.model_dump()
credentials_dict_jsonified = jsonify_object(credentials_dict)
await prisma_client.db.litellm_credentialstable.create(
data={
**credentials_dict_jsonified,
"created_by": user_api_key_dict.user_id,
"updated_by": user_api_key_dict.user_id,
}
)
return {"success": True, "message": "Credential created successfully"}
except Exception as e:
return handle_exception_on_proxy(e)
verbose_proxy_logger.exception(e)
raise handle_exception_on_proxy(e)
@router.get(

View file

@ -29,6 +29,18 @@ model LiteLLM_BudgetTable {
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
}
// Models on proxy
model LiteLLM_CredentialsTable {
credential_id String @id @default(uuid())
credential_name String @unique
credential_values Json
credential_info Json?
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
updated_by String
}
// Models on proxy
model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid())

View file

@ -18,13 +18,11 @@ from openai.types.moderation import (
CategoryScores,
)
from openai.types.moderation_create_response import Moderation, ModerationCreateResponse
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr, Secret
from pydantic import BaseModel, ConfigDict, Field, PrivateAttr
from typing_extensions import Callable, Dict, Required, TypedDict, override
import litellm
SecretDict = Secret[dict]
from ..litellm_core_utils.core_helpers import map_finish_reason
from .guardrails import GuardrailEventHooks
from .llms.openai import (
@ -2018,5 +2016,5 @@ class RawRequestTypedDict(TypedDict, total=False):
class CredentialItem(BaseModel):
credential_name: str
credential_values: SecretDict
credential_values: dict
credential_info: dict

View file

@ -29,6 +29,18 @@ model LiteLLM_BudgetTable {
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization
}
// Models on proxy
model LiteLLM_CredentialsTable {
credential_id String @id @default(uuid())
credential_name String @unique
credential_values Json
credential_info Json?
created_at DateTime @default(now()) @map("created_at")
created_by String
updated_at DateTime @default(now()) @updatedAt @map("updated_at")
updated_by String
}
// Models on proxy
model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid())