Fix OpenAI reasoning caps and Gemini transparent images

pixelagent/core/limits.py

@@ -2,6 +2,7 @@
 
 Each provider spells the cap differently, and a caller who already set it in
 `chat_kwargs` / `tool_kwargs` keeps that value.
+OpenAI defaults to `max_completion_tokens`, which also supports reasoning models.
 
 `openai_model_kwargs` is non-deterministic on purpose. `chat_completions`
 writes `tools` into the dict it receives, so a cached dict would leak that
@@ -21,7 +22,7 @@ def openai_model_kwargs(options: Optional[dict], max_tokens: int, variant: int =
     # chat_completions mutates.
     merged = dict(options or {})
     if "max_tokens" not in merged and "max_completion_tokens" not in merged:
-        merged["max_tokens"] = max_tokens
+        merged["max_completion_tokens"] = max_tokens
     return merged
 
 

pixelagent/gemini/utils.py

@@ -15,7 +15,8 @@ def create_content(
     """Flatten the conversation to Gemini content parts.
 
     Text-only turns are a one-element list. An image turn prepends an
-    `inline_data` part. `generate_content` accepts that list as `contents`.
+    `inline_data` PNG part, preserving transparency. `generate_content` accepts
+    that list as `contents`.
     """
     context = ""
     for msg in memory_context:
@@ -26,10 +27,11 @@ def create_content(
         return [context]
 
     buf = io.BytesIO()
-    image.save(buf, format="JPEG")
+    # PNG preserves transparency; RGBA also handles modes such as CMYK that PNG cannot encode directly.
+    image.convert("RGBA").save(buf, format="PNG")
     encoded = base64.b64encode(buf.getvalue()).decode("ascii")
     return [
-        {"inline_data": {"mime_type": "image/jpeg", "data": encoded}},
+        {"inline_data": {"mime_type": "image/png", "data": encoded}},
         context,
     ]
 

tests/test_gemini.py

@@ -1,7 +1,16 @@
+import asyncio
+import base64
+import io
+import json
 import os
 
+import httpx
 import pixeltable as pxt
+import pixeltable.functions.gemini as pxt_gemini
 import pytest
+from google import genai
+from google.genai import types
+from PIL import Image
 
 from pixelagent.gemini import Agent
 
@@ -10,6 +19,42 @@
 )
 
 
+@pytest.mark.parametrize("mode", ["RGB", "RGBA", "LA", "P", "CMYK"])
+def test_gemini_image_request(monkeypatch: pytest.MonkeyPatch, mode: str) -> None:
+    """Image turns reach the SDK with valid PNG bytes and preserve transparency."""
+    image = Image.new("RGBA", (8, 8), color=(255, 0, 0, 64)).convert(mode)
+    requests = []
+
+    def respond(request: httpx.Request) -> httpx.Response:
+        requests.append(json.loads(request.content))
+        return httpx.Response(200, json={
+            "candidates": [{"content": {"role": "model", "parts": [{"text": "Red."}]}}],
+        })
+
+    client = genai.Client(
+        api_key="test-placeholder",
+        http_options=types.HttpOptions(
+            httpx_async_client=httpx.AsyncClient(transport=httpx.MockTransport(respond)),
+        ),
+    )
+    monkeypatch.setattr(pxt_gemini, "_genai_client", lambda: client)
+    try:
+        agent = Agent(name="gemini_image_request", system_prompt="Be helpful.", reset=True)
+        assert agent.chat("What color is this?", image=image) == "Red."
+    finally:
+        asyncio.run(client.aio.aclose())
+        client.close()
+
+    assert len(requests) == 1
+    blob = requests[0]["contents"][0]["parts"][0]["inlineData"]
+    assert blob["mimeType"] == "image/png"
+    with Image.open(io.BytesIO(base64.urlsafe_b64decode(blob["data"]))) as decoded:
+        assert decoded.format == "PNG"
+        assert decoded.size == image.size
+        assert decoded.convert("RGBA").tobytes() == image.convert("RGBA").tobytes()
+    assert agent.memory.count() == 2
+
+
 @pxt.udf
 def weather(city: str) -> str:
     """Returns the weather in a given city."""

tests/test_openai.py

@@ -1,6 +1,12 @@
+import asyncio
+import json
 import os
+from collections.abc import Iterator
 
+import httpx2
+import openai
 import pixeltable as pxt
+import pixeltable.functions.openai as pxt_openai
 import pytest
 
 from pixelagent.openai import Agent
@@ -10,6 +16,83 @@
 )
 
 
+@pytest.fixture
+def openai_requests(monkeypatch: pytest.MonkeyPatch) -> Iterator[list[dict]]:
+    """Capture actual SDK request bodies without making network calls."""
+    requests = []
+
+    def respond(request: httpx2.Request) -> httpx2.Response:
+        body = json.loads(request.content)
+        requests.append(body)
+        message = {"role": "assistant", "content": "The price is 5."}
+        if body.get("tools"):
+            message = {
+                "role": "assistant",
+                "content": None,
+                "tool_calls": [{
+                    "id": "call_stock_price",
+                    "type": "function",
+                    "function": {
+                        "name": "stock_price_dict",
+                        "arguments": json.dumps({"ticker": "NVDA"}),
+                    },
+                }],
+            }
+        return httpx2.Response(200, json={
+            "id": "test_completion",
+            "object": "chat.completion",
+            "created": 0,
+            "model": body["model"],
+            "choices": [{"index": 0, "message": message, "finish_reason": "stop"}],
+        })
+
+    client = openai.AsyncOpenAI(
+        api_key="test-placeholder",
+        http_client=httpx2.AsyncClient(transport=httpx2.MockTransport(respond)),
+        max_retries=0,
+    )
+    monkeypatch.setattr(pxt_openai, "_openai_client", lambda: client)
+    yield requests
+    asyncio.run(client.close())
+
+
+@pytest.mark.parametrize("model", ["o3", "gpt-4o-mini"])
+def test_openai_default_token_cap(openai_requests: list[dict], model: str) -> None:
+    agent = Agent(name="openai_token_cap", system_prompt="Be helpful.", model=model, max_tokens=128, reset=True)
+    assert agent.chat("Hello") == "The price is 5."
+    assert openai_requests[0]["max_completion_tokens"] == 128
+    assert "max_tokens" not in openai_requests[0]
+
+
+@pytest.mark.parametrize("token_option", ["max_tokens", "max_completion_tokens"])
+def test_openai_explicit_token_cap(openai_requests: list[dict], token_option: str) -> None:
+    options = {token_option: 64}
+    agent = Agent(
+        name="openai_explicit_token_cap", system_prompt="Be helpful.",
+        max_tokens=128, chat_kwargs=options, reset=True,
+    )
+    agent.chat("Hello")
+    agent.chat("Hello again", model_kwargs={token_option: 32})
+    assert [body[token_option] for body in openai_requests] == [64, 32]
+    other_option = "max_completion_tokens" if token_option == "max_tokens" else "max_tokens"
+    assert all(other_option not in body for body in openai_requests)
+    assert options == {token_option: 64}
+
+
+def test_openai_tool_token_cap(openai_requests: list[dict], mock_stock_price_dict: pxt.func.Function) -> None:
+    agent = Agent(
+        name="openai_tool_token_cap", system_prompt="Be helpful.", model="o3",
+        tools=pxt.tools(mock_stock_price_dict), max_tokens=128, reset=True,
+    )
+    assert agent.tool_call("What is NVDA's price?") == "The price is 5."
+    assert len(openai_requests) == 2
+    assert all(body["max_completion_tokens"] == 128 for body in openai_requests)
+    assert all("max_tokens" not in body for body in openai_requests)
+    assert "tools" in openai_requests[0]
+    assert "tools" not in openai_requests[1]
+    assert "regularMarketPrice" in openai_requests[1]["messages"][0]["content"]
+
+
 @pytest.fixture
 def openai_agent(mock_stock_price_dict):
     """Fixture to create an OpenAI agent with tools"""