@ai-sdk/cohere 3.0.8 → 3.0.9

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

@@ -0,0 +1,243 @@
+import { createTestServer } from '@ai-sdk/test-server/with-vitest';
+import { createCohere } from '../cohere-provider';
+import { describe, it, expect, beforeEach } from 'vitest';
+import fs from 'node:fs';
+import { CohereRerankingOptions } from './cohere-reranking-options';
+
+const provider = createCohere({ apiKey: 'test-api-key' });
+const model = provider.rerankingModel('rerank-english-v3.0');
+
+describe('doRerank', () => {
+  const server = createTestServer({
+    'https://api.cohere.com/v2/rerank': {},
+  });
+
+  function prepareJsonFixtureResponse(filename: string) {
+    server.urls['https://api.cohere.com/v2/rerank'].response = {
+      type: 'json-value',
+      body: JSON.parse(
+        fs.readFileSync(`src/reranking/__fixtures__/${filename}.json`, 'utf8'),
+      ),
+    };
+    return;
+  }
+
+  describe('json documents', () => {
+    let result: Awaited<ReturnType<typeof model.doRerank>>;
+
+    beforeEach(async () => {
+      prepareJsonFixtureResponse('cohere-reranking.1');
+
+      result = await model.doRerank({
+        documents: {
+          type: 'object',
+          values: [
+            { example: 'sunny day at the beach' },
+            { example: 'rainy day in the city' },
+          ],
+        },
+        query: 'rainy day',
+        topN: 2,
+        providerOptions: {
+          cohere: {
+            maxTokensPerDoc: 1000,
+            priority: 1,
+          } satisfies CohereRerankingOptions,
+        },
+      });
+    });
+
+    it('should send request with stringified json documents', async () => {
+      expect(await server.calls[0].requestBodyJson).toMatchInlineSnapshot(`
+        {
+          "documents": [
+            "{"example":"sunny day at the beach"}",
+            "{"example":"rainy day in the city"}",
+          ],
+          "max_tokens_per_doc": 1000,
+          "model": "rerank-english-v3.0",
+          "priority": 1,
+          "query": "rainy day",
+          "top_n": 2,
+        }
+      `);
+    });
+
+    it('should send request with the correct headers', async () => {
+      expect(server.calls[0].requestHeaders).toMatchInlineSnapshot(`
+        {
+          "authorization": "Bearer test-api-key",
+          "content-type": "application/json",
+        }
+      `);
+    });
+
+    it('should return result with warnings', async () => {
+      expect(result.warnings).toMatchInlineSnapshot(`
+        [
+          {
+            "details": "Object documents are converted to strings.",
+            "feature": "object documents",
+            "type": "compatibility",
+          },
+        ]
+      `);
+    });
+
+    it('should return result with the correct ranking', async () => {
+      expect(result.ranking).toMatchInlineSnapshot(`
+        [
+          {
+            "index": 1,
+            "relevanceScore": 0.10183054,
+          },
+          {
+            "index": 0,
+            "relevanceScore": 0.03762639,
+          },
+        ]
+      `);
+    });
+
+    it('should not return provider metadata (use response body instead)', async () => {
+      expect(result.providerMetadata).toMatchInlineSnapshot(`undefined`);
+    });
+
+    it('should return result with the correct response', async () => {
+      expect(result.response).toMatchInlineSnapshot(`
+        {
+          "body": {
+            "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72",
+            "meta": {
+              "api_version": {
+                "version": "2",
+              },
+              "billed_units": {
+                "search_units": 1,
+              },
+            },
+            "results": [
+              {
+                "index": 1,
+                "relevance_score": 0.10183054,
+              },
+              {
+                "index": 0,
+                "relevance_score": 0.03762639,
+              },
+            ],
+          },
+          "headers": {
+            "content-length": "212",
+            "content-type": "application/json",
+          },
+          "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72",
+        }
+      `);
+    });
+  });
+
+  describe('text documents', () => {
+    let result: Awaited<ReturnType<typeof model.doRerank>>;
+
+    beforeEach(async () => {
+      prepareJsonFixtureResponse('cohere-reranking.1');
+
+      result = await model.doRerank({
+        documents: {
+          type: 'text',
+          values: ['sunny day at the beach', 'rainy day in the city'],
+        },
+        query: 'rainy day',
+        topN: 2,
+        providerOptions: {
+          cohere: {
+            maxTokensPerDoc: 1000,
+            priority: 1,
+          } satisfies CohereRerankingOptions,
+        },
+      });
+    });
+
+    it('should send request with text documents', async () => {
+      expect(await server.calls[0].requestBodyJson).toMatchInlineSnapshot(`
+        {
+          "documents": [
+            "sunny day at the beach",
+            "rainy day in the city",
+          ],
+          "max_tokens_per_doc": 1000,
+          "model": "rerank-english-v3.0",
+          "priority": 1,
+          "query": "rainy day",
+          "top_n": 2,
+        }
+      `);
+    });
+
+    it('should send request with the correct headers', async () => {
+      expect(server.calls[0].requestHeaders).toMatchInlineSnapshot(`
+        {
+          "authorization": "Bearer test-api-key",
+          "content-type": "application/json",
+        }
+      `);
+    });
+
+    it('should return result without warnings', async () => {
+      expect(result.warnings).toMatchInlineSnapshot(`[]`);
+    });
+
+    it('should return result with the correct ranking', async () => {
+      expect(result.ranking).toMatchInlineSnapshot(`
+        [
+          {
+            "index": 1,
+            "relevanceScore": 0.10183054,
+          },
+          {
+            "index": 0,
+            "relevanceScore": 0.03762639,
+          },
+        ]
+      `);
+    });
+
+    it('should not return provider metadata (use response body instead)', async () => {
+      expect(result.providerMetadata).toMatchInlineSnapshot(`undefined`);
+    });
+
+    it('should return result with the correct response', async () => {
+      expect(result.response).toMatchInlineSnapshot(`
+        {
+          "body": {
+            "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72",
+            "meta": {
+              "api_version": {
+                "version": "2",
+              },
+              "billed_units": {
+                "search_units": 1,
+              },
+            },
+            "results": [
+              {
+                "index": 1,
+                "relevance_score": 0.10183054,
+              },
+              {
+                "index": 0,
+                "relevance_score": 0.03762639,
+              },
+            ],
+          },
+          "headers": {
+            "content-length": "212",
+            "content-type": "application/json",
+          },
+          "id": "b44fe75b-e3d3-489a-b61e-1a1aede3ef72",
+        }
+      `);
+    });
+  });
+});

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

@@ -0,0 +1,107 @@
+import { RerankingModelV3, SharedV3Warning } from '@ai-sdk/provider';
+import {
+  combineHeaders,
+  createJsonResponseHandler,
+  FetchFunction,
+  parseProviderOptions,
+  postJsonToApi,
+} from '@ai-sdk/provider-utils';
+import { cohereFailedResponseHandler } from '../cohere-error';
+import {
+  CohereRerankingInput,
+  cohereRerankingResponseSchema,
+} from './cohere-reranking-api';
+import {
+  CohereRerankingModelId,
+  cohereRerankingOptionsSchema,
+} from './cohere-reranking-options';
+
+type CohereRerankingConfig = {
+  provider: string;
+  baseURL: string;
+  headers: () => Record<string, string | undefined>;
+  fetch?: FetchFunction;
+};
+
+export class CohereRerankingModel implements RerankingModelV3 {
+  readonly specificationVersion = 'v3';
+  readonly modelId: CohereRerankingModelId;
+
+  private readonly config: CohereRerankingConfig;
+
+  constructor(modelId: CohereRerankingModelId, config: CohereRerankingConfig) {
+    this.modelId = modelId;
+    this.config = config;
+  }
+
+  get provider(): string {
+    return this.config.provider;
+  }
+
+  // current implementation is based on v2 of the API: https://docs.cohere.com/v2/reference/rerank
+  async doRerank({
+    documents,
+    headers,
+    query,
+    topN,
+    abortSignal,
+    providerOptions,
+  }: Parameters<RerankingModelV3['doRerank']>[0]): Promise<
+    Awaited<ReturnType<RerankingModelV3['doRerank']>>
+  > {
+    const rerankingOptions = await parseProviderOptions({
+      provider: 'cohere',
+      providerOptions,
+      schema: cohereRerankingOptionsSchema,
+    });
+
+    const warnings: SharedV3Warning[] = [];
+
+    if (documents.type === 'object') {
+      warnings.push({
+        type: 'compatibility',
+        feature: 'object documents',
+        details: 'Object documents are converted to strings.',
+      });
+    }
+
+    const {
+      responseHeaders,
+      value: response,
+      rawValue,
+    } = await postJsonToApi({
+      url: `${this.config.baseURL}/rerank`,
+      headers: combineHeaders(this.config.headers(), headers),
+      body: {
+        model: this.modelId,
+        query,
+        documents:
+          documents.type === 'text'
+            ? documents.values
+            : documents.values.map(value => JSON.stringify(value)),
+        top_n: topN,
+        max_tokens_per_doc: rerankingOptions?.maxTokensPerDoc,
+        priority: rerankingOptions?.priority,
+      } satisfies CohereRerankingInput,
+      failedResponseHandler: cohereFailedResponseHandler,
+      successfulResponseHandler: createJsonResponseHandler(
+        cohereRerankingResponseSchema,
+      ),
+      abortSignal,
+      fetch: this.config.fetch,
+    });
+
+    return {
+      ranking: response.results.map(result => ({
+        index: result.index,
+        relevanceScore: result.relevance_score,
+      })),
+      warnings,
+      response: {
+        id: response.id ?? undefined,
+        headers: responseHeaders,
+        body: rawValue,
+      },
+    };
+  }
+}

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

@@ -0,0 +1,35 @@
+import { FlexibleSchema, lazySchema, zodSchema } from '@ai-sdk/provider-utils';
+import { z } from 'zod/v4';
+
+// https://docs.cohere.com/docs/rerank
+export type CohereRerankingModelId =
+  | 'rerank-v3.5'
+  | 'rerank-english-v3.0'
+  | 'rerank-multilingual-v3.0'
+  | (string & {});
+
+export type CohereRerankingOptions = {
+  /**
+   * Long documents will be automatically truncated to the specified number of tokens.
+   *
+   * @default 4096
+   */
+  maxTokensPerDoc?: number;
+
+  /**
+   * The priority of the request.
+   *
+   * @default 0
+   */
+  priority?: number;
+};
+
+export const cohereRerankingOptionsSchema: FlexibleSchema<CohereRerankingOptions> =
+  lazySchema(() =>
+    zodSchema(
+      z.object({
+        maxTokensPerDoc: z.number().optional(),
+        priority: z.number().optional(),
+      }),
+    ),
+  );

package/src/version.ts

@@ -0,0 +1,6 @@
+// Version string of this package injected at build time.
+declare const __PACKAGE_VERSION__: string | undefined;
+export const VERSION: string =
+  typeof __PACKAGE_VERSION__ !== 'undefined'
+    ? __PACKAGE_VERSION__
+    : '0.0.0-test';