Support JSON schema for function parameters in chat models

ai-mocks-core/src/commonMain/kotlin/me/kpavlov/aimocks/core/json/schema/SchemaHelper.kt

@@ -0,0 +1,116 @@
+package me.kpavlov.aimocks.core.json.schema
+
+import kotlinx.serialization.SerializationException
+import kotlinx.serialization.json.Json
+import kotlinx.serialization.json.JsonElement
+
+/**
+ * Helper object for working with JSON schemas.
+ *
+ * Provides utility functions for parsing and validating JSON Schema definitions.
+ */
+public object SchemaHelper {
+    private val json =
+        Json {
+            ignoreUnknownKeys = true
+        }
+
+    /**
+     * Parses a [JsonElement] into a [SchemaDefinition].
+     *
+     * @param schemaJson The JSON element containing the schema
+     * @return The parsed schema definition, or null if parsing fails
+     */
+    public fun parseSchema(schemaJson: JsonElement?): SchemaDefinition? =
+        schemaJson?.let {
+            try {
+                json.decodeFromJsonElement(SchemaDefinition.serializer(), it)
+            } catch (_: SerializationException) {
+                null
+            }
+        }
+
+    /**
+     * Checks if a schema has a property with the specified name.
+     *
+     * @param schema The schema definition to check
+     * @param propertyName The name of the property
+     * @return true if the property exists, false otherwise
+     */
+    public fun hasProperty(
+        schema: SchemaDefinition,
+        propertyName: String,
+    ): Boolean = schema.properties.containsKey(propertyName)
+
+    /**
+     * Gets the types of a property in the schema.
+     *
+     * @param schema The schema definition
+     * @param propertyName The name of the property
+     * @return The list of types for the property, or null if not found or not a value property
+     */
+    public fun getPropertyType(
+        schema: SchemaDefinition,
+        propertyName: String,
+    ): List<String>? =
+        schema.properties[propertyName]?.let { prop ->
+            when (prop) {
+                is ValuePropertyDefinition -> prop.type
+                else -> null
+            }
+        }
+
+    /**
+     * Checks if a property is required in the schema.
+     *
+     * @param schema The schema definition
+     * @param propertyName The name of the property
+     * @return true if the property is required, false otherwise
+     */
+    public fun isPropertyRequired(
+        schema: SchemaDefinition,
+        propertyName: String,
+    ): Boolean = schema.required.contains(propertyName)
+
+    /**
+     * Checks if all specified properties are required in the schema.
+     *
+     * @param schema The schema definition
+     * @param propertyNames The names of properties to check
+     * @return true if all properties are required, false otherwise
+     */
+    public fun hasAllRequiredProperties(
+        schema: SchemaDefinition,
+        vararg propertyNames: String,
+    ): Boolean = propertyNames.all { it in schema.required }
+
+    /**
+     * Gets the property definition for a given property name.
+     *
+     * @param schema The schema definition
+     * @param propertyName The name of the property
+     * @return The property definition, or null if not found
+     */
+    public fun getProperty(
+        schema: SchemaDefinition,
+        propertyName: String,
+    ): PropertyDefinition? = schema.properties[propertyName]
+
+    /**
+     * Gets the description of a property in the schema.
+     *
+     * @param schema The schema definition
+     * @param propertyName The name of the property
+     * @return The description of the property, or null if not found
+     */
+    public fun getPropertyDescription(
+        schema: SchemaDefinition,
+        propertyName: String,
+    ): String? =
+        schema.properties[propertyName]?.let { prop ->
+            when (prop) {
+                is ValuePropertyDefinition -> prop.description
+                else -> null
+            }
+        }
+}

ai-mocks-openai/src/commonMain/kotlin/me/kpavlov/aimocks/openai/Model.kt

@@ -9,6 +9,7 @@ import kotlinx.serialization.ExperimentalSerializationApi
 import kotlinx.serialization.Required
 import kotlinx.serialization.SerialName
 import kotlinx.serialization.Serializable
+import kotlinx.serialization.json.JsonElement
 import me.kpavlov.aimocks.core.json.schema.JsonSchema
 import me.kpavlov.aimocks.openai.model.ChatCompletionRole
 import me.kpavlov.aimocks.openai.model.ChatCompletionStreamOptions
@@ -238,7 +239,7 @@ public data class FunctionObject(
     @SerialName(value = "name") @Required val name: String,
     @SerialName(value = "description") val description: String? = null,
     @SerialName(value = "parameters")
-    val parameters: Map<String, String>? = null,
+    val parameters: JsonElement? = null,
     @SerialName(value = "strict") val strict: Boolean? = false,
 )
 

ai-mocks-openai/src/commonMain/kotlin/me/kpavlov/aimocks/openai/completions/OpenaiChatCompletionRequestSpecification.kt

@@ -39,4 +39,71 @@ public open class OpenaiChatCompletionRequestSpecification(
     override fun userMessageContains(substring: String) {
         requestBody.add(OpenaiCompletionsMatchers.userMessageContains(substring))
     }
+
+    /**
+     * Adds a matcher to verify that the request includes a tool with the specified function name.
+     *
+     * @param functionName The name of the function to match
+     */
+    public fun hasToolWithFunction(functionName: String) {
+        requestBody.add(OpenaiCompletionsMatchers.hasToolWithFunction(functionName))
+    }
+
+    /**
+     * Adds a matcher to verify that a tool's function has a parameter with the specified name.
+     *
+     * @param functionName The name of the function
+     * @param parameterName The name of the parameter to check
+     */
+    public fun toolHasParameter(
+        functionName: String,
+        parameterName: String,
+    ) {
+        requestBody.add(OpenaiCompletionsMatchers.toolHasParameter(functionName, parameterName))
+    }
+
+    /**
+     * Adds a matcher to verify that a tool's function has a parameter with the specified name and description.
+     *
+     * @param functionName The name of the function
+     * @param parameterName The name of the parameter to check
+     * @param description The expected description of the parameter
+     */
+    public fun toolHasParameter(
+        functionName: String,
+        parameterName: String,
+        description: String,
+    ) {
+        requestBody.add(OpenaiCompletionsMatchers.toolHasParameter(functionName, parameterName, description))
+    }
+
+    /**
+     * Adds a matcher to verify that a tool's function parameter has the specified type.
+     *
+     * @param functionName The name of the function
+     * @param parameterName The name of the parameter
+     * @param expectedType The expected type (e.g., "string", "integer", "number", "boolean", "object", "array")
+     */
+    public fun toolParameterHasType(
+        functionName: String,
+        parameterName: String,
+        expectedType: String,
+    ) {
+        requestBody.add(
+            OpenaiCompletionsMatchers.toolParameterHasType(functionName, parameterName, expectedType),
+        )
+    }
+
+    /**
+     * Adds a matcher to verify that a tool's function requires specific parameters.
+     *
+     * @param functionName The name of the function
+     * @param requiredParams The parameter names that should be required
+     */
+    public fun toolRequiresParameters(
+        functionName: String,
+        vararg requiredParams: String,
+    ) {
+        requestBody.add(OpenaiCompletionsMatchers.toolRequiresParameters(functionName, *requiredParams))
+    }
 }

ai-mocks-openai/src/commonMain/kotlin/me/kpavlov/aimocks/openai/completions/OpenaiCompletionsMatchers.kt

@@ -2,6 +2,7 @@ package me.kpavlov.aimocks.openai.completions
 
 import io.kotest.matchers.Matcher
 import io.kotest.matchers.MatcherResult
+import me.kpavlov.aimocks.core.json.schema.SchemaHelper
 import me.kpavlov.aimocks.openai.ChatCompletionRequest
 import me.kpavlov.aimocks.openai.model.ChatCompletionRole
 
@@ -45,4 +46,174 @@ internal object OpenaiCompletionsMatchers {
 
             override fun toString(): String = "User message should contain \"$string\""
         }
+
+    /**
+     * Matches requests that have a tool with the specified function name.
+     *
+     * @param functionName The name of the function to match
+     * @return A matcher that checks if the request contains a tool with the specified function name
+     */
+    fun hasToolWithFunction(functionName: String): Matcher<ChatCompletionRequest?> =
+        object : Matcher<ChatCompletionRequest?> {
+            override fun test(value: ChatCompletionRequest?): MatcherResult =
+                MatcherResult.Companion(
+                    value?.tools?.any { tool ->
+                        tool.function.name == functionName
+                    } == true,
+                    { "Request should have tool with function name \"$functionName\"" },
+                    { "Request should not have tool with function name \"$functionName\"" },
+                )
+
+            override fun toString(): String =
+                "Request should have tool with function name \"$functionName\""
+        }
+
+    /**
+     * Matches requests where a tool's function has a parameter with the specified name.
+     *
+     * @param functionName The name of the function
+     * @param parameterName The name of the parameter to check
+     * @return A matcher that checks if the specified function has the named parameter
+     */
+    fun toolHasParameter(
+        functionName: String,
+        parameterName: String,
+    ): Matcher<ChatCompletionRequest?> =
+        object : Matcher<ChatCompletionRequest?> {
+            override fun test(value: ChatCompletionRequest?): MatcherResult {
+                val tool = value?.tools?.find { it.function.name == functionName }
+                val schema = tool?.function?.parameters?.let { SchemaHelper.parseSchema(it) }
+                val hasParameter =
+                    schema?.let { SchemaHelper.hasProperty(it, parameterName) } == true
+
+                return MatcherResult.Companion(
+                    hasParameter,
+                    { "Function \"$functionName\" should have parameter \"$parameterName\"" },
+                    { "Function \"$functionName\" should not have parameter \"$parameterName\"" },
+                )
+            }
+
+            override fun toString(): String =
+                "Function \"$functionName\" should have parameter \"$parameterName\""
+        }
+
+    /**
+     * Matches requests where a tool's function has a parameter with the specified name and description.
+     *
+     * @param functionName The name of the function
+     * @param parameterName The name of the parameter to check
+     * @param description The expected description of the parameter
+     * @return A matcher that checks if the specified function has the named parameter with the given description
+     */
+    fun toolHasParameter(
+        functionName: String,
+        parameterName: String,
+        description: String,
+    ): Matcher<ChatCompletionRequest?> =
+        object : Matcher<ChatCompletionRequest?> {
+            override fun test(value: ChatCompletionRequest?): MatcherResult {
+                val tool = value?.tools?.find { it.function.name == functionName }
+                val schema = tool?.function?.parameters?.let { SchemaHelper.parseSchema(it) }
+                val actualDescription = schema?.let { SchemaHelper.getPropertyDescription(it, parameterName) }
+                val matches = actualDescription == description
+
+                return MatcherResult.Companion(
+                    matches,
+                    {
+                        "Function \"$functionName\" parameter \"$parameterName\" should have description \"$description\""
+                    },
+                    {
+                        "Function \"$functionName\" parameter \"$parameterName\" should not have description \"$description\""
+                    },
+                )
+            }
+
+            override fun toString(): String =
+                "Function \"$functionName\" parameter \"$parameterName\" should have description \"$description\""
+        }
+
+    /**
+     * Matches requests where a tool's function parameter has the specified type.
+     *
+     * @param functionName The name of the function
+     * @param parameterName The name of the parameter
+     * @param expectedType The expected type (e.g., "string", "integer", "number", "boolean", "object", "array")
+     * @return A matcher that checks if the parameter has the expected type
+     */
+    fun toolParameterHasType(
+        functionName: String,
+        parameterName: String,
+        expectedType: String,
+    ): Matcher<ChatCompletionRequest?> =
+        object : Matcher<ChatCompletionRequest?> {
+            override fun test(value: ChatCompletionRequest?): MatcherResult {
+                val tool = value?.tools?.find { it.function.name == functionName }
+                val schema = tool?.function?.parameters?.let { SchemaHelper.parseSchema(it) }
+                val actualTypes = schema?.let { SchemaHelper.getPropertyType(it, parameterName) }
+                val hasCorrectType = actualTypes?.contains(expectedType) == true
+
+                return MatcherResult.Companion(
+                    hasCorrectType,
+                    {
+                        "Function \"$functionName\" parameter \"$parameterName\" should have type \"$expectedType\""
+                    },
+                    {
+                        "Function \"$functionName\" parameter \"$parameterName\" should not have type \"$expectedType\""
+                    },
+                )
+            }
+
+            override fun toString(): String =
+                "Function \"$functionName\" parameter \"$parameterName\" should have type \"$expectedType\""
+        }
+
+    /**
+     * Matches requests where a tool's function schema requires specific parameters.
+     *
+     * @param functionName The name of the function
+     * @param requiredParams The list of parameter names that should be required
+     * @return A matcher that checks if all specified parameters are required
+     */
+    fun toolRequiresParameters(
+        functionName: String,
+        vararg requiredParams: String,
+    ): Matcher<ChatCompletionRequest?> =
+        object : Matcher<ChatCompletionRequest?> {
+            override fun test(value: ChatCompletionRequest?): MatcherResult {
+                val tool = value?.tools?.find { it.function.name == functionName }
+                val schema = tool?.function?.parameters?.let { SchemaHelper.parseSchema(it) }
+                val hasAllRequired =
+                    schema?.let {
+                        SchemaHelper.hasAllRequiredProperties(
+                            it,
+                            *requiredParams,
+                        )
+                    } == true
+
+                return MatcherResult.Companion(
+                    hasAllRequired,
+                    {
+                        "Function \"$functionName\" should require parameters: ${
+                            requiredParams.joinToString(
+                                ", ",
+                            )
+                        }"
+                    },
+                    {
+                        "Function \"$functionName\" should not require parameters: ${
+                            requiredParams.joinToString(
+                                ", ",
+                            )
+                        }"
+                    },
+                )
+            }
+
+            override fun toString(): String =
+                "Function \"$functionName\" should require parameters: ${
+                    requiredParams.joinToString(
+                        ", ",
+                    )
+                }"
+        }
 }