hddlZddlmZmZmZmZmZmZmZddl m Z m Z m Z m Z mZmZmZddlmZdee efdeeeffdZdedee fd Zded edee fd Zded eedee fd ZdededeefdZdededeefdZdededee fdZdededee fdZded edefdZded efdZdZded edee fdZdede dede fdZ dededede fdZ!y)N)AnyDictIterableListOptionalSetUnion)LiteLLM_MCPServerTableLiteLLM_ObjectPermissionTableLiteLLM_TeamTableNewMCPServerRequestSpecialMCPServerNameUpdateMCPServerRequestUserAPIKeyAuth) PrismaClientdatareturncddlm}|j}d|vrt|dd|d<|j||j|d<|j ||j |d<|S)a Helper function to prepare MCP server data for database operations. Handles JSON field serialization for mcp_info and env fields. Args: data: NewMCPServerRequest or UpdateMCPServerRequest object Returns: Dict with properly serialized JSON fields r) safe_dumpsaliasNmcp_infoenv)*litellm.litellm_core_utils.safe_json_dumpsr model_dumpgetattrrr)rr data_dicts e/var/www/Befach/backend/env/lib/python3.12/site-packages/litellm/proxy/_experimental/mcp_server/db.py_prepare_mcp_server_datarswF!Ii$T7D9 ' }} *4== 9 * xx%dhh/ %  prisma_clientc K|jjjd{}|Dcgc]}tdi|j c}S7+cc}ww)z4 Returns all of the mcp servers from the db N)dblitellm_mcpservertable find_manyr r)r mcp_servers mcp_servers rget_all_mcp_serversr(2s^ &((??IIKKKMX >H9!6!6!89 L s (AAA AAA server_idcrK|jjjd|id{}|S7w)z@ Returns the matching mcp server from the db iff exists r)whereN)r#r$ find_unique)r r)r's rget_mcp_serverr.?sI55AAYB     ,757 server_idscvK|jjjdd|iid{}|S7w)zJ Returns the matching mcp servers from the db with the server_ids r)inr+N)r#r$r%)r r0r&s rget_mcp_serversr3OsN55??dJ/@     s .979tokencK|jjjd|iddid{}g}|"|j|jj}|xsgS70w)zH Returns the mcp servers from the db for the verification token r4object_permissionTr,includeN)r#litellm_verificationtokenr-r6r&)r r4verification_token_recordr&s r$get_mcp_servers_by_verificationtokenr;_s88DD$T E  (*K!- % 7 7 C/AAMM  "! /A$A"1A$team_idcK|jjjd|iddid{}g}|"|j|jj}|xsgS70w)z= Returns the mcp servers from the db for the team id r=r6Tr7N)r#litellm_teamtabler-r6r&)r r= team_recordr&s rget_mcp_servers_by_teamrAys00<<7$T =  (*K;#@#@#L!33??  " r<usercKt}g}|jr|t||jd{}|j|tj |vr;|j /t||j d{}|j|t|dkDrt||d{}|S77;7 w)z 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. Nr) setapi_keyr;updaterall_team_serversr=rAlenr3)r rBmcp_server_idsr&token_mcp_serversteam_mcp_serverss rget_all_mcp_servers_for_userrLs #uNK ||"F 4<<#   /0 ! 1 1^ C (%<t||&    ! !"2 3 >Q+M>JJ %   Ks42CB;A CB=2C4B?5C=C?C mcp_server_idc~K|jjjdd|iidddd{}|S7w)z Get all the object permissions records and the associated team and verficiationtoken records that have access to the mcp server r&hasT)teamsverification_tokensr7N)r#litellm_objectpermissiontabler%)r rMobject_permission_recordss r$get_objectpermissions_for_mcp_serverrTs[<<FF}5'+ G   %$ s 2=;=c~K|jjjdd|iid{}|gS|S7 w)zE Get all the virtual keys that have access to the mcp server r&rOr+N)r#r9r%)r r) virtual_keyss rget_virtualkeys_for_mcp_serverrWsV '))CCMM E9- NL  s .=; =c Kyw)z- Remove the mcp server from the team Nr")r r)s rdelete_mcp_server_from_teamrY  c Kyw)z4 Remove the mcp server from the virtual key Nr"r"rr!delete_mcp_server_from_virtualkeyr]rZr[crK|jjjd|id{}|S7w)z Delete the mcp server from the db by server_id Returns the deleted mcp server record if it exists, otherwise None r)r+N)r#r$delete)r r)deleted_servers rdelete_mcp_serverrasF)++BBII  JN  r/ touched_bycK|j"ttj|_t |}||d<||d<|j j j|d{}|S7w)z2 Create a new mcp server record in the db N created_by updated_by)r)r)struuiduuid4rr#r$create)r rrbrnew_mcp_servers rcreate_mcp_serverrks}  ~~TZZ\*).I)Il(Il(++BBII JN  sA-A8/A60A8cKt|}||d<|jjjd|ji|d{}|S7w)z2 Update a new mcp server record in the db rer))r,rN)rr#r$rFr))r rrbrupdated_mcp_servers rupdate_mcp_serverrnsa).I)Il,//FFMMDNN+) N   sAA A A)"rgtypingrrrrrrr litellm.proxy._typesr r r r rrrlitellm.proxy.utilsrrfrr(r.r3r;rArLrTrWrYr]rarkrnr"rrrrs BBB- #%;; < #s(^D   !   ,/  $%   -5c]  !  (+ #Y4*- #Y.!! ! !!H%%03% '(%*,/ " \ c    ,/  $%  ':HK.'=KNr