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 import traceback
from typing import Optional from typing import Optional
from fastapi import APIRouter, Depends, Request, Response from fastapi import APIRouter, Depends, HTTPException, Request, Response
import litellm 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.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 from litellm.types.utils import CredentialItem
router = APIRouter() router = APIRouter()
@ -28,11 +29,33 @@ async def create_credential(
credential: CredentialItem, credential: CredentialItem,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), 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: 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"} return {"success": True, "message": "Credential created successfully"}
except Exception as e: except Exception as e:
return handle_exception_on_proxy(e) verbose_proxy_logger.exception(e)
raise handle_exception_on_proxy(e)
@router.get( @router.get(

View file

@ -29,6 +29,18 @@ model LiteLLM_BudgetTable {
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization 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 // Models on proxy
model LiteLLM_ProxyModelTable { model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid()) model_id String @id @default(uuid())

View file

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

View file

@ -29,6 +29,18 @@ model LiteLLM_BudgetTable {
organization_membership LiteLLM_OrganizationMembership[] // budgets of Users within a Organization 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 // Models on proxy
model LiteLLM_ProxyModelTable { model LiteLLM_ProxyModelTable {
model_id String @id @default(uuid()) model_id String @id @default(uuid())