Expose native web search and image generation in the SDK (#294)

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.

Co-authored-by: Claude <[email protected]>

Author: adambalogh

README.md

@@ -46,6 +46,8 @@ OpenGradient enables developers to build AI applications with verifiable executi
 
 - **Verifiable LLM Inference**: Drop-in replacement for OpenAI and Anthropic APIs with cryptographic attestation
 - **Multi-Provider Support**: Access models from OpenAI, Anthropic, Google, and xAI through a unified interface
+- **Native Web Search**: Opt-in `web_search` flag enables each provider's built-in web search, billed per search
+- **Image Generation**: Native image-output models ("nano banana") return generated images directly on the response
 - **TEE Execution**: Trusted Execution Environment inference with cryptographic verification
 - **Model Hub Integration**: Registry for model discovery, versioning, and deployment
 - **Consensus-Based Verification**: End-to-end verified AI execution through the OpenGradient network
@@ -168,6 +170,37 @@ async for chunk in stream:
         print(chunk.choices[0].delta.content, end="")
 ```
 
+### Native Web Search
+
+Set `web_search=True` to let the model search the web while answering. Each search is billed per search on top of token usage, at the provider's list price. Supported by OpenAI, Anthropic, Google, and xAI models; other providers ignore the flag.
+```python
+completion = await llm.chat(
+    model=og.TEE_LLM.CLAUDE_SONNET_4_6,
+    messages=[{"role": "user", "content": "What are today's top tech headlines?"}],
+    max_tokens=500,
+    web_search=True,
+)
+print(completion.chat_output["content"])
+```
+
+### Image Generation
+
+Native image-output models ("nano banana") return generated images on the response. The generated images are available in `result.images` as `data:` URIs, while any text caption is in `chat_output["content"]`. Images travel out-of-band and are not part of the signed output hash.
+```python
+import base64
+
+result = await llm.chat(
+    model=og.TEE_LLM.GEMINI_3_1_FLASH_IMAGE,
+    messages=[{"role": "user", "content": "A friendly robot reading under a tree"}],
+    max_tokens=1024,
+)
+
+for i, image in enumerate(result.images or []):
+    payload = image.split(",", 1)[1]  # strip the "data:image/png;base64," prefix
+    with open(f"image_{i}.png", "wb") as f:
+        f.write(base64.b64decode(payload))
+```
+
 ### Verifiable LangChain Integration
 
 Use OpenGradient as a drop-in LLM provider for LangChain agents with network-verified execution:
@@ -219,6 +252,9 @@ The SDK provides access to models from multiple providers via the `og.TEE_LLM` e
 - Gemini 2.5 Flash Lite
 - Gemini 3 Pro
 - Gemini 3 Flash
+- Gemini 3.5 Flash
+- Gemini 2.5 Flash Image (native image generation, "nano banana")
+- Gemini 3.1 Flash Image (native image generation, "nano banana 2")
 
 #### xAI
 - Grok 4

examples/README.md

@@ -84,6 +84,30 @@ python examples/llm_tool_calling.py
 - The model decides when to invoke tools based on the user's query
 - Uses x402 protocol for payment processing
 
+#### `llm_web_search.py`
+Runs an LLM chat completion with native web search enabled.
+
+```bash
+python examples/llm_web_search.py
+```
+
+**What it does:**
+- Sets `web_search=True` to enable the provider's built-in web search
+- The model searches the web to answer the query, citing live sources
+- Each search is billed per search on top of token usage (supported by OpenAI, Anthropic, Google, and xAI models)
+
+#### `llm_image_generation.py`
+Generates an image using a native image-output model ("nano banana").
+
+```bash
+python examples/llm_image_generation.py
+```
+
+**What it does:**
+- Calls an image-output model (`og.TEE_LLM.GEMINI_3_1_FLASH_IMAGE`) with a text prompt
+- Reads the generated images from `result.images` (data URIs) and writes them to disk
+- Image output is billed as completion tokens; images travel out-of-band and are not part of the signed output hash
+
 ## Alpha Testnet Examples
 
 Examples for features only available on the **Alpha Testnet** are located in the [`alpha/`](./alpha/) folder. These include:

examples/llm_image_generation.py

@@ -0,0 +1,52 @@
+import asyncio
+import base64
+import logging
+import os
+import re
+
+import opengradient as og
+
+logging.basicConfig()
+logging.getLogger("opengradient").setLevel(logging.DEBUG)
+
+_DATA_URI_RE = re.compile(r"^data:(?P<mime>[^;,]+)?(?:;base64)?,(?P<data>.*)$", re.DOTALL)
+
+
+def save_data_uri(data_uri: str, path: str) -> None:
+    """Decode a ``data:image/...;base64,...`` URI and write it to ``path``."""
+    match = _DATA_URI_RE.match(data_uri)
+    payload = match.group("data") if match else data_uri
+    with open(path, "wb") as f:
+        f.write(base64.b64decode(payload))
+
+
+async def main():
+    llm = og.LLM(private_key=os.environ.get("OG_PRIVATE_KEY"))
+    llm.ensure_opg_approval(min_allowance=0.1)
+
+    messages = [
+        {"role": "user", "content": "Generate an image of a friendly robot reading a book under a tree."},
+    ]
+
+    # Image-output models ("nano banana") return generated images on the response.
+    # The text caption (if any) is in chat_output["content"]; the generated images
+    # are in result.images as data: URIs. Images travel out-of-band and are not part
+    # of the signed output hash.
+    result = await llm.chat(
+        model=og.TEE_LLM.GEMINI_3_1_FLASH_IMAGE,
+        messages=messages,
+        max_tokens=1024,
+    )
+
+    if result.chat_output and result.chat_output.get("content"):
+        print(result.chat_output["content"])
+
+    images = result.images or []
+    print(f"Generated {len(images)} image(s)")
+    for index, image in enumerate(images):
+        path = f"generated_image_{index + 1}.png"
+        save_data_uri(image, path)
+        print(f"Saved {path}")
+
+
+asyncio.run(main())

examples/llm_web_search.py

@@ -0,0 +1,31 @@
+import asyncio
+import logging
+import os
+
+import opengradient as og
+
+logging.basicConfig()
+logging.getLogger("opengradient").setLevel(logging.DEBUG)
+
+
+async def main():
+    llm = og.LLM(private_key=os.environ.get("OG_PRIVATE_KEY"))
+    llm.ensure_opg_approval(min_allowance=0.1)
+
+    messages = [
+        {"role": "user", "content": "What are the top technology headlines today? Cite your sources."},
+    ]
+
+    # Enable the provider's native web search with web_search=True. Each search is
+    # billed per search on top of token usage. Web search is supported by OpenAI,
+    # Anthropic, Google, and xAI models; other providers ignore the flag.
+    result = await llm.chat(
+        model=og.TEE_LLM.CLAUDE_SONNET_4_6,
+        messages=messages,
+        max_tokens=500,
+        web_search=True,
+    )
+    print(result.chat_output["content"])
+
+
+asyncio.run(main())

src/opengradient/agents/og_langchain.py

@@ -135,6 +135,7 @@ class OpenGradientChatModel(BaseChatModel):
     model_cid: Union[TEE_LLM, str]
     max_tokens: int = 300
     temperature: float = 0.0
+    web_search: bool = False
     x402_settlement_mode: x402SettlementMode = x402SettlementMode.BATCH_HASHED
 
     _llm: LLM = PrivateAttr()
@@ -149,6 +150,7 @@ def __init__(
         model: Optional[Union[TEE_LLM, str]] = None,
         max_tokens: int = 300,
         temperature: float = 0.0,
+        web_search: bool = False,
         x402_settlement_mode: x402SettlementMode = x402SettlementMode.BATCH_HASHED,
         client: Optional[LLM] = None,
         rpc_url: Optional[str] = None,
@@ -165,6 +167,7 @@ def __init__(
             model_cid=resolved_model_cid,
             max_tokens=max_tokens,
             temperature=temperature,
+            web_search=web_search,
             x402_settlement_mode=x402_settlement_mode,
             **kwargs,
         )
@@ -307,6 +310,7 @@ def _build_chat_kwargs(
             "temperature": kwargs.get("temperature", self.temperature),
             "tools": kwargs.get("tools", self._tools),
             "tool_choice": kwargs.get("tool_choice", self._tool_choice),
+            "web_search": kwargs.get("web_search", self.web_search),
             "x402_settlement_mode": x402_settlement_mode,
             "stream": stream,
         }