Contents

CustomStage

A custom processing stage the Spotlight search tool uses to identify search results.

Declaration

protocol CustomStage : Generable, Decodable, Encodable, Sendable

Mentioned in

Overview

A custom stage is a generable type that implements app-specific data transformations for queries. When using Foundation Models, you can use the Spotlight search tool to find app-specific content related to a prompt. The model uses the Spotlight search tool to create queries, each of which might involve require several steps to deliver the final results. For example, a query might fetch items from your app’s Spotlight index, count the number of items it fetched, and assign relevance scores to each item. Each of these steps is a stage in the query pipeline, and a custom stage lets you integrate your app’s custom transformations.

Define custom stages as a generable type, and implement your stage’s behavior using the properties and methods of this protocol. A custom stage includes static properties that the model uses to assess how to apply the stage to queries. It also includes execute methods to perform the actual data transformations. Each execute method takes one of the input types your stage supports and delivers the specified output type.

The following example shows an implementation of this type that accepts Spotlight searchable items as input and produces scored items as output. The execute method in the implementation uses a custom SentimentAnalyzer type to calculate the score for each item, based on whether its content conveys a positive, negative, or neutral tone. The example also includes an extension with a static sentiment function, which simplifies the creation of the custom stage later.

@Generable
struct SentimentStage: CustomStage {
    static var name: String { "sentiment" }
    static var description: String { "Scores search results by sentiment.” }
    static var inputTypes: [SearchPipelineDataType] { [.items] }
    static var outputType: SearchPipelineDataType { .scoredItems }

    @Guide(description: “The sentiment to consider when scoring the text of a search result.”)
    var mode: String

    func execute(items: [SearchableItem]) async throws -> SearchPipelineData {
        let scored = items.map { item in
            ScoredSearchableItem(item: item,
                                   score: SentimentAnalyzer.score(item, mode: mode))
        }
        return .scoredItems(scored)
    }
}

extension CustomStage where Self == SentimentStage {
     static func sentiment(mode: String = "all") -> Self {
         SentimentStage(mode: mode)
     }
}

To make your custom stage available to a model, include it in the configuration of the Spotlight search tool you associate with your Foundation model’s session. The following example configures the Spotlight search tool with two separate instances of the sentiment stage from the previous example. The first instance scores items across all sentiments while the second instance scores items only on the positivity scale.

let tool = SpotlightSearchTool(configuration: .init(
      customStages: [.sentiment(), .sentiment(mode: "positive")]
))

The model builds tool pipelines dynamically, and can run multiple stages in parallel, so implement custom stage types to run independently. Treat the input data your stage receives as immutable, and don’t consider the state or contents of other stages when making decisions. If you do require additional data to generate results, make sure you access the data in a deterministic way.

Topics

Getting the stage metadata

Performing the stage behavior

Instance Methods

See Also

Tool customization