Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/_experimental/mcp_server/toolset_db.py: 83%

1import json 

2from collections.abc import Mapping, Sequence 

3from datetime import datetime 

4from typing import Final, Protocol 

5 

6from typing_extensions import NotRequired, ReadOnly, TypedDict 

7 

8from litellm._logging import verbose_proxy_logger 

9from litellm._uuid import uuid 

10from litellm.proxy.utils import PrismaClient 

11from litellm.repositories.table_repositories import MCPToolsetRepository 

12from litellm.types.mcp_server.mcp_toolset import ( 

13 MCPToolset, 

14 MCPToolsetTool, 

15 NewMCPToolsetRequest, 

16 UpdateMCPToolsetRequest, 

17) 

18 

19 

20class MCPToolsetFields(TypedDict): 

21 """The ``MCPToolset`` constructor keywords a toolset row expands into.""" 

22 

23 toolset_id: ReadOnly[str] 

24 toolset_name: ReadOnly[str] 

25 description: NotRequired[ReadOnly[str | None]] 

26 tools: NotRequired[ReadOnly[list[MCPToolsetTool]]] 

27 created_at: NotRequired[ReadOnly[datetime | None]] 

28 created_by: NotRequired[ReadOnly[str | None]] 

29 updated_at: NotRequired[ReadOnly[datetime | None]] 

30 updated_by: NotRequired[ReadOnly[str | None]] 

31 

32 

33class MCPToolsetRowData(TypedDict): 

34 """A toolset table row, whose ``tools`` column is stored as JSON.""" 

35 

36 toolset_id: ReadOnly[str] 

37 toolset_name: ReadOnly[str] 

38 description: NotRequired[ReadOnly[str | None]] 

39 tools: NotRequired[ReadOnly[str | list[MCPToolsetTool]]] 

40 created_at: NotRequired[ReadOnly[datetime | None]] 

41 created_by: NotRequired[ReadOnly[str | None]] 

42 updated_at: NotRequired[ReadOnly[datetime | None]] 

43 updated_by: NotRequired[ReadOnly[str | None]] 

44 

45 

46class MCPToolsetRow(Protocol): 

47 """A row of the toolset table, as the prisma client returns it.""" 

48 

49 def model_dump(self) -> MCPToolsetRowData: ... 49 ↛ exitline 49 didn't return from function 'model_dump' because

50 

51 

52class MCPToolsetTable(Protocol): 

53 """The prisma table actions this module runs against the toolset table.""" 

54 

55 async def create(self, data: Mapping[str, object]) -> MCPToolsetRow: ... 55 ↛ exitline 55 didn't return from function 'create' because

56 

57 async def find_unique(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... 57 ↛ exitline 57 didn't return from function 'find_unique' because

58 

59 async def find_first(self, where: Mapping[str, object]) -> MCPToolsetRow | None: ... 59 ↛ exitline 59 didn't return from function 'find_first' because

60 

61 async def find_many(self, where: Mapping[str, object]) -> Sequence[MCPToolsetRow]: ... 61 ↛ exitline 61 didn't return from function 'find_many' because

62 

63 async def update(self, where: Mapping[str, object], data: Mapping[str, object]) -> MCPToolsetRow: ... 63 ↛ exitline 63 didn't return from function 'update' because

64 

65 async def delete(self, where: Mapping[str, object]) -> MCPToolsetRow: ... 65 ↛ exitline 65 didn't return from function 'delete' because

66 

67 

68def _toolset_table(prisma_client: PrismaClient) -> MCPToolsetTable: 

69 """The toolset table actions of the prisma client.""" 

70 return MCPToolsetRepository(prisma_client).table 

71 

72 

73def _toolset_from_row(row: MCPToolsetRow) -> MCPToolset: 

74 data: Final = row.model_dump() 

75 tools: Final = data.get("tools") or [] 

76 resolved: Final[MCPToolsetFields] = { 

77 **data, 

78 "tools": json.loads(tools) if isinstance(tools, str) else tools, 

79 } 

80 return MCPToolset(**resolved) 

81 

82 

83async def create_mcp_toolset( 

84 prisma_client: PrismaClient, 

85 data: NewMCPToolsetRequest, 

86 touched_by: str, 

87) -> MCPToolset: 

88 data_dict: Final = data.model_dump(exclude_none=True) 

89 data_dict["toolset_id"] = str(uuid.uuid4()) 

90 data_dict["tools"] = json.dumps(data_dict.get("tools", [])) 

91 data_dict["created_by"] = touched_by 

92 data_dict["updated_by"] = touched_by 

93 row: Final = await _toolset_table(prisma_client).create(data=data_dict) 

94 return _toolset_from_row(row) 

95 

96 

97async def get_mcp_toolset( 

98 prisma_client: PrismaClient, 

99 toolset_id: str, 

100) -> MCPToolset | None: 

101 row: Final = await _toolset_table(prisma_client).find_unique(where={"toolset_id": toolset_id}) 

102 if row is None: 

103 return None 

104 return _toolset_from_row(row) 

105 

106 

107async def list_mcp_toolsets( 

108 prisma_client: PrismaClient, 

109 toolset_ids: Sequence[str] | None = None, 

110) -> Sequence[MCPToolset]: 

111 try: 

112 where: Final[Mapping[str, object]] = {} if toolset_ids is None else {"toolset_id": {"in": toolset_ids}} 

113 rows: Final = await _toolset_table(prisma_client).find_many(where=where) 

114 return [_toolset_from_row(r) for r in rows] 

115 except Exception as e: 

116 verbose_proxy_logger.warning("litellm.proxy._experimental.mcp_server.toolset_db::list_mcp_toolsets - %s", e) 

117 return [] 

118 

119 

120async def get_mcp_toolset_by_name( 

121 prisma_client: PrismaClient, 

122 toolset_name: str, 

123) -> MCPToolset | None: 

124 row: Final = await _toolset_table(prisma_client).find_first(where={"toolset_name": toolset_name}) 

125 if row is None: 

126 return None 

127 return _toolset_from_row(row) 

128 

129 

130async def update_mcp_toolset( 

131 prisma_client: PrismaClient, 

132 data: UpdateMCPToolsetRequest, 

133 touched_by: str, 

134) -> MCPToolset | None: 

135 """A partial update: absent keeps, null clears. A toolset always has a name and a 

136 tool list, so a null ``toolset_name`` or ``tools`` is a no-op rather than a clear; 

137 emptying the tool selection is an explicit ``[]``, which cannot be mistaken for a 

138 caller that left the field out.""" 

139 data_dict: Final = dict( # mutable-ok: Prisma requires a plain dict for JSON query serialization 

140 ( 

141 (field, json.dumps(value) if field == "tools" else value) 

142 for field, value in data.model_dump(exclude_unset=True).items() 

143 if field != "toolset_id" and (field not in ("toolset_name", "tools") or value is not None) 

144 ), 

145 updated_by=touched_by, 

146 ) 

147 try: 

148 row: Final = await _toolset_table(prisma_client).update( 

149 where={"toolset_id": data.toolset_id}, 

150 data=data_dict, 

151 ) 

152 except Exception as e: 

153 from prisma.errors import RecordNotFoundError 

154 

155 if isinstance(e, RecordNotFoundError): 155 ↛ 156line 155 didn't jump to line 156 because the condition on line 155 was never true

156 return None 

157 raise 

158 return _toolset_from_row(row) 

159 

160 

161async def delete_mcp_toolset( 

162 prisma_client: PrismaClient, 

163 toolset_id: str, 

164) -> MCPToolset | None: 

165 try: 

166 row: Final = await _toolset_table(prisma_client).delete(where={"toolset_id": toolset_id}) 

167 except Exception as e: 

168 from prisma.errors import RecordNotFoundError 

169 

170 if isinstance(e, RecordNotFoundError): 

171 return None 

172 raise 

173 return _toolset_from_row(row)