Spaces:
Running
Running
File size: 2,471 Bytes
09d1b8e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 | import inspect
from collections.abc import Awaitable, Callable
from typing import TypeAlias
import mcp.types
import pydantic
from mcp import ClientSession
from mcp.client.session import ListRootsFnT
from mcp.shared.context import LifespanContextT, RequestContext
RootsList: TypeAlias = list[str] | list[mcp.types.Root] | list[str | mcp.types.Root]
RootsHandler: TypeAlias = (
Callable[[RequestContext[ClientSession, LifespanContextT]], RootsList]
| Callable[[RequestContext[ClientSession, LifespanContextT]], Awaitable[RootsList]]
)
def convert_roots_list(roots: RootsList) -> list[mcp.types.Root]:
roots_list = []
for r in roots:
if isinstance(r, mcp.types.Root):
roots_list.append(r)
elif isinstance(r, pydantic.FileUrl):
roots_list.append(mcp.types.Root(uri=r))
elif isinstance(r, str):
roots_list.append(mcp.types.Root(uri=pydantic.FileUrl(r)))
else:
raise ValueError(f"Invalid root: {r}")
return roots_list
def create_roots_callback(
handler: RootsList | RootsHandler,
) -> ListRootsFnT:
if isinstance(handler, list):
return _create_roots_callback_from_roots(handler)
elif inspect.isfunction(handler):
return _create_roots_callback_from_fn(handler)
else:
raise ValueError(f"Invalid roots handler: {handler}")
def _create_roots_callback_from_roots(
roots: RootsList,
) -> ListRootsFnT:
roots = convert_roots_list(roots)
async def _roots_callback(
context: RequestContext[ClientSession, LifespanContextT],
) -> mcp.types.ListRootsResult:
return mcp.types.ListRootsResult(roots=roots)
return _roots_callback
def _create_roots_callback_from_fn(
fn: Callable[[RequestContext[ClientSession, LifespanContextT]], RootsList]
| Callable[[RequestContext[ClientSession, LifespanContextT]], Awaitable[RootsList]],
) -> ListRootsFnT:
async def _roots_callback(
context: RequestContext[ClientSession, LifespanContextT],
) -> mcp.types.ListRootsResult | mcp.types.ErrorData:
try:
roots = fn(context)
if inspect.isawaitable(roots):
roots = await roots
return mcp.types.ListRootsResult(roots=convert_roots_list(roots))
except Exception as e:
return mcp.types.ErrorData(
code=mcp.types.INTERNAL_ERROR,
message=str(e),
)
return _roots_callback
|