idZddlmZddlmZmZmZddlmZddl m Z m Z ddl m Z mZddlmZerddlmZed ej& ZGd d Zy) u8 Experimental server session features for server→client task operations. This module provides the server-side equivalent of ExperimentalClientFeatures, allowing the server to send task-augmented requests to the client and poll for results. WARNING: These APIs are experimental and may change without notice. ) AsyncIterator) TYPE_CHECKINGAnyTypeVarN)validate_sampling_tools!validate_tool_use_result_messages)"require_task_augmented_elicitationrequire_task_augmented_sampling)poll_until_terminal) ServerSessionResultT)boundceZdZdZddZdedejfdZdede e de fdZ dede ejfd Z d d d ed ejdedej fdZd ddddddddd deej&dedededzdej(dzdedzdeedzdeeefdzdej0dzdeej2dzdej4dzdej6fdZy)!ExperimentalServerSessionFeaturesu Experimental server session features for server→client task operations. This provides the server-side equivalent of ExperimentalClientFeatures, allowing the server to send task-augmented requests to the client and poll for results. WARNING: These APIs are experimental and may change without notice. Access via session.experimental: result = await session.experimental.elicit_as_task(...) returnNc||_y)N)_session)selfsessions h/root/aria/mcps/aria-brain/venv/lib/python3.12/site-packages/mcp/server/experimental/session_features.py__init__z*ExperimentalServerSessionFeatures.__init__)s  task_idc K|jjtjtjtj |tj d{S7w)z Send tasks/get to the client to get task status. Args: task_id: The task identifier Returns: GetTaskResult containing the task status taskIdparamsN)r send_requesttypes ServerRequestGetTaskRequestGetTaskRequestParams GetTaskResult)rrs rget_taskz*ExperimentalServerSessionFeatures.get_task,sW]]//    4 4E>cg *;7"mm88   ## 88 '(7"//C8   " "    $$++~~g.  ! !   .))'53E3EFFFFsBCD(C:+D(0C>4C<5C>8D(<C>>$D("D%#D() r/ system_promptinclude_context temperaturestop_sequencesmetadatamodel_preferencestools tool_choicemessages max_tokensr?r@rArBrCrDrErFc K|jjr |jjjnd} t| t | | | t ||jj tjtjtj|||||||| | | tj| tjd{} | jj}|j|23d{} 7576|j!|tj"d{7Sw)a5 Send a task-augmented sampling request and poll until complete. The client will create a local task, process the sampling request asynchronously, and return the result when ready. Args: messages: The conversation messages for sampling max_tokens: Maximum tokens in the response ttl: Task time-to-live in milliseconds system_prompt: Optional system prompt include_context: Context inclusion strategy temperature: Sampling temperature stop_sequences: Stop sequences metadata: Additional metadata model_preferences: Model selection preferences tools: Optional list of tools the LLM can use during sampling tool_choice: Optional control over tool usage behavior Returns: The sampling result from the client Raises: McpError: If client doesn't support task-augmented sampling or tools ValueError: If tool_use or tool_result message structure is invalid Nr.) rG maxTokens systemPromptincludeContextrA stopSequencesrCmodelPreferencesrE toolChoicer3r)rr4r5r rrrr r!CreateMessageRequestCreateMessageRequestParamsr8r9r3rr-r*CreateMessageResult)rrGrHr/r?r@rArBrCrDrErFr;r<rr=s rcreate_message_as_taskz8ExperimentalServerSessionFeatures.create_message_as_tasks#RCG--B]B]dmm11>>cg ' 4 UK@)(3"mm88   ** ;;!)",%2'6$/&4!)):##."//C8  "  " "%  * $$++~~g.  ! 1 . .))'53L3LMMMMsBC"E$D%+EDDDED$EEE)rr rN)__name__ __module__ __qualname____doc__rstrr r$r%typer r*rr-ElicitRequestedSchemaintr:r>listSamplingMessageIncludeContextfloatdictrModelPreferencesTool ToolChoicerRrSrrrrs    c  e.A.A    ']    (s}U=P=P/Q. 1G1G441G  1G    1Gp$(7;$(+/*.;?)-/3HNu,,-HN HN  HN Tz HN--4HNT\HNS D(HNsCx.4'HN!11D8HNEJJ$&HN%%,HN  " "HNrr)rWcollections.abcrtypingrrr mcp.typesr mcp.server.validationrr*mcp.shared.experimental.tasks.capabilitiesr r %mcp.shared.experimental.tasks.pollingr mcp.server.sessionr Resultr rrdrrrmsJ*..\F0 )5<< 0ANANr