Add support for structured outputs to the provider (#1159)

edgedb · Dec 17, 2024 · 5de9abd · 5de9abd
1 parent d153164
commit 5de9abd
Show file tree

Hide file tree

Showing 2 changed files with 37 additions and 27 deletions.
diff --git a/packages/vercel-ai-provider/src/edgedb-chat-language-model.ts b/packages/vercel-ai-provider/src/edgedb-chat-language-model.ts
@@ -78,6 +78,7 @@ export class EdgeDBChatLanguageModel implements LanguageModelV1 {
     // providerMetadata: exists in the Vercel SDK d.ts but none of the providers use it
   }: Parameters<LanguageModelV1["doGenerate"]>[0]) {
     const type = mode.type;
+    const isAnthropic = isAnthropicModel(this.modelId);
 
     const warnings: LanguageModelV1CallWarning[] = [];
 
@@ -102,41 +103,33 @@ export class EdgeDBChatLanguageModel implements LanguageModelV1 {
       });
     }
 
-    if (isAnthropicModel(this.modelId) && seed != null) {
+    if (isAnthropic && seed != null) {
       warnings.push({
         type: "unsupported-setting",
         setting: "seed",
       });
     }
 
-    if (!isAnthropicModel(this.modelId) && topK != null) {
+    if (!isAnthropic && topK != null) {
       warnings.push({
         type: "unsupported-setting",
         setting: "topK",
       });
     }
 
-    if (responseFormat != null && responseFormat.type !== "text") {
+    if (
+      (isAnthropic && responseFormat?.type !== "text") ||
+      (!isAnthropic &&
+        responseFormat?.type === "json" &&
+        responseFormat?.schema)
+    ) {
       warnings.push({
         type: "unsupported-setting",
         setting: "responseFormat",
-        details: "JSON response format is not supported.",
+        details: "JSON response format schema is not supported.",
       });
     }
 
-    // if (
-    //   responseFormat != null &&
-    //   responseFormat.type === "json"
-    //   // && responseFormat.schema != null
-    // ) {
-    //   warnings.push({
-    //     type: "unsupported-setting",
-    //     setting: "responseFormat",
-    //     details: "JSON response is not supported",
-    //     // details: "JSON response format schema is not supported",
-    //   });
-    // }
-
     const baseArgs = {
       model: this.modelId,
       messages: convertToEdgeDBMessages(prompt),
@@ -161,17 +154,30 @@ export class EdgeDBChatLanguageModel implements LanguageModelV1 {
 
     switch (type) {
       case "regular": {
+        const { tools, tool_choice, toolWarnings } = prepareTools(
+          mode,
+          this.modelId,
+        );
+
         return {
           args: {
             ...baseArgs,
-            ...prepareTools(mode, this.modelId),
+            tools,
+            tool_choice,
           },
-          warnings,
+          warnings: [...warnings, ...toolWarnings],
         };
       }
 
-      // case 'object-json': {
-      // }
+      case "object-json": {
+        return {
+          args: {
+            ...baseArgs,
+            response_format: { type: "json_object" },
+          },
+          warnings,
+        };
+      }
 
       // case 'object-tool': {
       // }

diff --git a/packages/vercel-ai-provider/src/edgedb-prepare-tools.ts b/packages/vercel-ai-provider/src/edgedb-prepare-tools.ts
@@ -1,7 +1,8 @@
-import type {
-  JSONSchema7,
-  LanguageModelV1,
-  LanguageModelV1CallWarning,
+import {
+  type JSONSchema7,
+  type LanguageModelV1,
+  type LanguageModelV1CallWarning,
+  UnsupportedFunctionalityError,
 } from "@ai-sdk/provider";
 import {
   type EdgeDBChatModelId,
@@ -110,6 +111,7 @@ export function prepareTools(
           : isOpenAI
             ? "required"
             : "any",
+        toolWarnings,
       };
 
     // mistral does not support tool mode directly,
@@ -121,8 +123,8 @@ export function prepareTools(
             tool_choice: {
               type: "tool",
               name: toolChoice.toolName,
-              toolWarnings,
             },
+            toolWarnings,
           }
         : isOpenAI
           ? {
@@ -145,7 +147,9 @@ export function prepareTools(
 
     default: {
       const _exhaustiveCheck: never = type;
-      throw new Error(`Unsupported tool choice type: ${_exhaustiveCheck}`);
+      throw new UnsupportedFunctionalityError({
+        functionality: `Unsupported tool choice type: ${_exhaustiveCheck}`,
+      });
     }
   }
 }