Skip to content

Conversation

@agcty
Copy link
Contributor

@agcty agcty commented Jan 20, 2026

Summary

Adds embedding model support to ai-gateway-provider, allowing users to route embedding requests through AI Gateway.

Changes

  • Add AiGatewayEmbeddingModel class implementing EmbeddingModelV3
  • Add embedding(), textEmbedding(), and textEmbeddingModel() methods to the provider factory
  • Extract shared types and utilities to shared.ts to avoid circular imports
  • Add comprehensive test coverage for embedding functionality

Usage

import { createAiGateway } from "ai-gateway-provider";
import { createOpenAI } from "@ai-sdk/openai";
import { embed } from "ai";

const gateway = createAiGateway({
  accountId: "your-account-id",
  gateway: "your-gateway",
  apiKey: "your-api-key",
});

const openai = createOpenAI({ apiKey: "..." });

const { embedding } = await embed({
  model: gateway.embedding(openai.textEmbeddingModel("text-embedding-3-small")),
  value: "Hello, world!",
});

Test plan

  • All existing tests pass (13 tests)
  • New embedding tests pass (6 tests)
  • Type checking passes
  • Changeset included

🤖 Generated with Claude Code

Add AiGatewayEmbeddingModel class that wraps underlying embedding models
and routes requests through AI Gateway, similar to the chat model.

- Create new AiGatewayEmbeddingModel implementing EmbeddingModelV3
- Add embedding(), textEmbedding(), and textEmbeddingModel() methods to AiGateway interface
- Extract shared types, errors, and utilities to shared.ts to avoid circular imports
- Add comprehensive test coverage for embedding functionality
- Support all gateway options (caching, logging, retries, metadata)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@changeset-bot
Copy link

changeset-bot bot commented Jan 20, 2026

🦋 Changeset detected

Latest commit: 744b7a3

The changes in this PR will be included in the next version bump.

This PR includes changesets to release 1 package
Name Type
ai-gateway-provider Minor

Not sure what this means? Click here to learn what changesets are.

Click here if you're a maintainer who wants to add another changeset to this PR

agcty and others added 2 commits January 20, 2026 23:27
Add tests for AiGatewayDoesNotExist (400/2001) and AiGatewayUnauthorizedError
(401/2009) error cases to match the chat model test coverage.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
The EmbeddingModelV3Result requires a `warnings` array (not optional).
Explicitly construct the return object to ensure V3 compliance regardless
of what the underlying model returns.

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>
@agcty
Copy link
Contributor Author

agcty commented Jan 21, 2026

Testing Notes

I tested this implementation locally by routing Google embedding requests through AI Gateway. The routing works correctly - requests are intercepted, reformatted for the Universal Endpoint, and return successful responses.

Debug logs from local testing (not included in final code):

[ai-gateway-embedding] doEmbed called, models: 1
[ai-gateway-embedding] Processing model: google.generative-ai gemini-embedding-001
[ai-gateway-embedding] FETCH INTERCEPTED: https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:embedContent
[ai-gateway-embedding] Requests captured: 1
[ai-gateway-embedding] Request body: [
  {
    "endpoint": "v1beta/models/gemini-embedding-001:embedContent",
    "headers": {
      "content-type": "application/json",
      "x-goog-api-key": "..."
    },
    "provider": "google-ai-studio",
    "query": {
      "model": "models/gemini-embedding-001",
      "content": { "parts": [{ "text": "..." }] },
      "outputDimensionality": 1024
    }
  }
]
[ai-gateway-embedding] Sending to AI Gateway: https://gateway.ai.cloudflare.com/v1/{accountId}/test
[ai-gateway-embedding] AI Gateway response status: 200

Note: Embedding requests don't appear in the AI Gateway dashboard, though they route correctly and return 200 responses. This may require backend support to display embedding analytics alongside chat completions.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant