1616from typing import Any , TypeAlias
1717
1818import anyio
19- from pydantic import BaseModel
19+ from pydantic import BaseModel , AnyUrl
2020from typing_extensions import Self
2121
2222import mcp
@@ -100,6 +100,7 @@ class _ComponentNames(BaseModel):
100100 # Client-server connection management.
101101 _sessions : dict [mcp .ClientSession , _ComponentNames ]
102102 _tool_to_session : dict [str , mcp .ClientSession ]
103+ _resource_to_session : dict [str , mcp .ClientSession ]
103104 _exit_stack : contextlib .AsyncExitStack
104105 _session_exit_stacks : dict [mcp .ClientSession , contextlib .AsyncExitStack ]
105106
@@ -116,20 +117,16 @@ def __init__(
116117 ) -> None :
117118 """Initializes the MCP client."""
118119
119- self ._tools = {}
120- self ._resources = {}
120+ self ._exit_stack = exit_stack or contextlib .AsyncExitStack ()
121+ self ._owns_exit_stack = exit_stack is None
122+ self ._session_exit_stacks = {}
123+ self ._component_name_hook = component_name_hook
121124 self ._prompts = {}
122-
125+ self ._resources = {}
126+ self ._tools = {}
123127 self ._sessions = {}
124128 self ._tool_to_session = {}
125- if exit_stack is None :
126- self ._exit_stack = contextlib .AsyncExitStack ()
127- self ._owns_exit_stack = True
128- else :
129- self ._exit_stack = exit_stack
130- self ._owns_exit_stack = False
131- self ._session_exit_stacks = {}
132- self ._component_name_hook = component_name_hook
129+ self ._resource_to_session = {} # New mapping
133130
134131 async def __aenter__ (self ) -> Self :
135132 # Enter the exit stack only if we created it ourselves
@@ -174,6 +171,16 @@ def tools(self) -> dict[str, types.Tool]:
174171 """Returns the tools as a dictionary of names to tools."""
175172 return self ._tools
176173
174+ @property
175+ def resource_templates (self ) -> list [types .ResourceTemplate ]:
176+ """Return all unique resource templates from the resources."""
177+ templates : list [types .ResourceTemplate ] = []
178+ for r in self ._resources .values ():
179+ t = getattr (r , "template" , None )
180+ if t is not None and t not in templates :
181+ templates .append (t )
182+ return templates
183+
177184 async def call_tool (self , name : str , args : dict [str , Any ]) -> types .CallToolResult :
178185 """Executes a tool given its name and arguments."""
179186 session = self ._tool_to_session [name ]
@@ -296,8 +303,8 @@ async def _aggregate_components(
296303 resources_temp : dict [str , types .Resource ] = {}
297304 tools_temp : dict [str , types .Tool ] = {}
298305 tool_to_session_temp : dict [str , mcp .ClientSession ] = {}
306+ resource_to_session_temp : dict [str , mcp .ClientSession ] = {}
299307
300- # Query the server for its prompts and aggregate to list.
301308 try :
302309 prompts = (await session .list_prompts ()).prompts
303310 for prompt in prompts :
@@ -314,6 +321,7 @@ async def _aggregate_components(
314321 name = self ._component_name (resource .name , server_info )
315322 resources_temp [name ] = resource
316323 component_names .resources .add (name )
324+ resource_to_session_temp [name ] = session
317325 except McpError as err :
318326 logging .warning (f"Could not fetch resources: { err } " )
319327
@@ -365,8 +373,20 @@ async def _aggregate_components(
365373 self ._resources .update (resources_temp )
366374 self ._tools .update (tools_temp )
367375 self ._tool_to_session .update (tool_to_session_temp )
376+ self ._resource_to_session .update (resource_to_session_temp )
368377
369378 def _component_name (self , name : str , server_info : types .Implementation ) -> str :
370379 if self ._component_name_hook :
371380 return self ._component_name_hook (name , server_info )
372381 return name
382+
383+ async def read_resource (self , uri : AnyUrl ) -> types .ReadResourceResult :
384+ """Read a resource from the appropriate session based on the URI."""
385+ print (self ._resources )
386+ print (self ._resource_to_session )
387+ for name , resource in self ._resources .items ():
388+ if resource .uri == uri :
389+ session = self ._resource_to_session .get (name )
390+ if session :
391+ return await session .read_resource (uri )
392+ raise ValueError (f"Resource not found: { uri } " )
0 commit comments