Expose native web search and image generation in the SDK#294
Merged
Conversation
The TEE gateway recently added an opt-in `web_search` flag and native
Gemini image generation. This surfaces both through the LLM client.
Web search:
- `LLM.chat` and `LLM.completion` gain a `web_search: bool = False` arg that
forwards `web_search` to the gateway when enabled. Supported by OpenAI,
Anthropic, Google, and xAI models; billed per search on top of tokens.
- LangChain adapter forwards `web_search` (constructor + per-call kwarg).
- CLI `chat`/`completion` gain a `--web-search` flag.
Image generation:
- Add `GEMINI_2_5_FLASH_IMAGE` and `GEMINI_3_1_FLASH_IMAGE` ("nano banana")
plus `GEMINI_3_5_FLASH` to the `TEE_LLM` enum.
- Surface generated images on `TextGenerationOutput.images` (non-streaming)
and `StreamChunk.images` (final stream chunk) as `data:` URIs.
- CLI saves generated images to disk (`--image-output-dir`) instead of
dumping base64.
Adds examples (`llm_web_search.py`, `llm_image_generation.py`), README and
docstring updates, and unit tests for payload forwarding and image surfacing.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
The TEE gateway recently added an opt-in
web_searchflag and native Gemini image generation. This PR surfaces both capabilities through the SDK'sLLMclient (plus the LangChain adapter and CLI). No gateway changes are needed — this is purely client-side wiring.Web search
Opt-in
web_searchflag, billed per search on top of token usage. Supported by OpenAI, Anthropic, Google, and xAI models; other providers ignore the flag.LLM.chat()andLLM.completion()gain aweb_search: bool = Falseargument that forwardsweb_search: trueto the gateway when enabled.OpenGradientChatModel) forwards it via both the constructor and per-call kwargs.chat/completiongain a--web-searchflag.Image generation ("nano banana")
GEMINI_2_5_FLASH_IMAGE,GEMINI_3_1_FLASH_IMAGE, andGEMINI_3_5_FLASHto theTEE_LLMenum.data:URIs onTextGenerationOutput.images(non-streaming) andStreamChunk.images(final stream chunk).--image-output-dir) instead of dumping base64.Images travel out-of-band and are not part of the signed TEE output hash, matching the gateway's existing contract (documented in the docstrings).
Supporting work
examples/llm_web_search.pyandexamples/llm_image_generation.py, both referenced inexamples/README.md.Verification
uv run pytest tests/→ 131 passeduv run mypy srcanduv run mypy examples→ cleanruff formatcleanhttps://claude.ai/code/session_015BV6hef5Ld82rbov858L58
Generated by Claude Code