Add tests for Gemini via openai package (#382)

* Add openai_gemini marker. Add GEMINI_API_KEY to .env.template

* Add tests for gemini via openai

* Remove assumption of one usage per stream

* Use fixture in openai gemini tests

* Add VCRs for gemini tests

* Add docs section on Gemini via OpenAI

* Add list of str test case for xAI grok

.env.template

@@ -1,6 +1,7 @@
 # To set API keys, copy this as .env and fill in the values
 # API keys are required to rewrite the VCR cassettes for the tests
 ANTHROPIC_API_KEY="anthropic_api_key"
+GEMINI_API_KEY="gemini_api_key"
 MISTRAL_API_KEY="mistral_api_key"
 OPENAI_API_KEY="openai_api_key"
 XAI_API_KEY="xai_api_key"

docs/configuration.md

@@ -16,6 +16,26 @@ from magentic import OpenaiChatModel
 model = OpenaiChatModel("gpt-4o")
 ```
 
+### Gemini via OpenAI
+
+Gemini supports an OpenAI-compatible API, allowing you to use their models through the OpenAI backend. See https://ai.google.dev/gemini-api/docs/openai
+
+First, ensure you have the appropriate API key set as an environment variable `GEMINI_API_KEY`.
+
+Then, specify the model name, Google APIs `base_url`, and API key when creating the `OpenaiChatModel` instance.
+
+```python
+import os
+
+from magentic import OpenaiChatModel
+
+model = OpenaiChatModel(
+    "gemini-2.0-flash",
+    base_url="https://generativelanguage.googleapis.com/v1beta/openai/",
+    api_key=os.environ["GEMINI_API_KEY"],
+)
+```
+
 #### Ollama via OpenAI
 
 Ollama supports an OpenAI-compatible API, which allows you to use Ollama models via the OpenAI backend.

pyproject.toml

@@ -91,6 +91,7 @@ markers = [
     "litellm_openai: Tests that query the OpenAI API via litellm. Requires the OPENAI_API_KEY environment variable to be set.",
     "mistral: Tests that query the Mistral API (via openai). Requires the MISTRAL_API_KEY environment variable to be set.",
     "openai: Tests that query the OpenAI API. Requires the OPENAI_API_KEY environment variable to be set.",
+    "openai_gemini: Tests that query the Gemini API via openai. Requires the GEMINI_API_KEY environment variable to be set.",
     "openai_ollama: Tests that query Ollama via openai. Requires ollama to be installed and running on localhost:11434.",
     "openai_xai: Tests that query xAI via openai. Requires the XAI_API_KEY environment variable to be set.",
 ]

src/magentic/chat_model/openai_chat_model.py

@@ -350,7 +350,10 @@ def update(self, item: ChatCompletionChunk) -> None:
                 tool_call_chunk.index = self._current_tool_call_index
         self._chat_completion_stream_state.handle_chunk(item)
         if item.usage:
-            # xAI Grok provides usage on every chunk, so cannot assert usage is None here
+            # Only keep the last usage
+            # xAI Grok and Gemini openai-compatible API includes usage in all streamed chunks
+            # OpenAI only includes this in the last chunk
+            self.usage_ref.clear()
             self.usage_ref.append(
                 Usage(
                     input_tokens=item.usage.prompt_tokens,

tests/chat_model/cassettes/test_openai_chat_model_gemini/test_openai_chat_model_gemini_acomplete[Return True.-output_types1-bool].yaml

@@ -0,0 +1,78 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": "Return True."}], "model": "gemini-2.0-flash",
+      "parallel_tool_calls": false, "stream": true, "stream_options": {"include_usage":
+      true}, "tool_choice": {"type": "function", "function": {"name": "return_bool"}},
+      "tools": [{"type": "function", "function": {"name": "return_bool", "parameters":
+      {"properties": {"value": {"title": "Value", "type": "boolean"}}, "required":
+      ["value"], "type": "object"}}}]}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '442'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+      user-agent:
+      - AsyncOpenAI/Python 1.59.3
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - async:asyncio
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.59.3
+      x-stainless-retry-count:
+      - '0'
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.10.15
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/openai/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAC/0SQ3WrDMAxG7/sUQ9dpSEvYOr/KWoLqKIlX/wRbHi3B7z6ZQXb7nU9HSCMyqrcN
+        9BKMpgTqa4ORLCOoDWKwBAowJZMYPUMDHIIdNFr7V52y12yCr22Mc3bkWQhsV/hBm+kKimOmIpMe
+        XZVF4hz9cBcPlAbMKFn1vtZKd1+5CZyMN2kZImGqKyBxWKVr/EhPUF25NaAFMonk9NF3l8/+ve8b
+        cEFOkP5MTgzHc9sdJ4tpkdlw/ybNwvSC3OrgVkt1YauX7B9SyAlnquf8s4HDg7ycJeo1Ssx7cqkf
+        YbR7cDqXcjj8AgAA//9KAYdqtIu/n2ssFxcAAAD//wMAxiRLqWQBAAA=
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - text/event-stream
+      Date:
+      - Sun, 02 Mar 2025 05:50:44 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=507
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/chat_model/cassettes/test_openai_chat_model_gemini/test_openai_chat_model_gemini_acomplete[Return [1, 2, 3, 4, 5]-output_types2-list].yaml

@@ -0,0 +1,79 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": "Return [1, 2, 3, 4, 5]"}], "model":
+      "gemini-2.0-flash", "parallel_tool_calls": false, "stream": true, "stream_options":
+      {"include_usage": true}, "tool_choice": {"type": "function", "function": {"name":
+      "return_list_of_int"}}, "tools": [{"type": "function", "function": {"name":
+      "return_list_of_int", "parameters": {"properties": {"value": {"items": {"type":
+      "integer"}, "title": "Value", "type": "array"}}, "required": ["value"], "type":
+      "object"}}}]}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '494'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+      user-agent:
+      - AsyncOpenAI/Python 1.59.3
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - async:asyncio
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.59.3
+      x-stainless-retry-count:
+      - '0'
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.10.15
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/openai/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAC/0SQy26DMBBF9/mKatYOSni0Cb/SIGtqBnBjbISHqhXyv3esSnR77/G5GvfI2L7s
+        YKZgDUVo33foyTFCu8MaHEELGKONjJ5BAYfgtEHn/tBh84Zt8JnGddxm8iwN7A/4QrfRQ6irKlWl
+        atV0SQQe5+xcibfVaydiHQZtRZ4U2F6qvPKzZOiwp07KwXobJ70SxjwIkcMirPU9fUN7SZ0CIyWT
+        SK5v9eV2r1/rRsEc5CDhR5rFcC6Ly3lwGCd5Gz4+ybB0ZkIuTJgXR3mwMNPmnwJsEUfKx/13msOT
+        vBx5U7CsEvORlE3+IEZ3JFWV0un0CwAA//9KAQdytIu/n2ssFxcAAAD//wMAifa8XnMBAAA=
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - text/event-stream
+      Date:
+      - Sun, 02 Mar 2025 05:50:45 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=613
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/chat_model/cassettes/test_openai_chat_model_gemini/test_openai_chat_model_gemini_acomplete[Return a list of fruit-output_types3-list].yaml

@@ -0,0 +1,80 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": "Return a list of fruit"}], "model":
+      "gemini-2.0-flash", "parallel_tool_calls": false, "stream": true, "stream_options":
+      {"include_usage": true}, "tool_choice": {"type": "function", "function": {"name":
+      "return_list_of_str"}}, "tools": [{"type": "function", "function": {"name":
+      "return_list_of_str", "parameters": {"properties": {"value": {"items": {"type":
+      "string"}, "title": "Value", "type": "array"}}, "required": ["value"], "type":
+      "object"}}}]}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '493'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+      user-agent:
+      - AsyncOpenAI/Python 1.59.3
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - async:asyncio
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.59.3
+      x-stainless-retry-count:
+      - '0'
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.10.15
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/openai/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAC/0SQ0W6DMAxF3/sVk59TVCrWbfxKQZEbDGQNCUqcaRPi3+dsUqu8WPee+NoekLF9
+        2cDMwRpK0F43GMgxQrtBDI6gBUzJJkbPoIBDcNqgc//omL1hG3yhMU55Ic/iwNbBF7pMnVAd4Lo6
+        KVUHN/Ty/soQ0U+i9ru09biUpEico9dO4nQYdeIIuwI7iFWyf9YCPTL3XszReptmHQlTGQMSh1VY
+        6wf6hva09wqMmEzSpH5rTu8fzaW5KFiCrCn8RIt0OJ6r03F0mGb5G26fZFg8MyNXJiwyfQmszJz9
+        XYCccKKy8tPTHO7kZfW6VrBG0fkpNeVujO6hnF/3/XD4BQAA//9KAYd9tIu/n2ssFxcAAAD//wMA
+        uRfCRIoBAAA=
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - text/event-stream
+      Date:
+      - Sun, 02 Mar 2025 05:50:46 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=575
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+    status:
+      code: 200
+      message: OK
+version: 1

tests/chat_model/cassettes/test_openai_chat_model_gemini/test_openai_chat_model_gemini_acomplete[Say hello!-output_types0-str].yaml

@@ -0,0 +1,75 @@
+interactions:
+- request:
+    body: '{"messages": [{"role": "user", "content": "Say hello!"}], "model": "gemini-2.0-flash",
+      "stream": true, "stream_options": {"include_usage": true}}'
+    headers:
+      accept:
+      - application/json
+      accept-encoding:
+      - gzip, deflate
+      connection:
+      - keep-alive
+      content-length:
+      - '145'
+      content-type:
+      - application/json
+      host:
+      - generativelanguage.googleapis.com
+      user-agent:
+      - AsyncOpenAI/Python 1.59.3
+      x-stainless-arch:
+      - arm64
+      x-stainless-async:
+      - async:asyncio
+      x-stainless-lang:
+      - python
+      x-stainless-os:
+      - MacOS
+      x-stainless-package-version:
+      - 1.59.3
+      x-stainless-retry-count:
+      - '0'
+      x-stainless-runtime:
+      - CPython
+      x-stainless-runtime-version:
+      - 3.10.15
+    method: POST
+    uri: https://generativelanguage.googleapis.com/v1beta/openai/chat/completions
+  response:
+    body:
+      string: !!binary |
+        H4sIAAAAAAAC/0SOXQrCMBCE3z2F7HNaqhR/cgLvICJruraxabY0WxBK7u6KqI/7zezMNCho1wu4
+        jr2jBPa8QENBEKxCjkJRwMKJQmAwMHEgPTElnwRVygZ8bOgJtsoXA24iFGrAbvZ1dTjWu7o2MLAm
+        6ldLg4++2JZVcQ+YOs3j24Pcu8B1KKXjYQwknmPpujn2apgTtvTZ8tWuwj1FnVoZGCfF8iNaJiwY
+        /iDn1eoFAAD//0ohzo8KJRmpRalD1ZMAAAAA//8i0pOKCh755QrJiXkKngoZqTkFCpX5pQol+SmJ
+        lfYxeTj8ngb0UHFGPNDXxfl5QMnikvwCpQEMEkNDjDAxRg8TQ3CgAAAAAP//ggRKtIu/n2ssFxcA
+        AAD//wMA/RvkTOwCAAA=
+    headers:
+      Alt-Svc:
+      - h3=":443"; ma=2592000,h3-29=":443"; ma=2592000
+      Content-Encoding:
+      - gzip
+      Content-Type:
+      - text/event-stream
+      Date:
+      - Sun, 02 Mar 2025 05:50:44 GMT
+      Server:
+      - scaffolding on HTTPServer2
+      Server-Timing:
+      - gfet4t7; dur=373
+      Transfer-Encoding:
+      - chunked
+      Vary:
+      - Origin
+      - X-Origin
+      - Referer
+      X-Content-Type-Options:
+      - nosniff
+      X-Frame-Options:
+      - SAMEORIGIN
+      X-XSS-Protection:
+      - '0'
+    status:
+      code: 200
+      message: OK
+version: 1