Skip to content

Commit e31db39

Browse files
Loule95450Loule95450
authored
feat: add first-class embedding model support for OpenRouter (v4 & v5) (#236)
* Add OpenRouter embedding model implementation and tests Introduces the OpenRouterEmbeddingModel class for embedding generation, including schema validation and API integration. Adds comprehensive tests for model instantiation, embedding functionality, custom settings, and response handling. * Add OpenRouter embedding model support Introduces text embedding model support to the OpenRouter provider, including a new textEmbeddingModel method and a deprecated embedding alias for backward compatibility. Updates both the facade and provider to expose and implement these methods. * Add OpenRouter embedding settings types Introduced new type definitions for OpenRouter embedding settings in src/types/openrouter-embedding-settings.ts and re-exported them from index.ts. These types provide configuration options for embedding model requests, including provider routing preferences and user identification. * Add documentation for embedding model support Expanded the README with details on using embedding models with OpenRouter, including usage examples for AI SDK v5 and v4, batch embeddings, and a list of supported embedding models.
1 parent 55f4aee commit e31db39

File tree

8 files changed

+562
-0
lines changed

8 files changed

+562
-0
lines changed

‎README.md‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,6 +55,64 @@ This list is not a definitive list of models supported by OpenRouter, as it cons
5555

5656
You can find the latest list of tool-supported models supported by OpenRouter [here](https://openrouter.ai/models?order=newest&supported_parameters=tools). (Note: This list may contain models that are not compatible with the AI SDK.)
5757

58+
## Embeddings
59+
60+
OpenRouter supports embedding models for semantic search, RAG pipelines, and vector-native features. The provider exposes embeddings compatible with both AI SDK v5 and v4.
61+
62+
### AI SDK v5 (Recommended)
63+
64+
```ts
65+
import { embed } from 'ai';
66+
import { openrouter } from '@openrouter/ai-sdk-provider';
67+
68+
const { embedding } = await embed({
69+
model: openrouter.textEmbeddingModel('openai/text-embedding-3-small'),
70+
value: 'sunny day at the beach',
71+
});
72+
73+
console.log(embedding); // Array of numbers representing the embedding
74+
```
75+
76+
### Batch Embeddings
77+
78+
```ts
79+
import { embedMany } from 'ai';
80+
import { openrouter } from '@openrouter/ai-sdk-provider';
81+
82+
const { embeddings } = await embedMany({
83+
model: openrouter.textEmbeddingModel('openai/text-embedding-3-small'),
84+
values: [
85+
'sunny day at the beach',
86+
'rainy day in the city',
87+
'snowy mountain peak',
88+
],
89+
});
90+
91+
console.log(embeddings); // Array of embedding arrays
92+
```
93+
94+
### AI SDK v4 (Deprecated)
95+
96+
For backwards compatibility, the `embedding` method is also available:
97+
98+
```ts
99+
import { embed } from 'ai';
100+
import { openrouter } from '@openrouter/ai-sdk-provider';
101+
102+
const { embedding } = await embed({
103+
model: openrouter.embedding('openai/text-embedding-3-small'),
104+
value: 'sunny day at the beach',
105+
});
106+
```
107+
108+
### Supported Embedding Models
109+
110+
OpenRouter supports various embedding models including:
111+
- `openai/text-embedding-3-small`
112+
- `openai/text-embedding-3-large`
113+
- `openai/text-embedding-ada-002`
114+
- And more available on [OpenRouter](https://openrouter.ai/models?output_modalities=embeddings)
115+
58116
## Passing Extra Body to OpenRouter
59117

60118
There are 3 ways to pass extra body to OpenRouter:

‎src/embedding/index.test.ts‎

Lines changed: 253 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,253 @@
1+
import { describe, expect, it } from 'vitest';
2+
import { createOpenRouter } from '../provider';
3+
import { OpenRouterEmbeddingModel } from './index';
4+
5+
describe('OpenRouterEmbeddingModel', () => {
6+
const mockFetch = async (
7+
_url: URL | RequestInfo,
8+
_init?: RequestInit,
9+
): Promise<Response> => {
10+
return new Response(
11+
JSON.stringify({
12+
id: 'test-id',
13+
object: 'list',
14+
data: [
15+
{
16+
object: 'embedding',
17+
embedding: new Array(1536).fill(0.1),
18+
index: 0,
19+
},
20+
],
21+
model: 'openai/text-embedding-3-small',
22+
usage: {
23+
prompt_tokens: 5,
24+
total_tokens: 5,
25+
cost: 0.00001,
26+
},
27+
}),
28+
{
29+
status: 200,
30+
headers: {
31+
'content-type': 'application/json',
32+
},
33+
},
34+
);
35+
};
36+
37+
describe('provider methods', () => {
38+
it('should expose textEmbeddingModel method', () => {
39+
const provider = createOpenRouter({ apiKey: 'test-key' });
40+
expect(provider.textEmbeddingModel).toBeDefined();
41+
expect(typeof provider.textEmbeddingModel).toBe('function');
42+
});
43+
44+
it('should expose embedding method (deprecated)', () => {
45+
const provider = createOpenRouter({ apiKey: 'test-key' });
46+
expect(provider.embedding).toBeDefined();
47+
expect(typeof provider.embedding).toBe('function');
48+
});
49+
50+
it('should create an embedding model instance', () => {
51+
const provider = createOpenRouter({ apiKey: 'test-key' });
52+
const model = provider.textEmbeddingModel(
53+
'openai/text-embedding-3-small',
54+
);
55+
expect(model).toBeInstanceOf(OpenRouterEmbeddingModel);
56+
expect(model.modelId).toBe('openai/text-embedding-3-small');
57+
expect(model.provider).toBe('openrouter');
58+
expect(model.specificationVersion).toBe('v2');
59+
});
60+
});
61+
62+
describe('doEmbed', () => {
63+
it('should embed a single value', async () => {
64+
const provider = createOpenRouter({
65+
apiKey: 'test-key',
66+
fetch: mockFetch,
67+
});
68+
const model = provider.textEmbeddingModel(
69+
'openai/text-embedding-3-small',
70+
);
71+
72+
const result = await model.doEmbed({
73+
values: ['sunny day at the beach'],
74+
});
75+
76+
expect(result.embeddings).toHaveLength(1);
77+
expect(result.embeddings[0]).toHaveLength(1536);
78+
expect(result.usage).toEqual({ tokens: 5 });
79+
expect(
80+
(result.providerMetadata?.openrouter as { usage?: { cost?: number } })
81+
?.usage?.cost,
82+
).toBe(0.00001);
83+
});
84+
85+
it('should embed multiple values', async () => {
86+
const mockFetchMultiple = async (
87+
_url: URL | RequestInfo,
88+
_init?: RequestInit,
89+
): Promise<Response> => {
90+
return new Response(
91+
JSON.stringify({
92+
object: 'list',
93+
data: [
94+
{
95+
object: 'embedding',
96+
embedding: new Array(1536).fill(0.1),
97+
index: 0,
98+
},
99+
{
100+
object: 'embedding',
101+
embedding: new Array(1536).fill(0.2),
102+
index: 1,
103+
},
104+
{
105+
object: 'embedding',
106+
embedding: new Array(1536).fill(0.3),
107+
index: 2,
108+
},
109+
],
110+
model: 'openai/text-embedding-3-small',
111+
usage: {
112+
prompt_tokens: 15,
113+
total_tokens: 15,
114+
},
115+
}),
116+
{
117+
status: 200,
118+
headers: {
119+
'content-type': 'application/json',
120+
},
121+
},
122+
);
123+
};
124+
125+
const provider = createOpenRouter({
126+
apiKey: 'test-key',
127+
fetch: mockFetchMultiple,
128+
});
129+
const model = provider.textEmbeddingModel(
130+
'openai/text-embedding-3-small',
131+
);
132+
133+
const result = await model.doEmbed({
134+
values: [
135+
'sunny day at the beach',
136+
'rainy day in the city',
137+
'snowy mountain peak',
138+
],
139+
});
140+
141+
expect(result.embeddings).toHaveLength(3);
142+
expect(result.embeddings[0]).toHaveLength(1536);
143+
expect(result.embeddings[1]).toHaveLength(1536);
144+
expect(result.embeddings[2]).toHaveLength(1536);
145+
expect(result.usage).toEqual({ tokens: 15 });
146+
});
147+
148+
it('should pass custom settings to API', async () => {
149+
let capturedRequest: Record<string, unknown> | undefined;
150+
151+
const mockFetchWithCapture = async (
152+
_url: URL | RequestInfo,
153+
init?: RequestInit,
154+
): Promise<Response> => {
155+
capturedRequest = JSON.parse(init?.body as string);
156+
return new Response(
157+
JSON.stringify({
158+
object: 'list',
159+
data: [
160+
{
161+
object: 'embedding',
162+
embedding: new Array(1536).fill(0.1),
163+
index: 0,
164+
},
165+
],
166+
model: 'openai/text-embedding-3-small',
167+
usage: {
168+
prompt_tokens: 5,
169+
total_tokens: 5,
170+
},
171+
}),
172+
{
173+
status: 200,
174+
headers: {
175+
'content-type': 'application/json',
176+
},
177+
},
178+
);
179+
};
180+
181+
const provider = createOpenRouter({
182+
apiKey: 'test-key',
183+
fetch: mockFetchWithCapture,
184+
});
185+
186+
const model = provider.textEmbeddingModel(
187+
'openai/text-embedding-3-small',
188+
{
189+
user: 'test-user-123',
190+
provider: {
191+
order: ['openai'],
192+
allow_fallbacks: false,
193+
},
194+
},
195+
);
196+
197+
await model.doEmbed({
198+
values: ['test input'],
199+
});
200+
201+
expect(capturedRequest?.user).toBe('test-user-123');
202+
expect(capturedRequest?.provider).toEqual({
203+
order: ['openai'],
204+
allow_fallbacks: false,
205+
});
206+
expect(capturedRequest?.model).toBe('openai/text-embedding-3-small');
207+
expect(capturedRequest?.input).toEqual(['test input']);
208+
});
209+
210+
it('should handle response without usage information', async () => {
211+
const mockFetchNoUsage = async (
212+
_url: URL | RequestInfo,
213+
_init?: RequestInit,
214+
): Promise<Response> => {
215+
return new Response(
216+
JSON.stringify({
217+
object: 'list',
218+
data: [
219+
{
220+
object: 'embedding',
221+
embedding: new Array(1536).fill(0.1),
222+
index: 0,
223+
},
224+
],
225+
model: 'openai/text-embedding-3-small',
226+
}),
227+
{
228+
status: 200,
229+
headers: {
230+
'content-type': 'application/json',
231+
},
232+
},
233+
);
234+
};
235+
236+
const provider = createOpenRouter({
237+
apiKey: 'test-key',
238+
fetch: mockFetchNoUsage,
239+
});
240+
const model = provider.textEmbeddingModel(
241+
'openai/text-embedding-3-small',
242+
);
243+
244+
const result = await model.doEmbed({
245+
values: ['test'],
246+
});
247+
248+
expect(result.embeddings).toHaveLength(1);
249+
expect(result.usage).toBeUndefined();
250+
expect(result.providerMetadata).toBeUndefined();
251+
});
252+
});
253+
});

0 commit comments

Comments
 (0)