Defining a Foundation Models Tool with arguments determined at runtime

I'm experimenting with Foundation Models and I'm trying to understand how to define a Tool whose input argument is defined at runtime. Specifically, I want a Tool that takes a single String parameter that can only take certain values defined at runtime.

I think my question is basically the same as this one: https://developer.apple.com/forums/thread/793471 However, the answer provided by the engineer doesn't actually demonstrate how to create the GenerationSchema. Trying to piece things together from the documentation that the engineer linked to, I came up with this:

let citiesDefinedAtRuntime = ["London", "New York", "Paris"]

let citySchema = DynamicGenerationSchema(
    name: "CityList",
    properties: [
        DynamicGenerationSchema.Property(
            name: "city",
            schema: DynamicGenerationSchema(
                name: "city",
                anyOf: citiesDefinedAtRuntime
            )
        )
    ]
)

let generationSchema = try GenerationSchema(root: citySchema, dependencies: [])

let tools = [CityInfo(parameters: generationSchema)]

let session = LanguageModelSession(tools: tools, instructions: "...")

With the CityInfo Tool defined like this:

struct CityInfo: Tool {
    let name: String = "getCityInfo"
    let description: String = "Get information about a city."
    let parameters: GenerationSchema
    
    func call(arguments: GeneratedContent) throws -> String {
        let cityName = try arguments.value(String.self, forProperty: "city")
        print("Requested info about \(cityName)")
        let cityInfo = getCityInfo(for: cityName)
        return cityInfo
    }
    
    func getCityInfo(for city: String) -> String {
    	// some backend that provides the info
    }
}

This compiles and usually seems to work. However, sometimes the model will try to request info about a city that is not in citiesDefinedAtRuntime. For example, if I prompt the model with "I want to travel to Tokyo in Japan, can you tell me about this city?", the model will try to request info about Tokyo, even though this is not in the citiesDefinedAtRuntime array.

My understanding is that this should not be possible – constrained generation should only allow the LLM to generate an input argument from the list of cities defined in the schema.

Am I missing something here or overcomplicating things?

What's the correct way to make sure the LLM can only call a Tool with an input parameter from a set of possible values defined at runtime?

Many thanks!

Defining a Foundation Models Tool with arguments determined at runtime
 
 
Q