Skip to content

Support VertexAI embedding task_type and title - #824

Closed
goianiense wants to merge 1 commit into
crmne:mainfrom
goianiense:fix/810-vertexai-embedding-task-type
Closed

goianiense wants to merge 1 commit into
crmne:mainfrom
goianiense:fix/810-vertexai-embedding-task-type

Conversation

@goianiense

Copy link
Copy Markdown

Summary

  • Plumbs provider-specific kwargs (**provider_params) through RubyLLM.embed → Provider#embed → Protocol#embed → render_embedding_payload
  • VertexAI embedding payloads can now include task_type (e.g. SEMANTIC_SIMILARITY, RETRIEVAL_QUERY, RETRIEVAL_DOCUMENT) and optional title on each instance
  • Other embedding providers ignore unknown kwargs via ** splat so extra options don't raise

Fixes #810

Usage

RubyLLM.embed(
  "hello",
  model: "gemini-embedding-001",
  provider: :vertexai,
  dimensions: 1536,
  task_type: "SEMANTIC_SIMILARITY"
)

Produces a VertexAI payload shaped like:

{
  "instances": [
    {
      "content": "hello",
      "task_type": "SEMANTIC_SIMILARITY"
    }
  ],
  "parameters": {
    "outputDimensionality": 1536
  }
}

For document embeddings:

RubyLLM.embed(
  "document body",
  model: "gemini-embedding-001",
  provider: :vertexai,
  task_type: "RETRIEVAL_DOCUMENT",
  title: "My Document"
)

Test plan

  • Added unit specs for VertexAI::Embeddings#render_embedding_payload covering task_type, title, batch instances, and dimensions
  • Added specs proving RubyLLM.embed forwards task_type/title to the provider (reproduced ArgumentError: unknown keyword: :task_type before the fix)
  • All new specs pass locally

Plumb provider-specific kwargs through RubyLLM.embed so VertexAI can
send task_type (and title for RETRIEVAL_DOCUMENT) on gemini-embedding-001
instances, matching the Vertex predict API shape.

Closes #810
@crmne

crmne commented Jul 3, 2026 •

Copy link
Copy Markdown
Owner

Thanks for the PR.

The underlying Vertex AI embedding issue is fixed on main, but through the general one-shot params API rather than provider-specific top-level keywords. The fix landed in 1bc6fc0, with the follow-up keyword cleanup in 5cdf01f.

Use params for Vertex AI request fields:

RubyLLM.embed(
  "hello",
  model: "gemini-embedding-001",
  provider: :vertexai,
  dimensions: 1536,
  params: { task_type: "SEMANTIC_SIMILARITY" }
)

For document embeddings, params can include both task_type and title. RubyLLM applies those fields to each Vertex AI embedding instance.

Closing this PR as superseded by the implementation on main.

@crmne crmne closed this Jul 3, 2026
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.

[BUG] VertexAI embeddings cannot send task_type for gemini-embedding-001

2 participants