Pool tokens for embed image

[Jurgee]

Summary by CodeRabbit

• Documentation

  • Added documentation describing how images are serialized, compressed, and sent for embedding, and how embedding responses are returned and reshaped.

• New Features

  • embed_image now accepts additional parameters that are transmitted as custom request headers.

[coderabbitai]

No actionable comments were generated in the recent review. 🎉

ℹ️ Recent review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: a77a96f0-2cec-4b75-8c0b-935a76571d58

📥 Commits

Reviewing files that changed from the base of the PR and between 2e9636e and b5cef7e.

📒 Files selected for processing (1)

• rationai/resources/models.py

🚧 Files skipped from review as they are similar to previous changes (1)

• rationai/resources/models.py

---

📝 Walkthrough

Walkthrough

Documentation and code updates add explicit embed_image protocol details (image.tobytes(), LZ4 frame, x-output-dtype, x-output-shape) and extend Models.embed_image and AsyncModels.embed_image to accept **headers which are converted to hyphenated x-* request headers.

Changes

Cohort / File(s)	Summary
Documentation docs/learn/how-it-works.md	Added embed_image section documenting request payload (image bytes, LZ4 compression, x-output-dtype, keyword args -> x-* headers) and response format (LZ4-compressed binary with x-output-shape header).
Embed Image Methods rationai/resources/models.py	Updated Models.embed_image and AsyncModels.embed_image signatures to accept **headers: str. Request header construction now merges x-output-dtype with caller-provided headers after underscore->hyphen conversion and x- prefixing; _post calls receive the composed headers.

Sequence Diagram(s)

(omitted)

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Possibly related PRs

• Embed image for Virchow2 #6: Modified the same embed_image methods and initial embed_image implementation; directly related.

Suggested reviewers

• matejpekar

Poem

🐰 I pack bytes tight and zip with delight,
Underscores hop, become dashes at night,
Headers wear x- like a jaunty new hat,
Images travel snug, coming back as a flat —
Hop, skip, compress — embed done just like that! 🎉

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 inconclusive)

Check name	Status	Explanation	Resolution
Title check	❓ Inconclusive	The PR title 'Pool tokens for embed image' is vague and doesn't clearly describe the actual changes, which involve adding **headers parameter support to embed_image methods and updating documentation.	Consider using a more descriptive title that clearly indicates the main change, such as 'Add custom headers support to embed_image methods' or 'Allow custom x-* headers in embed_image requests'.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Docstring Coverage	✅ Passed	Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

📝 Generate docstrings

• Create stacked PR
• Commit on current branch

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Commit unit tests in branch feature/tokens

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces support for custom headers in the embed_image method for both synchronous and asynchronous clients. It updates the documentation to explain the serialization and compression process, and modifies the implementation to convert keyword arguments into x- prefixed headers. The review feedback correctly identifies that the docstring examples in models.py are inconsistent with the code logic, as they omit the x- prefix in the illustrated output.

[Copilot]

Pull request overview

Adds support for passing additional x-* request headers to the embed_image model endpoint, enabling server-side options (e.g., token pooling) to be controlled from the SDK.

Changes:

• Extend Models.embed_image / AsyncModels.embed_image to accept **headers and send them as x-* request headers (underscores → hyphens).
• Update the “How it works” documentation to describe embed-image request/response behavior, including the new header mechanism and x-output-shape.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
rationai/resources/models.py	Adds **headers support to embed-image requests and constructs a merged request headers dict.
docs/learn/how-it-works.md	Documents embed-image wire format, including x-output-dtype, optional x-* keyword headers, and x-output-shape.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (1)

rationai/resources/models.py (1)

91-94: Consider extracting header-building into a shared helper.

The sync/async methods duplicate the same header normalization block; centralizing it reduces drift risk.

Also applies to: 190-193

🤖 Prompt for AI Agents

Verify each finding against the current code and only fix it if needed.

In `@rationai/resources/models.py` around lines 91 - 94, The header normalization
logic that builds request_headers (using np.dtype(output_dtype).name and mapping
headers via f"x-{k.replace('_', '-')"} is duplicated between the sync and async
request methods; extract this into a single helper function (e.g.,
build_request_headers or normalize_request_headers) that takes output_dtype and
headers and returns the normalized dict, then replace the inline blocks in the
methods that currently assign request_headers with calls to that helper (ensure
the helper imports numpy as np or accepts already-normalized dtype string). This
centralizes the logic and removes duplication referenced by the request_headers
usage in the existing methods.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Inline comments:
In `@rationai/resources/models.py`:
- Around line 82-84: Update the docstring examples that describe header
handling: change the example showing "pool-tokens: false" to include the x-
prefix (e.g., "x-pool-tokens: false") so it matches the implementation which
emits "x-" + header-name and converts underscores to hyphens; locate the
docstrings that contain the text "**headers: Additional x- headers. Keyword
underscores are converted to hyphens..." and update both occurrences (the one
around the shown diff and the other similar block) so the example and
explanation accurately reflect "x-" prefix behavior.
- Around line 91-94: The code that builds request_headers (using request_headers
= {"x-output-dtype": np.dtype(output_dtype).name} followed by
{f"x-{k.replace('_', '-')}": v for k, v in headers.items()}) should validate
custom header keys to avoid producing accidental "x-x-*" headers: before
constructing the prefixed header names, check each incoming key in headers and
raise a clear ValueError if key.lower().startswith("x") (covers "x_", "x-" and
"x...") so callers must not include the "x" prefix; apply the same validation
and behavior to the other identical block referenced around lines 190-193, and
keep the existing underscore-to-dash replacement for valid keys.

---

Nitpick comments:
In `@rationai/resources/models.py`:
- Around line 91-94: The header normalization logic that builds request_headers
(using np.dtype(output_dtype).name and mapping headers via f"x-{k.replace('_',
'-')"} is duplicated between the sync and async request methods; extract this
into a single helper function (e.g., build_request_headers or
normalize_request_headers) that takes output_dtype and headers and returns the
normalized dict, then replace the inline blocks in the methods that currently
assign request_headers with calls to that helper (ensure the helper imports
numpy as np or accepts already-normalized dtype string). This centralizes the
logic and removes duplication referenced by the request_headers usage in the
existing methods.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 14e30c20-3c12-4a6a-93db-880df62cb623

📥 Commits

Reviewing files that changed from the base of the PR and between 5e2d64c and 2e9636e.

📒 Files selected for processing (2)

• docs/learn/how-it-works.md
• rationai/resources/models.py