pkgs/dart_mcp/lib/src/api/sampling.dart

// Copyright (c) 2025, the Dart project authors.  Please see the AUTHORS file
// for details. All rights reserved. Use of this source code is governed by a
// BSD-style license that can be found in the LICENSE file.

part of 'api.dart';

/// The types of context in which should be included in a prompt.
///
/// The schema deprecates [thisServer] and [allServers], asking a server to omit
/// the field or send [none] unless the client declares the `sampling.context`
/// capability.
enum IncludeContext {
  none,
  thisServer,
  allServers;

  @Deprecated('Use `IncludeContext.thisServer` instead.')
  static const thisService = thisServer;
}

/// A request from the server to sample an LLM via the client.
///
/// The client has full discretion over which model to select. The client should
/// also inform the user before beginning sampling, to allow them to inspect
/// the request (human in the loop) and decide whether to approve it.
extension type CreateMessageRequest.fromMap(Map<String, Object?> _value)
    implements Request {
  static const methodName = 'sampling/createMessage';

  factory CreateMessageRequest({
    required List<SamplingMessage> messages,
    ModelPreferences? modelPreferences,
    String? systemPrompt,
    IncludeContext? includeContext,
    num? temperature,
    required int maxTokens,
    List<String>? stopSequences,
    ToolChoice? toolChoice,
    List<Tool>? tools,
    Map<String, Object?>? metadata,
    MetaWithProgressToken? meta,
  }) => CreateMessageRequest.fromMap({
    Keys.messages: messages,
    if (modelPreferences != null) Keys.modelPreferences: modelPreferences,
    if (systemPrompt != null) Keys.systemPrompt: systemPrompt,
    if (includeContext != null) Keys.includeContext: includeContext.name,
    if (temperature != null) Keys.temperature: temperature,
    Keys.maxTokens: maxTokens,
    if (stopSequences != null) Keys.stopSequences: stopSequences,
    if (toolChoice != null) Keys.toolChoice: toolChoice,
    if (tools != null) Keys.tools: tools,
    if (metadata != null) Keys.metadata: metadata,
    if (meta != null) Keys.meta: meta,
  });

  /// The messages to send to the LLM.
  List<SamplingMessage> get messages {
    final messages = _value[Keys.messages] as List?;
    if (messages == null) {
      throw ArgumentError('Missing messages field in $CreateMessageRequest.');
    }
    return messages.cast<SamplingMessage>();
  }

  /// The server's preferences for which model to select.
  ///
  /// The client MAY ignore these preferences.
  ModelPreferences? get modelPreferences =>
      _value[Keys.modelPreferences] as ModelPreferences?;

  /// An optional system prompt the server wants to use for sampling.
  ///
  /// The client MAY modify or omit this prompt.
  String? get systemPrompt => _value[Keys.systemPrompt] as String?;

  /// A request to include context from one or more MCP servers (including
  /// the caller), to be attached to the prompt.
  ///
  /// The client MAY ignore this request.
  IncludeContext? get includeContext {
    var includeContext = _value[Keys.includeContext] as String?;
    if (includeContext == null) return null;
    // This package wrote `thisService` for the schema's `thisServer` up to
    // 0.5.2. Treat it as an alias in case a server is still on that version.
    if (includeContext == 'thisService') {
      includeContext = IncludeContext.thisServer.name;
    }
    return IncludeContext.values.firstWhere(
      (value) => value.name == includeContext,
    );
  }

  /// The temperature to use for sampling.
  double? get temperature => (_value[Keys.temperature] as num?)?.toDouble();

  /// The maximum number of tokens to sample, as requested by the server.
  ///
  /// The client MAY choose to sample fewer tokens than requested.
  int get maxTokens {
    final maxTokens = _value[Keys.maxTokens] as int?;
    if (maxTokens == null) {
      throw ArgumentError(
        'Missing ${Keys.maxTokens} field in $CreateMessageRequest.',
      );
    }
    return maxTokens;
  }

  /// Note: This has no documentation in the specification or schema.
  List<String>? get stopSequences =>
      (_value[Keys.stopSequences] as List?)?.cast<String>();

  /// Controls how the model uses tools (if available).
  ToolChoice? get toolChoice => _value[Keys.toolChoice] as ToolChoice?;

  /// Tools the model may call during this request.
  List<Tool>? get tools => (_value[Keys.tools] as List?)?.cast<Tool>();

  /// Optional metadata to pass through to the LLM provider.
  ///
  /// The format of this metadata is provider-specific.
  Map<String, Object?>? get metadata =>
      (_value[Keys.metadata] as Map?)?.cast<String, Object?>();
}

/// The client's response to a sampling/create_message request from the
/// server.
///
/// The client should inform the user before returning the sampled message, to
/// allow them to inspect the response (human in the loop) and decide whether
/// to allow the server to see it.
extension type CreateMessageResult.fromMap(Map<String, Object?> _value)
    implements Result, SamplingMessage {
  factory CreateMessageResult({
    required Role role,
    required List<SamplingMessageContentBlock> content,
    required String model,
    String? stopReason,
    Meta? meta,
  }) => CreateMessageResult.fromMap({
    Keys.role: role.name,
    Keys.content: _encodeContent(content),
    Keys.model: model,
    if (stopReason != null) Keys.stopReason: stopReason,
    if (meta != null) Keys.meta: meta,
  });

  /// The name of the model that generated the message.
  String get model => _value[Keys.model] as String;

  /// The reason why sampling stopped, if known.
  ///
  /// Known reasons are "endTurn", "stopSequence", "maxTokens", "toolUse", or
  /// any other reason.
  String? get stopReason => _value[Keys.stopReason] as String?;

  /// The JSON representation of this object.
  Map<String, Object?> toJson() => _value;
}

/// Describes a message issued to or received from an LLM API.
extension type SamplingMessage.fromMap(Map<String, Object?> _value) {
  factory SamplingMessage({
    required Role role,
    required List<SamplingMessageContentBlock> content,
  }) => SamplingMessage.fromMap({
    Keys.role: role.name,
    Keys.content: _encodeContent(content),
  });

  /// The role of the message.
  Role get role =>
      Role.values.firstWhere((value) => value.name == _value[Keys.role]);

  /// The content of the message.
  ///
  /// The schema allows one block or a list of them under `content`, and both
  /// read as a list here. One block comes back as a single-element list.
  List<SamplingMessageContentBlock> get content {
    final content = _value[Keys.content];
    if (content == null) {
      throw ArgumentError('Missing ${Keys.content} field in $SamplingMessage.');
    }
    if (content is List) {
      return content.cast<SamplingMessageContentBlock>();
    }
    return [content as SamplingMessageContentBlock];
  }
}

/// Writes [content] under `content` as one of the schema's two shapes.
///
/// A single block goes on the wire as that block. Two or more go as a list.
Object? _encodeContent(List<SamplingMessageContentBlock> content) {
  if (content.length == 1) return content.single;
  return content;
}

/// One block of a [SamplingMessage]'s content, sent to or received from an
/// LLM.
///
/// Could be either [TextContent], [ImageContent], [AudioContent],
/// [ToolUseContent] or [ToolResultContent].
///
/// Switch on the [type] before casting to the more specific types. The two
/// arms a `tools/call` result cannot carry have their own checks,
/// [isToolUse] and [isToolResult]. A [TextContent], [ImageContent] or
/// [AudioContent] read as a [Content] keeps the checks declared there.
///
/// This does not implement [Content]. A plain `tools/call` result never
/// carries [ToolUseContent] or [ToolResultContent]. Keeping sampling content
/// on its own type stops the two from being swapped by accident.
///
/// Doing `is` checks does not work because these are just extension types,
/// they all have the same runtime type (`Map<String, Object?>`).
extension type SamplingMessageContentBlock._(Map<String, Object?> _value) {
  factory SamplingMessageContentBlock.fromMap(Map<String, Object?> value) {
    assert(value.containsKey(Keys.type));
    return SamplingMessageContentBlock._(value);
  }

  /// Alias for [TextContent.new].
  static const text = TextContent.new;

  /// Alias for [ImageContent.new].
  static const image = ImageContent.new;

  /// Alias for [AudioContent.new].
  static const audio = AudioContent.new;

  /// Alias for [ToolUseContent.new].
  static const toolUse = ToolUseContent.new;

  /// Alias for [ToolResultContent.new].
  static const toolResult = ToolResultContent.new;

  /// Whether or not this is a [ToolUseContent].
  bool get isToolUse => _value[Keys.type] == ToolUseContent.expectedType;

  /// Whether or not this is a [ToolResultContent].
  bool get isToolResult => _value[Keys.type] == ToolResultContent.expectedType;

  /// The type of content.
  ///
  /// Switch on this to handle each case (see the static `expectedType`
  /// getters). [isToolUse] and [isToolResult] cover the two arms a
  /// `tools/call` result cannot carry.
  String get type => _value[Keys.type] as String;
}

/// A request from the assistant to call a tool.
///
/// From the 2025-11-25 revision.
extension type ToolUseContent.fromMap(Map<String, Object?> _value)
    implements SamplingMessageContentBlock, WithMetadata {
  static const expectedType = 'tool_use';

  factory ToolUseContent({
    required String id,
    required String name,
    required Map<String, Object?> input,
    Meta? meta,
  }) => ToolUseContent.fromMap({
    Keys.id: id,
    Keys.input: input,
    Keys.name: name,
    Keys.type: expectedType,
    if (meta != null) Keys.meta: meta,
  });

  /// The content type, always [expectedType].
  String get type {
    final type = _value[Keys.type] as String;
    assert(type == expectedType);
    return type;
  }

  /// The unique identifier for this tool use.
  String get id => _value[Keys.id] as String;

  /// The name of the tool to call.
  String get name => _value[Keys.name] as String;

  /// The arguments to pass to the tool.
  Map<String, Object?> get input =>
      (_value[Keys.input] as Map).cast<String, Object?>();
}

/// The result of a tool use, provided by the user back to the assistant.
///
/// From the 2025-11-25 revision.
extension type ToolResultContent.fromMap(Map<String, Object?> _value)
    implements SamplingMessageContentBlock, WithMetadata {
  static const expectedType = 'tool_result';

  factory ToolResultContent({
    required List<Content> content,
    required String toolUseId,
    Map<String, Object?>? structuredContent,
    bool? isError,
    Meta? meta,
  }) => ToolResultContent.fromMap({
    Keys.content: content,
    Keys.toolUseId: toolUseId,
    Keys.type: expectedType,
    if (structuredContent != null) Keys.structuredContent: structuredContent,
    if (isError != null) Keys.isError: isError,
    if (meta != null) Keys.meta: meta,
  });

  /// The content type, always [expectedType].
  String get type {
    final type = _value[Keys.type] as String;
    assert(type == expectedType);
    return type;
  }

  /// The content returned by the tool, either [TextContent], [ImageContent],
  /// [AudioContent], [ResourceLink] or [EmbeddedResource].
  List<Content> get content {
    final content = (_value[Keys.content] as List?)?.cast<Content>();
    if (content == null) {
      throw ArgumentError(
        'Missing ${Keys.content} field in $ToolResultContent',
      );
    }
    return content;
  }

  /// The structured result returned by the tool.
  Map<String, Object?>? get structuredContent =>
      _value[Keys.structuredContent] as Map<String, Object?>?;

  /// Whether the tool use resulted in an error.
  bool? get isError => _value[Keys.isError] as bool?;

  /// The identifier of the tool use this result corresponds to.
  String get toolUseId => _value[Keys.toolUseId] as String;
}

/// The server's preferences for model selection, requested of the client
/// during sampling.
///
/// Because LLMs can vary along multiple dimensions, choosing the "best" model
/// is rarely straightforward.  Different models excel in different areas—some
/// are faster but less capable, others are more capable but more expensive,
/// and so on. This interface allows servers to express their priorities
/// across multiple dimensions to help clients make an appropriate selection
/// for their use case.
///
/// These preferences are always advisory. The client MAY ignore them. It is
/// also up to the client to decide how to interpret these preferences and
/// how to balance them against other considerations.
extension type ModelPreferences.fromMap(Map<String, Object?> _value) {
  factory ModelPreferences({
    List<ModelHint>? hints,
    double? costPriority,
    double? speedPriority,
    double? intelligencePriority,
  }) => ModelPreferences.fromMap({
    if (hints != null) Keys.hints: hints,
    if (costPriority != null) Keys.costPriority: costPriority,
    if (speedPriority != null) Keys.speedPriority: speedPriority,
    if (intelligencePriority != null)
      Keys.intelligencePriority: intelligencePriority,
  });

  /// Optional hints to use for model selection.
  ///
  /// If multiple hints are specified, the client MUST evaluate them in order
  /// (such that the first match is taken).
  ///
  /// The client SHOULD prioritize these hints over the numeric priorities,
  /// but MAY still use the priorities to select from ambiguous matches.
  List<ModelHint>? get hints =>
      (_value[Keys.hints] as List?)?.cast<ModelHint>();

  /// How much to prioritize cost when selecting a model.
  ///
  /// A value of 0 means cost is not important, while a value of 1 means cost
  /// is the most important factor.
  double? get costPriority => (_value[Keys.costPriority] as num?)?.toDouble();

  /// How much to prioritize sampling speed (latency) when selecting a model.
  ///
  /// A value of 0 means speed is not important, while a value of 1 means speed
  /// is the most important factor.
  double? get speedPriority => (_value[Keys.speedPriority] as num?)?.toDouble();

  /// How much to prioritize intelligence and capabilities when selecting a
  /// model.
  ///
  /// A value of 0 means intelligence is not important, while a value of 1
  /// means intelligence is the most important factor.
  double? get intelligencePriority =>
      (_value[Keys.intelligencePriority] as num?)?.toDouble();
}

/// Hints to use for model selection.
///
/// Keys not declared here are currently left unspecified by the spec and are
/// up to the client to interpret.
extension type ModelHint.fromMap(Map<String, Object?> _value) {
  factory ModelHint({String? name}) =>
      ModelHint.fromMap({if (name != null) Keys.name: name});

  /// A hint for a model name.
  ///
  /// The client SHOULD treat this as a substring of a model name; for
  /// example:
  ///  - `claude-3-5-sonnet` should match `claude-3-5-sonnet-20241022`
  ///  - `sonnet` should match `claude-3-5-sonnet-20241022`,
  ///    `claude-3-sonnet-20240229`, etc.
  ///  - `claude` should match any Claude model
  ///
  /// The client MAY also map the string to a different provider's model name
  /// or a different model family, as long as it fills a similar niche; for
  /// example:
  ///  - `gemini-1.5-flash` could match `claude-3-haiku-20240307`
  String? get name => _value[Keys.name] as String?;
}

/// Controls tool selection behavior for sampling requests.
extension type ToolChoice.fromMap(Map<String, Object?> _value) {
  factory ToolChoice({required ToolChoiceMode mode}) =>
      ToolChoice.fromMap({Keys.mode: mode.name});

  /// Controls the tool use ability of the model:
  ToolChoiceMode get mode {
    final mode = _value[Keys.mode] as String?;
    if (mode == null) {
      throw ArgumentError('Missing ${Keys.mode} field in $ToolChoice');
    }
    return ToolChoiceMode.values.firstWhere((value) => value.name == mode);
  }
}

/// The tool selection mode for sampling requests.
enum ToolChoiceMode {
  /// Model decides whether to use tools (default).
  auto,

  /// Model MUST use at least one tool before completing.
  ///
  /// On the wire, this is represented as "required", but that is a reserved
  /// keyword in Dart, so we use "require" instead.
  require(Keys.required),

  /// Model MUST NOT use any tools.
  none;

  const ToolChoiceMode([this._name]);

  final String? _name;

  String get name => _name ?? EnumName(this).name;
}