@ai-sdk/cohere 3.0.7 → 3.0.9

package/src/cohere-provider.ts

@@ -0,0 +1,169 @@
+import {
+  EmbeddingModelV3,
+  LanguageModelV3,
+  NoSuchModelError,
+  RerankingModelV3,
+  ProviderV3,
+} from '@ai-sdk/provider';
+
+import {
+  FetchFunction,
+  generateId,
+  loadApiKey,
+  withoutTrailingSlash,
+  withUserAgentSuffix,
+} from '@ai-sdk/provider-utils';
+import { CohereChatLanguageModel } from './cohere-chat-language-model';
+import { CohereChatModelId } from './cohere-chat-options';
+import { CohereEmbeddingModel } from './cohere-embedding-model';
+import { CohereRerankingModelId } from './reranking/cohere-reranking-options';
+import { CohereRerankingModel } from './reranking/cohere-reranking-model';
+import { CohereEmbeddingModelId } from './cohere-embedding-options';
+import { VERSION } from './version';
+
+export interface CohereProvider extends ProviderV3 {
+  (modelId: CohereChatModelId): LanguageModelV3;
+
+  /**
+Creates a model for text generation.
+*/
+  languageModel(modelId: CohereChatModelId): LanguageModelV3;
+
+  /**
+   * Creates a model for text embeddings.
+   */
+  embedding(modelId: CohereEmbeddingModelId): EmbeddingModelV3;
+
+  /**
+   * Creates a model for text embeddings.
+   */
+  embeddingModel(modelId: CohereEmbeddingModelId): EmbeddingModelV3;
+
+  /**
+   * @deprecated Use `embedding` instead.
+   */
+  textEmbedding(modelId: CohereEmbeddingModelId): EmbeddingModelV3;
+
+  /**
+   * @deprecated Use `embeddingModel` instead.
+   */
+  textEmbeddingModel(modelId: CohereEmbeddingModelId): EmbeddingModelV3;
+
+  /**
+   * Creates a model for reranking.
+   */
+  reranking(modelId: CohereRerankingModelId): RerankingModelV3;
+
+  /**
+   * Creates a model for reranking.
+   */
+  rerankingModel(modelId: CohereRerankingModelId): RerankingModelV3;
+}
+
+export interface CohereProviderSettings {
+  /**
+Use a different URL prefix for API calls, e.g. to use proxy servers.
+The default prefix is `https://api.cohere.com/v2`.
+   */
+  baseURL?: string;
+
+  /**
+API key that is being send using the `Authorization` header.
+It defaults to the `COHERE_API_KEY` environment variable.
+   */
+  apiKey?: string;
+
+  /**
+Custom headers to include in the requests.
+     */
+  headers?: Record<string, string>;
+
+  /**
+Custom fetch implementation. You can use it as a middleware to intercept requests,
+or to provide a custom fetch implementation for e.g. testing.
+    */
+  fetch?: FetchFunction;
+
+  /**
+Optional function to generate a unique ID for each request.
+     */
+  generateId?: () => string;
+}
+
+/**
+Create a Cohere AI provider instance.
+ */
+export function createCohere(
+  options: CohereProviderSettings = {},
+): CohereProvider {
+  const baseURL =
+    withoutTrailingSlash(options.baseURL) ?? 'https://api.cohere.com/v2';
+
+  const getHeaders = () =>
+    withUserAgentSuffix(
+      {
+        Authorization: `Bearer ${loadApiKey({
+          apiKey: options.apiKey,
+          environmentVariableName: 'COHERE_API_KEY',
+          description: 'Cohere',
+        })}`,
+        ...options.headers,
+      },
+      `ai-sdk/cohere/${VERSION}`,
+    );
+
+  const createChatModel = (modelId: CohereChatModelId) =>
+    new CohereChatLanguageModel(modelId, {
+      provider: 'cohere.chat',
+      baseURL,
+      headers: getHeaders,
+      fetch: options.fetch,
+      generateId: options.generateId ?? generateId,
+    });
+
+  const createEmbeddingModel = (modelId: CohereEmbeddingModelId) =>
+    new CohereEmbeddingModel(modelId, {
+      provider: 'cohere.textEmbedding',
+      baseURL,
+      headers: getHeaders,
+      fetch: options.fetch,
+    });
+
+  const createRerankingModel = (modelId: CohereRerankingModelId) =>
+    new CohereRerankingModel(modelId, {
+      provider: 'cohere.reranking',
+      baseURL,
+      headers: getHeaders,
+      fetch: options.fetch,
+    });
+
+  const provider = function (modelId: CohereChatModelId) {
+    if (new.target) {
+      throw new Error(
+        'The Cohere model function cannot be called with the new keyword.',
+      );
+    }
+
+    return createChatModel(modelId);
+  };
+
+  provider.specificationVersion = 'v3' as const;
+  provider.languageModel = createChatModel;
+  provider.embedding = createEmbeddingModel;
+  provider.embeddingModel = createEmbeddingModel;
+  provider.textEmbedding = createEmbeddingModel;
+  provider.textEmbeddingModel = createEmbeddingModel;
+  provider.reranking = createRerankingModel;
+  provider.rerankingModel = createRerankingModel;
+
+  provider.imageModel = (modelId: string) => {
+    throw new NoSuchModelError({ modelId, modelType: 'imageModel' });
+  };
+
+  return provider;
+}
+
+/**
+Default Cohere provider instance.
+ */
+export const cohere = createCohere();

package/src/convert-cohere-usage.ts

@@ -0,0 +1,45 @@
+import { LanguageModelV3Usage } from '@ai-sdk/provider';
+
+export type CohereUsageTokens = {
+  input_tokens: number;
+  output_tokens: number;
+};
+
+export function convertCohereUsage(
+  tokens: CohereUsageTokens | undefined | null,
+): LanguageModelV3Usage {
+  if (tokens == null) {
+    return {
+      inputTokens: {
+        total: undefined,
+        noCache: undefined,
+        cacheRead: undefined,
+        cacheWrite: undefined,
+      },
+      outputTokens: {
+        total: undefined,
+        text: undefined,
+        reasoning: undefined,
+      },
+      raw: undefined,
+    };
+  }
+
+  const inputTokens = tokens.input_tokens;
+  const outputTokens = tokens.output_tokens;
+
+  return {
+    inputTokens: {
+      total: inputTokens,
+      noCache: inputTokens,
+      cacheRead: undefined,
+      cacheWrite: undefined,
+    },
+    outputTokens: {
+      total: outputTokens,
+      text: outputTokens,
+      reasoning: undefined,
+    },
+    raw: tokens,
+  };
+}

package/src/convert-to-cohere-chat-prompt.test.ts

@@ -0,0 +1,175 @@
+import { convertToCohereChatPrompt } from './convert-to-cohere-chat-prompt';
+import { describe, it, expect } from 'vitest';
+
+describe('convert to cohere chat prompt', () => {
+  describe('file processing', () => {
+    it('should extract documents from file parts', () => {
+      const result = convertToCohereChatPrompt([
+        {
+          role: 'user',
+          content: [
+            { type: 'text', text: 'Analyze this file: ' },
+            {
+              type: 'file',
+              data: Buffer.from('This is file content'),
+              mediaType: 'text/plain',
+              filename: 'test.txt',
+            },
+          ],
+        },
+      ]);
+
+      expect(result).toEqual({
+        messages: [
+          {
+            role: 'user',
+            content: 'Analyze this file: ',
+          },
+        ],
+        documents: [
+          {
+            data: {
+              text: 'This is file content',
+              title: 'test.txt',
+            },
+          },
+        ],
+        warnings: [],
+      });
+    });
+
+    it('should throw error for unsupported media types', () => {
+      expect(() => {
+        convertToCohereChatPrompt([
+          {
+            role: 'user',
+            content: [
+              {
+                type: 'file',
+                data: Buffer.from('PDF content'),
+                mediaType: 'application/pdf',
+                filename: 'test.pdf',
+              },
+            ],
+          },
+        ]);
+      }).toThrow("Media type 'application/pdf' is not supported");
+    });
+  });
+
+  describe('tool messages', () => {
+    it('should convert a tool call into a cohere chatbot message', async () => {
+      const result = convertToCohereChatPrompt([
+        {
+          role: 'assistant',
+          content: [
+            {
+              type: 'text',
+              text: 'Calling a tool',
+            },
+            {
+              type: 'tool-call',
+              toolName: 'tool-1',
+              toolCallId: 'tool-call-1',
+              input: { test: 'This is a tool message' },
+            },
+          ],
+        },
+      ]);
+
+      expect(result).toEqual({
+        messages: [
+          {
+            content: undefined,
+            role: 'assistant',
+            tool_calls: [
+              {
+                id: 'tool-call-1',
+                type: 'function',
+                function: {
+                  name: 'tool-1',
+                  arguments: JSON.stringify({ test: 'This is a tool message' }),
+                },
+              },
+            ],
+          },
+        ],
+        documents: [],
+        warnings: [],
+      });
+    });
+
+    it('should convert a single tool result into a cohere tool message', async () => {
+      const result = convertToCohereChatPrompt([
+        {
+          role: 'tool',
+          content: [
+            {
+              type: 'tool-result',
+              toolName: 'tool-1',
+              toolCallId: 'tool-call-1',
+              output: {
+                type: 'json',
+                value: { test: 'This is a tool message' },
+              },
+            },
+          ],
+        },
+      ]);
+
+      expect(result).toEqual({
+        messages: [
+          {
+            role: 'tool',
+            content: JSON.stringify({ test: 'This is a tool message' }),
+            tool_call_id: 'tool-call-1',
+          },
+        ],
+        documents: [],
+        warnings: [],
+      });
+    });
+
+    it('should convert multiple tool results into a cohere tool message', async () => {
+      const result = convertToCohereChatPrompt([
+        {
+          role: 'tool',
+          content: [
+            {
+              type: 'tool-result',
+              toolName: 'tool-1',
+              toolCallId: 'tool-call-1',
+              output: {
+                type: 'json',
+                value: { test: 'This is a tool message' },
+              },
+            },
+            {
+              type: 'tool-result',
+              toolName: 'tool-2',
+              toolCallId: 'tool-call-2',
+              output: { type: 'json', value: { something: 'else' } },
+            },
+          ],
+        },
+      ]);
+
+      expect(result).toEqual({
+        messages: [
+          {
+            role: 'tool',
+            content: JSON.stringify({ test: 'This is a tool message' }),
+            tool_call_id: 'tool-call-1',
+          },
+          {
+            role: 'tool',
+            content: JSON.stringify({ something: 'else' }),
+            tool_call_id: 'tool-call-2',
+          },
+        ],
+        documents: [],
+        warnings: [],
+      });
+    });
+  });
+});

package/src/convert-to-cohere-chat-prompt.ts

@@ -0,0 +1,156 @@
+import {
+  SharedV3Warning,
+  LanguageModelV3Prompt,
+  UnsupportedFunctionalityError,
+} from '@ai-sdk/provider';
+import { CohereAssistantMessage, CohereChatPrompt } from './cohere-chat-prompt';
+
+export function convertToCohereChatPrompt(prompt: LanguageModelV3Prompt): {
+  messages: CohereChatPrompt;
+  documents: Array<{
+    data: { text: string; title?: string };
+  }>;
+  warnings: SharedV3Warning[];
+} {
+  const messages: CohereChatPrompt = [];
+  const documents: Array<{ data: { text: string; title?: string } }> = [];
+  const warnings: SharedV3Warning[] = [];
+
+  for (const { role, content } of prompt) {
+    switch (role) {
+      case 'system': {
+        messages.push({ role: 'system', content });
+        break;
+      }
+
+      case 'user': {
+        messages.push({
+          role: 'user',
+          content: content
+            .map(part => {
+              switch (part.type) {
+                case 'text': {
+                  return part.text;
+                }
+                case 'file': {
+                  // Extract documents for RAG
+                  let textContent: string;
+
+                  if (typeof part.data === 'string') {
+                    // Base64 or text data
+                    textContent = part.data;
+                  } else if (part.data instanceof Uint8Array) {
+                    // Check if the media type is supported for text extraction
+                    if (
+                      !(
+                        part.mediaType?.startsWith('text/') ||
+                        part.mediaType === 'application/json'
+                      )
+                    ) {
+                      throw new UnsupportedFunctionalityError({
+                        functionality: `document media type: ${part.mediaType}`,
+                        message: `Media type '${part.mediaType}' is not supported. Supported media types are: text/* and application/json.`,
+                      });
+                    }
+                    textContent = new TextDecoder().decode(part.data);
+                  } else {
+                    throw new UnsupportedFunctionalityError({
+                      functionality: 'File URL data',
+                      message:
+                        'URLs should be downloaded by the AI SDK and not reach this point. This indicates a configuration issue.',
+                    });
+                  }
+
+                  documents.push({
+                    data: {
+                      text: textContent,
+                      title: part.filename,
+                    },
+                  });
+
+                  // Files are handled separately via the documents parameter
+                  // Return empty string to not include file content in message text
+                  return '';
+                }
+              }
+            })
+            .join(''),
+        });
+        break;
+      }
+
+      case 'assistant': {
+        let text = '';
+        const toolCalls: CohereAssistantMessage['tool_calls'] = [];
+
+        for (const part of content) {
+          switch (part.type) {
+            case 'text': {
+              text += part.text;
+              break;
+            }
+            case 'tool-call': {
+              toolCalls.push({
+                id: part.toolCallId,
+                type: 'function' as const,
+                function: {
+                  name: part.toolName,
+                  arguments: JSON.stringify(part.input),
+                },
+              });
+              break;
+            }
+          }
+        }
+
+        messages.push({
+          role: 'assistant',
+          content: toolCalls.length > 0 ? undefined : text,
+          tool_calls: toolCalls.length > 0 ? toolCalls : undefined,
+          tool_plan: undefined,
+        });
+
+        break;
+      }
+      case 'tool': {
+        messages.push(
+          ...content
+            .filter(toolResult => toolResult.type !== 'tool-approval-response')
+            .map(toolResult => {
+              const output = toolResult.output;
+
+              let contentValue: string;
+              switch (output.type) {
+                case 'text':
+                case 'error-text':
+                  contentValue = output.value;
+                  break;
+                case 'execution-denied':
+                  contentValue = output.reason ?? 'Tool execution denied.';
+                  break;
+                case 'content':
+                case 'json':
+                case 'error-json':
+                  contentValue = JSON.stringify(output.value);
+                  break;
+              }
+
+              return {
+                role: 'tool' as const,
+                content: contentValue,
+                tool_call_id: toolResult.toolCallId,
+              };
+            }),
+        );
+
+        break;
+      }
+      default: {
+        const _exhaustiveCheck: never = role;
+        throw new Error(`Unsupported role: ${_exhaustiveCheck}`);
+      }
+    }
+  }
+
+  return { messages, documents, warnings };
+}

package/src/index.ts

@@ -0,0 +1,5 @@
+export type { CohereChatModelOptions } from './cohere-chat-options';
+export { cohere, createCohere } from './cohere-provider';
+export type { CohereProvider, CohereProviderSettings } from './cohere-provider';
+export type { CohereRerankingOptions } from './reranking/cohere-reranking-options';
+export { VERSION } from './version';

package/src/map-cohere-finish-reason.ts

@@ -0,0 +1,23 @@
+import { LanguageModelV3FinishReason } from '@ai-sdk/provider';
+
+export function mapCohereFinishReason(
+  finishReason: string | null | undefined,
+): LanguageModelV3FinishReason['unified'] {
+  switch (finishReason) {
+    case 'COMPLETE':
+    case 'STOP_SEQUENCE':
+      return 'stop';
+
+    case 'MAX_TOKENS':
+      return 'length';
+
+    case 'ERROR':
+      return 'error';
+
+    case 'TOOL_CALL':
+      return 'tool-calls';
+
+    default:
+      return 'other';
+  }
+}

package/src/reranking/__fixtures__/cohere-reranking.1.json

@@ -0,0 +1,21 @@
+{
+  "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72",
+  "results": [
+    {
+      "index": 1,
+      "relevance_score": 0.10183054
+    },
+    {
+      "index": 0,
+      "relevance_score": 0.03762639
+    }
+  ],
+  "meta": {
+    "api_version": {
+      "version": "2"
+    },
+    "billed_units": {
+      "search_units": 1
+    }
+  }
+}

package/src/reranking/cohere-reranking-api.ts

@@ -0,0 +1,27 @@
+import { lazySchema, zodSchema } from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+// https://docs.cohere.com/v2/reference/rerank
+export type CohereRerankingInput = {
+  model: string;
+  query: string;
+  documents: string[];
+  top_n: number | undefined;
+  max_tokens_per_doc: number | undefined;
+  priority: number | undefined;
+};
+
+export const cohereRerankingResponseSchema = lazySchema(() =>
+  zodSchema(
+    z.object({
+      id: z.string().nullish(),
+      results: z.array(
+        z.object({
+          index: z.number(),
+          relevance_score: z.number(),
+        }),
+      ),
+      meta: z.any(),
+    }),
+  ),
+);