@ai-sdk/cohere 3.0.8 → 3.0.9

package/src/cohere-embedding-model.test.ts

@@ -0,0 +1,143 @@
+import { EmbeddingModelV3Embedding } from '@ai-sdk/provider';
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import { createCohere } from './cohere-provider';
+import { describe, it, expect, vi } from 'vitest';
+
+vi.mock('./version', () => ({
+  VERSION: '0.0.0-test',
+}));
+
+const dummyEmbeddings = [
+  [0.1, 0.2, 0.3, 0.4, 0.5],
+  [0.6, 0.7, 0.8, 0.9, 1.0],
+];
+const testValues = ['sunny day at the beach', 'rainy day in the city'];
+
+const provider = createCohere({ apiKey: 'test-api-key' });
+const model = provider.embeddingModel('embed-english-v3.0');
+
+const server = createTestServer({
+  'https://api.cohere.com/v2/embed': {},
+});
+
+describe('doEmbed', () => {
+  function prepareJsonResponse({
+    embeddings = dummyEmbeddings,
+    meta = { billed_units: { input_tokens: 8 } },
+    headers,
+  }: {
+    embeddings?: EmbeddingModelV3Embedding[];
+    meta?: { billed_units: { input_tokens: number } };
+    headers?: Record<string, string>;
+  } = {}) {
+    server.urls['https://api.cohere.com/v2/embed'].response = {
+      type: 'json-value',
+      headers,
+      body: {
+        id: 'test-id',
+        texts: testValues,
+        embeddings: { float: embeddings },
+        meta,
+      },
+    };
+  }
+
+  it('should extract embedding', async () => {
+    prepareJsonResponse();
+
+    const { embeddings } = await model.doEmbed({ values: testValues });
+
+    expect(embeddings).toStrictEqual(dummyEmbeddings);
+  });
+
+  it('should expose the raw response', async () => {
+    prepareJsonResponse({
+      headers: { 'test-header': 'test-value' },
+    });
+
+    const { response } = await model.doEmbed({ values: testValues });
+
+    expect(response?.headers).toStrictEqual({
+      // default headers:
+      'content-length': '185',
+      'content-type': 'application/json',
+
+      // custom header
+      'test-header': 'test-value',
+    });
+    expect(response).toMatchSnapshot();
+  });
+
+  it('should extract usage', async () => {
+    prepareJsonResponse({
+      meta: { billed_units: { input_tokens: 20 } },
+    });
+
+    const { usage } = await model.doEmbed({ values: testValues });
+
+    expect(usage).toStrictEqual({ tokens: 20 });
+  });
+
+  it('should pass the model and the values', async () => {
+    prepareJsonResponse();
+
+    await model.doEmbed({ values: testValues });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      model: 'embed-english-v3.0',
+      embedding_types: ['float'],
+      texts: testValues,
+      input_type: 'search_query',
+    });
+  });
+
+  it('should pass the input_type setting', async () => {
+    prepareJsonResponse();
+
+    await provider.embeddingModel('embed-english-v3.0').doEmbed({
+      values: testValues,
+      providerOptions: {
+        cohere: {
+          inputType: 'search_document',
+        },
+      },
+    });
+
+    expect(await server.calls[0].requestBodyJson).toStrictEqual({
+      model: 'embed-english-v3.0',
+      embedding_types: ['float'],
+      texts: testValues,
+      input_type: 'search_document',
+    });
+  });
+
+  it('should pass headers', async () => {
+    prepareJsonResponse();
+
+    const provider = createCohere({
+      apiKey: 'test-api-key',
+      headers: {
+        'Custom-Provider-Header': 'provider-header-value',
+      },
+    });
+
+    await provider.embeddingModel('embed-english-v3.0').doEmbed({
+      values: testValues,
+      headers: {
+        'Custom-Request-Header': 'request-header-value',
+      },
+    });
+
+    const requestHeaders = server.calls[0].requestHeaders;
+
+    expect(requestHeaders).toStrictEqual({
+      authorization: 'Bearer test-api-key',
+      'content-type': 'application/json',
+      'custom-provider-header': 'provider-header-value',
+      'custom-request-header': 'request-header-value',
+    });
+    expect(server.calls[0].requestUserAgent).toContain(
+      `ai-sdk/cohere/0.0.0-test`,
+    );
+  });
+});

package/src/cohere-embedding-model.ts

@@ -0,0 +1,112 @@
+import {
+  EmbeddingModelV3,
+  TooManyEmbeddingValuesForCallError,
+} from '@ai-sdk/provider';
+import {
+  combineHeaders,
+  createJsonResponseHandler,
+  FetchFunction,
+  parseProviderOptions,
+  postJsonToApi,
+} from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+import {
+  CohereEmbeddingModelId,
+  cohereEmbeddingOptions,
+} from './cohere-embedding-options';
+import { cohereFailedResponseHandler } from './cohere-error';
+
+type CohereEmbeddingConfig = {
+  provider: string;
+  baseURL: string;
+  headers: () => Record<string, string | undefined>;
+  fetch?: FetchFunction;
+};
+
+export class CohereEmbeddingModel implements EmbeddingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly modelId: CohereEmbeddingModelId;
+
+  readonly maxEmbeddingsPerCall = 96;
+  readonly supportsParallelCalls = true;
+
+  private readonly config: CohereEmbeddingConfig;
+
+  constructor(modelId: CohereEmbeddingModelId, config: CohereEmbeddingConfig) {
+    this.modelId = modelId;
+    this.config = config;
+  }
+
+  get provider(): string {
+    return this.config.provider;
+  }
+
+  async doEmbed({
+    values,
+    headers,
+    abortSignal,
+    providerOptions,
+  }: Parameters<EmbeddingModelV3['doEmbed']>[0]): Promise<
+    Awaited<ReturnType<EmbeddingModelV3['doEmbed']>>
+  > {
+    const embeddingOptions = await parseProviderOptions({
+      provider: 'cohere',
+      providerOptions,
+      schema: cohereEmbeddingOptions,
+    });
+
+    if (values.length > this.maxEmbeddingsPerCall) {
+      throw new TooManyEmbeddingValuesForCallError({
+        provider: this.provider,
+        modelId: this.modelId,
+        maxEmbeddingsPerCall: this.maxEmbeddingsPerCall,
+        values,
+      });
+    }
+
+    const {
+      responseHeaders,
+      value: response,
+      rawValue,
+    } = await postJsonToApi({
+      url: `${this.config.baseURL}/embed`,
+      headers: combineHeaders(this.config.headers(), headers),
+      body: {
+        model: this.modelId,
+        // The AI SDK only supports 'float' embeddings. Note that the Cohere API
+        // supports other embedding types, but they are not currently supported by the AI SDK.
+        // https://docs.cohere.com/v2/reference/embed#request.body.embedding_types
+        embedding_types: ['float'],
+        texts: values,
+        input_type: embeddingOptions?.inputType ?? 'search_query',
+        truncate: embeddingOptions?.truncate,
+      },
+      failedResponseHandler: cohereFailedResponseHandler,
+      successfulResponseHandler: createJsonResponseHandler(
+        cohereTextEmbeddingResponseSchema,
+      ),
+      abortSignal,
+      fetch: this.config.fetch,
+    });
+
+    return {
+      warnings: [],
+      embeddings: response.embeddings.float,
+      usage: { tokens: response.meta.billed_units.input_tokens },
+      response: { headers: responseHeaders, body: rawValue },
+    };
+  }
+}
+
+// minimal version of the schema, focussed on what is needed for the implementation
+// this approach limits breakages when the API changes and increases efficiency
+const cohereTextEmbeddingResponseSchema = z.object({
+  embeddings: z.object({
+    float: z.array(z.array(z.number())),
+  }),
+  meta: z.object({
+    billed_units: z.object({
+      input_tokens: z.number(),
+    }),
+  }),
+});

package/src/cohere-embedding-options.ts

@@ -0,0 +1,37 @@
+import { z } from 'zod/v4';
+
+export type CohereEmbeddingModelId =
+  | 'embed-english-v3.0'
+  | 'embed-multilingual-v3.0'
+  | 'embed-english-light-v3.0'
+  | 'embed-multilingual-light-v3.0'
+  | 'embed-english-v2.0'
+  | 'embed-english-light-v2.0'
+  | 'embed-multilingual-v2.0'
+  | (string & {});
+
+export const cohereEmbeddingOptions = z.object({
+  /**
+   * Specifies the type of input passed to the model. Default is `search_query`.
+   *
+   * - "search_document": Used for embeddings stored in a vector database for search use-cases.
+   * - "search_query": Used for embeddings of search queries run against a vector DB to find relevant documents.
+   * - "classification": Used for embeddings passed through a text classifier.
+   * - "clustering": Used for embeddings run through a clustering algorithm.
+   */
+  inputType: z
+    .enum(['search_document', 'search_query', 'classification', 'clustering'])
+    .optional(),
+
+  /**
+   * Specifies how the API will handle inputs longer than the maximum token length.
+   * Default is `END`.
+   *
+   * - "NONE": If selected, when the input exceeds the maximum input token length will return an error.
+   * - "START": Will discard the start of the input until the remaining input is exactly the maximum input token length for the model.
+   * - "END": Will discard the end of the input until the remaining input is exactly the maximum input token length for the model.
+   */
+  truncate: z.enum(['NONE', 'START', 'END']).optional(),
+});
+
+export type CohereEmbeddingOptions = z.infer<typeof cohereEmbeddingOptions>;

package/src/cohere-error.ts

@@ -0,0 +1,13 @@
+import { createJsonErrorResponseHandler } from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+const cohereErrorDataSchema = z.object({
+  message: z.string(),
+});
+
+export type CohereErrorData = z.infer<typeof cohereErrorDataSchema>;
+
+export const cohereFailedResponseHandler = createJsonErrorResponseHandler({
+  errorSchema: cohereErrorDataSchema,
+  errorToMessage: data => data.message,
+});

package/src/cohere-prepare-tools.test.ts

@@ -0,0 +1,152 @@
+import { prepareTools } from './cohere-prepare-tools';
+import { describe, it, expect } from 'vitest';
+
+it('should return undefined tools when no tools are provided', () => {
+  const result = prepareTools({
+    tools: [],
+  });
+
+  expect(result).toStrictEqual({
+    tools: undefined,
+    toolChoice: undefined,
+    toolWarnings: [],
+  });
+});
+
+it('should process function tools correctly', () => {
+  const functionTool = {
+    type: 'function' as const,
+    name: 'testFunction',
+    description: 'test description',
+    inputSchema: { type: 'object' as const, properties: {} },
+  };
+
+  const result = prepareTools({
+    tools: [functionTool],
+  });
+
+  expect(result).toStrictEqual({
+    tools: [
+      {
+        type: 'function',
+        function: {
+          name: 'testFunction',
+          description: 'test description',
+          parameters: { type: 'object' as const, properties: {} },
+        },
+      },
+    ],
+    toolChoice: undefined,
+    toolWarnings: [],
+  });
+});
+
+it('should add warnings for provider-defined tools', () => {
+  const result = prepareTools({
+    tools: [
+      {
+        type: 'provider' as const,
+        id: 'provider.tool',
+        name: 'tool',
+        args: {},
+      },
+    ],
+  });
+
+  expect(result).toMatchInlineSnapshot(`
+    {
+      "toolChoice": undefined,
+      "toolWarnings": [
+        {
+          "feature": "provider-defined tool provider.tool",
+          "type": "unsupported",
+        },
+      ],
+      "tools": [],
+    }
+  `);
+});
+
+describe('tool choice handling', () => {
+  const basicTool = {
+    type: 'function' as const,
+    name: 'testFunction',
+    description: 'test description',
+    inputSchema: { type: 'object' as const, properties: {} },
+  };
+
+  it('should handle auto tool choice', () => {
+    const result = prepareTools({
+      tools: [basicTool],
+      toolChoice: { type: 'auto' },
+    });
+
+    expect(result.toolChoice).toBe(undefined);
+  });
+
+  it('should handle none tool choice', () => {
+    const result = prepareTools({
+      tools: [basicTool],
+      toolChoice: { type: 'none' },
+    });
+
+    expect(result).toStrictEqual({
+      tools: [
+        {
+          type: 'function',
+          function: {
+            name: 'testFunction',
+            description: 'test description',
+            parameters: { type: 'object', properties: {} },
+          },
+        },
+      ],
+      toolChoice: 'NONE',
+      toolWarnings: [],
+    });
+  });
+
+  it('should handle required tool choice', () => {
+    const result = prepareTools({
+      tools: [basicTool],
+      toolChoice: { type: 'required' },
+    });
+
+    expect(result).toStrictEqual({
+      tools: [
+        {
+          type: 'function',
+          function: {
+            name: 'testFunction',
+            description: 'test description',
+            parameters: { type: 'object', properties: {} },
+          },
+        },
+      ],
+      toolChoice: 'REQUIRED',
+      toolWarnings: [],
+    });
+  });
+
+  it('should handle tool type tool choice by filtering tools', () => {
+    const result = prepareTools({
+      tools: [basicTool],
+      toolChoice: { type: 'tool', toolName: 'testFunction' },
+    });
+
+    expect(result).toStrictEqual({
+      tools: [
+        {
+          type: 'function',
+          function: {
+            name: 'testFunction',
+            description: 'test description',
+            parameters: { type: 'object', properties: {} },
+          },
+        },
+      ],
+      toolChoice: 'REQUIRED',
+      toolWarnings: [],
+    });
+  });
+});

package/src/cohere-prepare-tools.ts

@@ -0,0 +1,96 @@
+import {
+  LanguageModelV3CallOptions,
+  SharedV3Warning,
+  UnsupportedFunctionalityError,
+} from '@ai-sdk/provider';
+import { CohereToolChoice } from './cohere-chat-prompt';
+
+export function prepareTools({
+  tools,
+  toolChoice,
+}: {
+  tools: LanguageModelV3CallOptions['tools'];
+  toolChoice?: LanguageModelV3CallOptions['toolChoice'];
+}): {
+  tools:
+    | Array<{
+        type: 'function';
+        function: {
+          name: string | undefined;
+          description: string | undefined;
+          parameters: unknown;
+        };
+      }>
+    | undefined;
+  toolChoice: CohereToolChoice;
+  toolWarnings: SharedV3Warning[];
+} {
+  // when the tools array is empty, change it to undefined to prevent errors:
+  tools = tools?.length ? tools : undefined;
+
+  const toolWarnings: SharedV3Warning[] = [];
+
+  if (tools == null) {
+    return { tools: undefined, toolChoice: undefined, toolWarnings };
+  }
+
+  const cohereTools: Array<{
+    type: 'function';
+    function: {
+      name: string;
+      description: string | undefined;
+      parameters: unknown;
+    };
+  }> = [];
+
+  for (const tool of tools) {
+    if (tool.type === 'provider') {
+      toolWarnings.push({
+        type: 'unsupported',
+        feature: `provider-defined tool ${tool.id}`,
+      });
+    } else {
+      cohereTools.push({
+        type: 'function',
+        function: {
+          name: tool.name,
+          description: tool.description,
+          parameters: tool.inputSchema,
+        },
+      });
+    }
+  }
+
+  if (toolChoice == null) {
+    return { tools: cohereTools, toolChoice: undefined, toolWarnings };
+  }
+
+  const type = toolChoice.type;
+
+  switch (type) {
+    case 'auto':
+      return { tools: cohereTools, toolChoice: undefined, toolWarnings };
+
+    case 'none':
+      return { tools: cohereTools, toolChoice: 'NONE', toolWarnings };
+
+    case 'required':
+      return { tools: cohereTools, toolChoice: 'REQUIRED', toolWarnings };
+
+    case 'tool':
+      return {
+        tools: cohereTools.filter(
+          tool => tool.function.name === toolChoice.toolName,
+        ),
+        toolChoice: 'REQUIRED',
+        toolWarnings,
+      };
+
+    default: {
+      const _exhaustiveCheck: never = type;
+      throw new UnsupportedFunctionalityError({
+        functionality: `tool choice type: ${_exhaustiveCheck}`,
+      });
+    }
+  }
+}