feat: MCP Servers with CRUD operations

.github/workflows/test-litellm.yml

@@ -37,4 +37,4 @@ jobs:
         cd ..
     - name: Run tests
       run: |
-        poetry run pytest tests/test_litellm -x -vv -n 4 
+        poetry run pytest tests/test_litellm -x -vv -n 4 

docs/my-website/docs/mcp.md

@@ -18,10 +18,19 @@ This allows you to define tools that can be called by any MCP compatible client.
 
 #### How it works
 
+1. Allow proxy admin users to perform create, update, and delete operations on MCP servers stored in the db.
+2. Allows users to view and call tools to the MCP servers they have access to.
+
 LiteLLM exposes the following MCP endpoints:
 
-- `/mcp/tools/list` - List all available tools
-- `/mcp/tools/call` - Call a specific tool with the provided arguments
+- GET `/mcp/enabled` - Returns if MCP is enabled (python>=3.10 requirements are met)
+- GET `/mcp/tools/list` - List all available tools
+- POST `/mcp/tools/call` - Call a specific tool with the provided arguments
+- GET `/v1/mcp/server` - Returns all of the configured mcp servers in the db filtered by requestor's access
+- GET `/v1/mcp/server/{server_id}` - Returns the the specific mcp server in the db given `server_id` filtered by requestor's access
+- PUT  `/v1/mcp/server` - Updates an existing external mcp server.
+- POST `/v1/mcp/server` - Add a new external mcp server.
+- DELETE `/v1/mcp/server/{server_id}` - Deletes the mcp server given `server_id`.
 
 When MCP clients connect to LiteLLM they can follow this workflow:
 

litellm/proxy/_experimental/mcp_server/db.py

@@ -0,0 +1,203 @@
+from typing import Iterable, List, Optional, Set
+
+import uuid
+
+from prisma.models import LiteLLM_MCPServerTable, LiteLLM_ObjectPermissionTable, LiteLLM_TeamTable
+from litellm.proxy._types import NewMCPServerRequest, SpecialMCPServerName, UpdateMCPServerRequest, UserAPIKeyAuth
+from litellm.proxy.utils import PrismaClient
+
+
+async def get_all_mcp_servers(prisma_client: PrismaClient) -> List[LiteLLM_MCPServerTable]:
+    """
+    Returns all of the mcp servers from the db
+    """
+    mcp_servers = await prisma_client.db.litellm_mcpservertable.find_many()
+
+    return mcp_servers
+
+
+async def get_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optional[LiteLLM_MCPServerTable]:
+    """
+    Returns the matching mcp server from the db iff exists
+    """
+    mcp_server: Optional[LiteLLM_MCPServerTable] = await prisma_client.db.litellm_mcpservertable.find_unique(
+        where={
+            "server_id": server_id,
+        }
+    )
+    return mcp_server
+
+
+async def get_mcp_servers(prisma_client: PrismaClient, server_ids: Iterable[str]) -> List[LiteLLM_MCPServerTable]:
+    """
+    Returns the matching mcp servers from the db with the server_ids
+    """
+    mcp_servers: List[LiteLLM_MCPServerTable] = await prisma_client.db.litellm_mcpservertable.find_many(
+        where={
+            "server_id": {"in": server_ids},
+        }
+    )
+    return mcp_servers
+
+
+async def get_mcp_servers_by_verificationtoken(prisma_client: PrismaClient, token: str) -> List[str]:
+    """
+    Returns the mcp servers from the db for the verification token
+    """
+    verification_token_record: LiteLLM_TeamTable = await prisma_client.db.litellm_verificationtoken.find_unique(
+        where={
+            "token": token,
+        },
+        include={
+            "object_permission": True,
+        },
+    )
+
+    mcp_servers = []
+    if verification_token_record is not None and verification_token_record.object_permission is not None:
+        mcp_servers = verification_token_record.object_permission.mcp_servers
+    return mcp_servers
+
+
+async def get_mcp_servers_by_team(prisma_client: PrismaClient, team_id: str) -> List[str]:
+    """
+    Returns the mcp servers from the db for the team id
+    """
+    team_record: LiteLLM_TeamTable = await prisma_client.db.litellm_teamtable.find_unique(
+        where={
+            "team_id": team_id,
+        },
+        include={
+            "object_permission": True,
+        },
+    )
+
+    mcp_servers = []
+    if team_record is not None and team_record.object_permission is not None:
+        mcp_servers = team_record.object_permission.mcp_servers
+    return mcp_servers
+
+
+async def get_all_mcp_servers_for_user(
+    prisma_client: PrismaClient,
+    user: UserAPIKeyAuth,
+) -> List[LiteLLM_MCPServerTable]:
+    """
+    Get all the mcp servers filtered by the given user has access to.
+
+    Following Least-Privilege Principle - the requestor should only be able to see the mcp servers that they have access to.
+    """
+
+    mcp_server_ids: Set[str] = set()
+    mcp_servers = []
+
+    # Get the mcp servers for the key
+    if user.api_key:
+        token_mcp_servers = await get_mcp_servers_by_verificationtoken(prisma_client, user.api_key)
+        mcp_server_ids.update(token_mcp_servers)
+
+        # check for special team membership
+        if SpecialMCPServerName.all_team_servers in mcp_server_ids and user.team_id is not None:
+            team_mcp_servers = await get_mcp_servers_by_team(prisma_client, user.team_id)
+            mcp_server_ids.update(team_mcp_servers)
+
+    if len(mcp_server_ids) > 0:
+        mcp_servers = await get_mcp_servers(prisma_client, mcp_server_ids)
+
+    return mcp_servers
+
+
+async def get_objectpermissions_for_mcp_server(
+    prisma_client: PrismaClient, mcp_server_id: str
+) -> List[LiteLLM_ObjectPermissionTable]:
+    """
+    Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server
+    """
+    object_permission_records = await prisma_client.db.litellm_objectpermissiontable.find_many(
+        where={
+            "mcp_servers": {"has": mcp_server_id},
+        },
+        include={
+            "teams": True,
+            "verification_tokens": True,
+        },
+    )
+
+    return object_permission_records
+
+
+async def get_virtualkeys_for_mcp_server(prisma_client: PrismaClient, server_id: str) -> List:
+    """
+    Get all the virtual keys that have access to the mcp server
+    """
+    virtual_keys = await prisma_client.db.litellm_verificationtoken.find_many(
+        where={
+            "mcp_servers": {"has": server_id},
+        },
+    )
+
+    if virtual_keys is None:
+        return []
+    return virtual_keys
+
+
+async def delete_mcp_server_from_team(prisma_client: PrismaClient, server_id: str):
+    """
+    Remove the mcp server from the team
+    """
+    pass
+
+
+async def delete_mcp_server_from_virtualkey():
+    """
+    Remove the mcp server from the virtual key
+    """
+    pass
+
+
+async def delete_mcp_server(prisma_client: PrismaClient, server_id: str) -> Optional[LiteLLM_MCPServerTable]:
+    """
+    Delete the mcp server from the db by server_id
+
+    Returns the deleted mcp server record if it exists, otherwise None
+    """
+    deleted_server = await prisma_client.db.litellm_mcpservertable.delete(
+        where={
+            "server_id": server_id,
+        },
+    )
+    return deleted_server
+
+
+async def create_mcp_server(prisma_client: PrismaClient, data: NewMCPServerRequest, touched_by: str) -> LiteLLM_MCPServerTable:
+    """
+    Create a new mcp server record in the db
+    """
+    if data.server_id is None:
+        data.server_id = str(uuid.uuid4())
+
+    mcp_server_record = await prisma_client.db.litellm_mcpservertable.create(
+        data={
+            **data.model_dump(),
+            'created_by': touched_by,
+            'updated_by': touched_by,
+        }
+    )
+    return mcp_server_record
+
+
+async def update_mcp_server(prisma_client: PrismaClient, data: UpdateMCPServerRequest, touched_by: str) -> LiteLLM_MCPServerTable:
+    """
+    Update a new mcp server record in the db
+    """
+    mcp_server_record = await prisma_client.db.litellm_mcpservertable.update(
+        where={
+            'server_id': data.server_id,
+        },
+        data={
+            **data.model_dump(),
+            'created_by': touched_by,
+            'updated_by': touched_by,
+        }
+    )
+    return mcp_server_record