NotificationsYou must be signed in to change notification settings
Fork1k
Star11k

Commitd10f036

authored

Revert "Use contextvars for tracking the MCP sampling model" (#2132)

1 parent09e0821 commitd10f036Copy full SHA for d10f036

File tree

3 files changed

-26

lines changed

pydantic_ai_slim/pydantic_ai
- agent.py
- mcp.py
tests
- test_examples.py

3 files changed

-26

lines changed

`‎pydantic_ai_slim/pydantic_ai/agent.py`

Lines changed: 1 addition & 1 deletion

Original file line number	Diff line number	Diff line change
`@@ -1740,7 +1740,7 @@ async def run_mcp_servers(`
`1740`	`1740`	`try:`
`1741`	`1741`	`formcp_serverinself._mcp_servers:`
`1742`	`1742`	`ifsampling_modelisnotNone:# pragma: no branch`
`1743`		`-exit_stack.enter_context(mcp_server.override_sampling_model(sampling_model))`
	`1743`	`+mcp_server.sampling_model=sampling_model`
`1744`	`1744`	`awaitexit_stack.enter_async_context(mcp_server)`
`1745`	`1745`	`yield`
`1746`	`1746`	`finally:`

`‎pydantic_ai_slim/pydantic_ai/mcp.py`

Lines changed: 5 additions & 23 deletions

Original file line number	Diff line number	Diff line change
`@@ -3,9 +3,8 @@`
`3`	`3`	`importbase64`
`4`	`4`	`importfunctools`
`5`	`5`	`fromabcimportABC,abstractmethod`
`6`		`-fromcollections.abcimportAsyncIterator,Awaitable,Iterator,Sequence`
`7`		`-fromcontextlibimportAbstractAsyncContextManager,AsyncExitStack,asynccontextmanager,contextmanager`
`8`		`-fromcontextvarsimportContextVar`
	`6`	`+fromcollections.abcimportAsyncIterator,Awaitable,Sequence`
	`7`	`+fromcontextlibimportAbstractAsyncContextManager,AsyncExitStack,asynccontextmanager`
`9`	`8`	`fromdataclassesimportdataclass`
`10`	`9`	`frompathlibimportPath`
`11`	`10`	`fromtypesimportTracebackType`
`@@ -61,22 +60,6 @@ class MCPServer(ABC):`
`61`	`60`	`_exit_stack:AsyncExitStack`
`62`	`61`	`sampling_model:models.Model\|None=None`
`63`	`62`
`64`		`-def__post_init__(self):`
`65`		`-self._override_sampling_model:ContextVar[models.Model\|None]=ContextVar(`
`66`		`-'_override_sampling_model',default=None`
`67`		`- )`
`68`		`-`
`69`		`-@contextmanager`
`70`		`-defoverride_sampling_model(`
`71`		`-self,`
`72`		`-model:models.Model,`
`73`		`- )->Iterator[None]:`
`74`		`-token=self._override_sampling_model.set(model)`
`75`		`-try:`
`76`		`-yield`
`77`		`-finally:`
`78`		`-self._override_sampling_model.reset(token)`
`79`		`-`
`80`	`63`	`@abstractmethod`
`81`	`64`	`@asynccontextmanager`
`82`	`65`	`asyncdefclient_streams(`
`@@ -201,8 +184,7 @@ async def _sampling_callback(`
`201`	`184`	`self,context:RequestContext[ClientSession,Any],params:mcp_types.CreateMessageRequestParams`
`202`	`185`	`)->mcp_types.CreateMessageResult\|mcp_types.ErrorData:`
`203`	`186`	`"""MCP sampling callback."""`
`204`		`-sampling_model=self._override_sampling_model.get()orself.sampling_model`
`205`		`-ifsampling_modelisNone:`
	`187`	`+ifself.sampling_modelisNone:`
`206`	`188`	`raiseValueError('Sampling model is not set')# pragma: no cover`
`207`	`189`
`208`	`190`	`pai_messages=_mcp.map_from_mcp_params(params)`
`@@ -214,15 +196,15 @@ async def _sampling_callback(`
`214`	`196`	`ifstop_sequences:=params.stopSequences:# pragma: no branch`
`215`	`197`	`model_settings['stop_sequences']=stop_sequences`
`216`	`198`
`217`		`-model_response=awaitsampling_model.request(`
	`199`	`+model_response=awaitself.sampling_model.request(`
`218`	`200`	`pai_messages,`
`219`	`201`	`model_settings,`
`220`	`202`	`models.ModelRequestParameters(),`
`221`	`203`	`)`
`222`	`204`	`returnmcp_types.CreateMessageResult(`
`223`	`205`	`role='assistant',`
`224`	`206`	`content=_mcp.map_from_model_response(model_response),`
`225`		`-model=sampling_model.model_name,`
	`207`	`+model=self.sampling_model.model_name,`
`226`	`208`	`)`
`227`	`209`
`228`	`210`	`def_map_tool_result_part(`

`‎tests/test_examples.py`

Lines changed: 0 additions & 2 deletions

Original file line number	Diff line number	Diff line change
`@@ -6,7 +6,6 @@`
`6`	`6`	`importshutil`
`7`	`7`	`importsys`
`8`	`8`	`fromcollections.abcimportAsyncIterator,Iterable,Sequence`
`9`		`-fromcontextlibimportnullcontext`
`10`	`9`	`fromdataclassesimportdataclass`
`11`	`10`	`frominspectimportFrameInfo`
`12`	`11`	`fromioimportStringIO`
`@@ -259,7 +258,6 @@ def rich_prompt_ask(prompt: str, _args: Any, *_kwargs: Any) -> str:`
`259`	`258`
`260`	`259`	`classMockMCPServer:`
`261`	`260`	`is_running=True`
`262`		`-override_sampling_model=nullcontext`
`263`	`261`
`264`	`262`	`asyncdef__aenter__(self)->MockMCPServer:`
`265`	`263`	`returnself`

0 commit comments

Comments

(0)

Movatterモバイル変換

Navigation Menu

Search code, repositories, users, issues, pull requests...

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Commitd10f036

File tree

3 files changed

3 files changed

`‎pydantic_ai_slim/pydantic_ai/agent.py`

`‎pydantic_ai_slim/pydantic_ai/mcp.py`

`‎tests/test_examples.py`

0 commit comments